DTWorkflow.py

Go to the documentation of this file.

from __future__ import print_function
from __future__ import absolute_import
import os,sys
import glob
import logging
import argparse
import subprocess
import time, datetime
import urllib
import json

from . import tools
from .CLIHelper import CLIHelper
from .CrabHelper import CrabHelper
import FWCore.ParameterSet.Config as cms
log = logging.getLogger(__name__)

class DTWorkflow(CLIHelper, CrabHelper):
    """ This is the base class for all DTWorkflows and contains some
        common tasks """
    def __init__(self, options):
        self.options = options
        super( DTWorkflow, self ).__init__()
        self.digilabel = "muonDTDigis"
        # dict to hold required variables. Can not be marked in argparse to allow
        # loading of options from config
        self.required_options_dict = {}
        self.required_options_prepare_dict = {}
        self.fill_required_options_dict()
        self.fill_required_options_prepare_dict()
        # These variables are determined in the derived classes
        self.pset_name = ""
        self.outpath_command_tag = ""
        self.output_files = []
        self.input_files = []

        self.run_all_command = False
        self.files_reveived = False
        self._user = ""
        # change to working directory
        os.chdir(self.options.working_dir)

    def check_missing_options(self, requirements_dict):
        missing_options = []
        # check if all required options exist
        if self.options.command in requirements_dict:
            for option in requirements_dict[self.options.command]:
                if not (hasattr(self.options, option)
                    and ( (getattr(self.options,option))
                          or isinstance(getattr(self.options,option), bool) )):
                    missing_options.append(option)
        if len(missing_options) > 0:
            err = "The following CLI options are missing"
            err += " for command %s: " % self.options.command
            err += " ".join(missing_options)
            raise ValueError(err)

    def run(self):
        """ Generalized function to run workflow command"""
        msg = "Preparing %s workflow" % self.options.workflow
        if hasattr(self.options, "command"):
            msg += " for command %s" % self.options.command
        log.info(msg)
        if self.options.config_path:
            self.load_options( self.options.config_path )
        #check if all options to prepare the command are used
        self.check_missing_options(self.required_options_prepare_dict)
        self.prepare_workflow()
        # create output folder if they do not exist yet
        if not os.path.exists( self.local_path ):
            os.makedirs(self.local_path)
        # dump used options
        self.dump_options()
        #check if all options to run the command are used
        self.check_missing_options(self.required_options_dict)
        try:
            run_function = getattr(self, self.options.command)
        except AttributeError:
            errmsg = "Class `{}` does not implement `{}` for workflow %s" % self.options.workflow
            if hasattr(self.options, "workflow_mode"):
                errmsg += "and workflow mode %s" % self.options.workflow_mode
            raise NotImplementedError( errmsg.format(self.__class__.__name__,
                                                     self.options.command))
        log.debug("Running command %s" % self.options.command)
        # call chosen function
        run_function()

    def prepare_workflow(self):
        """ Abstract implementation of prepare workflow function"""
        errmsg = "Class `{}` does not implement `{}`"
        raise NotImplementedError( errmsg.format(self.__class__.__name__,
                                                     "prepare_workflow"))

    def all(self):
        """ generalized function to perform several workflow mode commands in chain.
            All commands mus be specified in self.all_commands list in workflow mode specific
            prepare function in child workflow objects.
        """
        self.run_all_command = True
        for command in self.all_commands:
            log.info(f"Will run command: {command}")
            self.options.command = command
            self.run()

    def submit(self):
        self.submit_crab_task()

    def check(self):
        """ Function to check status of submitted tasks """
        self.check_crabtask()

    def write(self):
        self.runCMSSWtask()

    def dump(self):
        self.runCMSSWtask()

    def correction(self):
        self.runCMSSWtask()

    def add_preselection(self):
        """ Add preselection to the process object stored in workflow_object"""
        if not hasattr(self, "process"):
            raise NameError("Process is not initalized in workflow object")
        pathsequence = self.options.preselection.split(':')[0]
        seqname = self.options.preselection.split(':')[1]
        self.process.load(pathsequence)
        tools.prependPaths(self.process, seqname)

    def add_raw_option(self):
        getattr(self.process, self.digilabel).inputLabel = self.options.raw_data_label
        tools.prependPaths(self.process,self.digilabel)

    def add_local_t0_db(self, local=False):
        """ Add a local t0 database as input. Use the option local is used
            if the pset is processed locally and not with crab.
        """
        if local:
            connect = os.path.abspath(self.options.inputT0DB)
        else:
            connect = os.path.basename(self.options.inputT0DB)
        self.addPoolDBESSource( process = self.process,
                                moduleName = 't0DB',
                                record = 'DTT0Rcd',
                                tag = 't0',
                                connect =  'sqlite_file:%s' % connect)
        self.input_files.append(os.path.abspath(self.options.inputT0DB))

    def add_local_vdrift_db(self, local=False):
        """ Add a local vdrift database as input. Use the option local is used
            if the pset is processed locally and not with crab.
         """
        if local:
            connect = os.path.abspath(self.options.inputVDriftDB)
        else:
            connect = os.path.basename(self.options.inputVDriftDB)
        self.addPoolDBESSource( process = self.process,
                                moduleName = 'vDriftDB',
                                record = 'DTMtimeRcd',
                                tag = 'vDrift',
                                connect = 'sqlite_file:%s' % connect)
        self.input_files.append( os.path.abspath(self.options.inputVDriftDB) )

    def add_local_calib_db(self, local=False):
        """ Add a local calib database as input. Use the option local is used
            if the pset is processed locally and not with crab.
         """
        label = ''
        if self.options.datasettype == "Cosmics":
            label = 'cosmics'
        if local:
            connect = os.path.abspath(self.options.inputCalibDB)
        else:
            connect = os.path.basename(self.options.inputCalibDB)
        self.addPoolDBESSource( process = self.process,
                                moduleName = 'calibDB',
                                record = 'DTTtrigRcd',
                                tag = 'ttrig',
                                connect = str("sqlite_file:%s" % connect),
                                label = label
                                )
        self.input_files.append( os.path.abspath(self.options.inputCalibDB) )

    def add_local_custom_db(self):
        for option in ('inputDBRcd', 'connectStrDBTag'):
            if hasattr(self.options, option) and not getattr(self.options, option):
                raise ValueError("Option %s needed for custom input db" % option)
        self.addPoolDBESSource( process = self.process,
                                    record = self.options.inputDBRcd,
                                    tag = self.options.inputDBTag,
                                    connect = self.options.connectStrDBTag,
                                    moduleName = 'customDB%s' % self.options.inputDBRcd
                                   )

    def prepare_common_submit(self):
        """ Common operations used in most prepare_[workflow_mode]_submit functions"""
        if not self.options.run:
            raise ValueError("Option run is required for submission!")
        if hasattr(self.options, "inputT0DB") and self.options.inputT0DB:
            self.add_local_t0_db()

        if hasattr(self.options, "inputVDriftDB") and self.options.inputVDriftDB:
            self.add_local_vdrift_db()

        if hasattr(self.options, "inputDBTag") and self.options.inputDBTag:
            self.add_local_custom_db()

        if self.options.run_on_RAW:
            self.add_raw_option()
        if self.options.preselection:
            self.add_preselection()

    def prepare_common_write(self, do_hadd=True):
        """ Common operations used in most prepare_[workflow_mode]_erite functions"""
        self.load_options_command("submit")
        output_path = os.path.join( self.local_path, "unmerged_results" )
        merged_file = os.path.join(self.result_path, self.output_file)
        crabtask = self.crabFunctions.CrabTask(crab_config = self.crab_config_filepath,
                                               initUpdate = False)
        if not (self.options.skip_stageout or self.files_reveived or self.options.no_exec):
            output_files =  self.get_output_files(crabtask, output_path)
            if "xrootd" not in output_files.keys():
                raise RuntimeError("Could not get output files. No xrootd key found.")
            if len(output_files["xrootd"]) == 0:
                raise RuntimeError("Could not get output files. Output file list is empty.")
            log.info("Received files from storage element")
            log.info("Using hadd to merge output files")
        if not self.options.no_exec and do_hadd:
            returncode = tools.haddLocal(output_files["xrootd"], merged_file)
            if returncode != 0:
                raise RuntimeError("Failed to merge files with hadd")
        return crabtask.crabConfig.Data.outputDatasetTag

    def prepare_common_dump(self, db_path):
        self.process = tools.loadCmsProcess(self.pset_template)
        self.process.calibDB.connect = 'sqlite_file:%s' % db_path
        try:
            path = self.result_path
        except:
            path = os.getcwd()
        print("path", path)
        out_path = os.path.abspath(os.path.join(path,
                                                os.path.splitext(db_path)[0] + ".txt"))

        self.process.dumpToFile.outputFileName = out_path

    @staticmethod
    def addPoolDBESSource( process,
                           moduleName,
                           record,
                           tag,
                           connect='sqlite_file:',
                           label='',):

        from CondCore.CondDB.CondDB_cfi import CondDB

        calibDB = cms.ESSource("PoolDBESSource",
                               CondDB,
                               timetype = cms.string('runnumber'),
                               toGet = cms.VPSet(cms.PSet(
                                   record = cms.string(record),
                                   tag = cms.string(tag),
                                   label = cms.untracked.string(label)
                                    )),
                               )
        calibDB.connect = cms.string( str(connect) )
        #if authPath: calibDB.DBParameters.authenticationPath = authPath
        if 'oracle:' in connect:
            calibDB.DBParameters.authenticationPath = '/afs/cern.ch/cms/DB/conddb'
        setattr(process,moduleName,calibDB)
        setattr(process,"es_prefer_" + moduleName,cms.ESPrefer('PoolDBESSource',
                                                                moduleName)
                                                                )

    def get_output_files(self, crabtask, output_path):
        res = self.crab.callCrabCommand( ["getoutput",
                                    "--dump",
                                    "--xrootd",
                                    crabtask.crabFolder ] )
        
        return res

    def runCMSSWtask(self, pset_path=""):
        """ Run a cmsRun job locally. The member variable self.pset_path is used
            if pset_path argument is not given"""
        if self.options.no_exec:
            return 0
        process = subprocess.Popen( "cmsRun %s" % self.pset_path,
                            stdout=subprocess.PIPE,
                            stderr=subprocess.STDOUT,
                            shell = True)
        stdout = process.communicate()[0]
        log.info(stdout)
        if process.returncode != 0:
            raise RuntimeError("Failed to use cmsRun for pset %s" % self.pset_name)
        return process.returncode

    @property
    def remote_out_path(self):
        """ Output path on remote excluding user base path
        Returns a dict if crab is used due to crab path setting policy"""
        if self.options.command =="submit":
            return {
                "outLFNDirBase" : os.path.join( "/store",
                                                "user",
                                                self.user,
                                                'DTCalibration/',
                                                self.outpath_command_tag,
                                                self.outpath_workflow_mode_tag),
                "outputDatasetTag" : self.tag
                    }
        else:
            return os.path.join( 'DTCalibration/',
                                 datasetstr,
                                 'Run' + str(self.options.run),
                                 self.outpath_command_tag,
                                 self.outpath_workflow_mode_tag,
                                 'v' + str(self.options.trial),
                                )
    @property
    def outpath_workflow_mode_tag(self):
        if not self.options.workflow_mode in self.outpath_workflow_mode_dict:
            raise NotImplementedError("%s missing in outpath_workflow_mode_dict" % self.options.workflow_mode)
        return self.outpath_workflow_mode_dict[self.options.workflow_mode]

    @property
    def tag(self):
        return 'Run' + str(self.options.run) + '_v' + str(self.options.trial)

    @property
    def user(self):
        if self._user:
            return self._user
        if hasattr(self.options, "user") and self.options.user:
            self._user = self.options.user
        else:
            self._user = self.crab.checkusername()
        return self._user

    @property
    def local_path(self):
        """ Output path on local machine """
        if self.options.run and self.options.label:
            prefix = "Run%d-%s_v%d" % ( self.options.run,
                                        self.options.label,
                                        self.options.trial)
        else:
            prefix = ""
        if self.outpath_workflow_mode_tag:
            path = os.path.join( self.options.working_dir,
                                 prefix,
                                 self.outpath_workflow_mode_tag)
        else:
            path =  os.path.join( self.options.working_dir,
                                  prefix,
                                  self.outpath_command_tag )
        return path

    @property
    def result_path(self):
        result_path = os.path.abspath(os.path.join(self.local_path,"results"))
        if not os.path.exists(result_path):
            os.makedirs(result_path)
        return result_path

    @property
    def pset_template_base_bath(self):
        """ Base path to folder containing pset files for cmsRun"""
        return os.path.expandvars(os.path.join("$CMSSW_BASE",
                                               "src",
                                               "CalibMuon",
                                               "test",
                                               )
                                 )

    @property
    def pset_path(self):
        """ full path to the pset file """
        basepath = os.path.join( self.local_path, "psets")
        if not os.path.exists( basepath ):
            os.makedirs( basepath )
        return os.path.join( basepath, self.pset_name )

    def write_pset_file(self):
        if not hasattr(self, "process"):
            raise NameError("Process is not initalized in workflow object")
        if not os.path.exists(self.local_path):
            os.makedirs(self.local_path)
        with open( self.pset_path,'w') as pfile:
            pfile.write(self.process.dumpPython())

    def get_config_name(self, command= ""):
        """ Create the name for the output json file which will be dumped"""
        if not command:
            command = self.options.command
        return "config_" + command + ".json"

    def dump_options(self):
        with open(os.path.join(self.local_path, self.get_config_name()),"w") as out_file:
            json.dump(vars(self.options), out_file, indent=4)

    def load_options(self, config_file_path):
        if not os.path.exists(config_file_path):
            raise IOError("File %s not found" % config_file_path)
        with open(config_file_path, "r") as input_file:
            config_json = json.load(input_file)
            for key, val in config_json.items():
                if not hasattr(self.options, key) or not getattr(self.options, key):
                    setattr(self.options, key, val)

    def load_options_command(self, command ):
        """Load options for previous command in workflow """
        if not self.options.config_path:
            if not self.options.run:
                raise RuntimeError("Option run is required if no config path specified")
            if not os.path.exists(self.local_path):
                raise IOError("Local path %s does not exist" % self.local_path)
            self.options.config_path = os.path.join(self.local_path,
                                                    self.get_config_name(command))
        self.load_options( self.options.config_path )