/
githubmirror
/
cmssw
Обзор
Документация
Войти
/
githubmirror
/
cmssw
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
Alignment/CommonAlignment/scripts/tkal_create_file_lists.py
1 177 строк
46 KB
Shahzad Malik Muzaffar
[ALCA-DB] py2/3 compatibility:drop use of __future__
22 ноя 2024, 20:15
22 ноя 2024, 20:15
c26340c
Код
Авторство
О чём код?
#!/usr/bin/env python3 from builtins import range import os import re import sys import glob import json import math import bisect import random import signal if sys.version_info[0]>2: import _pickle as cPickle else: import cPickle import difflib import argparse import functools import itertools import subprocess import collections import multiprocessing import FWCore.PythonUtilities.LumiList as LumiList import Utilities.General.cmssw_das_client as cmssw_das_client import Alignment.MillePedeAlignmentAlgorithm.mpslib.tools as mps_tools ################################################################################ def main(argv = None): """ Main routine. Not called, if this module is loaded via `import`. Arguments: - `argv`: Command line arguments passed to the script. """ if argv == None: argv = sys.argv[1:] file_list_creator = FileListCreator(argv) file_list_creator.create() ################################################################################ class FileListCreator(object): """Create file lists for alignment and validation for a given dataset. """ def __init__(self, argv): """Constructor taking the command line arguments. Arguments: - `args`: command line arguments """ self._first_dataset_ini = True self._parser = self._define_parser() self._args = self._parser.parse_args(argv) if not mps_tools.check_proxy(): print_msg( "Please create proxy via 'voms-proxy-init -voms cms -rfc'.") sys.exit(1) self._dataset_regex = re.compile(r"^/([^/]+)/([^/]+)/([^/]+)$") self._validate_input() if self._args.test_mode: import Configuration.PyReleaseValidation.relval_steps as rvs import Configuration.PyReleaseValidation.relval_production as rvp self._args.datasets = [rvs.steps[rvp.workflows[1000][1][0]]["INPUT"].dataSet] self._validate_input() # ensure that this change is valid self._datasets = sorted([dataset for pattern in self._args.datasets for dataset in get_datasets(pattern) if re.search(self._args.dataset_filter, dataset)]) if len(self._datasets) == 0: print_msg("Found no dataset matching the pattern(s):") for d in self._args.datasets: print_msg("\t"+d) sys.exit(1) self._formatted_dataset = merge_strings( [re.sub(self._dataset_regex, r"\1_\2_\3", dataset) for dataset in self._datasets]) self._output_dir = os.path.join(self._args.output_dir, self._formatted_dataset) self._output_dir = os.path.abspath(self._output_dir) self._cache = _DasCache(self._output_dir) self._prepare_iov_datastructures() self._prepare_run_datastructures() try: os.makedirs(self._output_dir) except OSError as e: if e.args == (17, "File exists"): if self._args.force: pass # do nothing, just clear the existing output elif self._args.use_cache: self._cache.load() # load cache before clearing the output else: print_msg("Directory '{}' already exists from previous runs" " of the script. Use '--use-cache' if you want to" " use the cached DAS-query results Or use " "'--force' to remove it." .format(self._output_dir)) sys.exit(1) files = glob.glob(os.path.join(self._output_dir, "*")) for f in files: os.remove(f) else: raise def create(self): """Creates file list. To be called by user of the class.""" self._request_dataset_information() self._create_file_lists() self._print_eventcounts() self._write_file_lists() _event_count_log = "event_count_info.log" def _define_parser(self): """Definition of command line argument parser.""" parser = argparse.ArgumentParser( description = "Create file lists for alignment", epilog = ("The tool will create a directory containing all file " "lists and a log file with all relevant event counts " "('{}').".format(FileListCreator._event_count_log))) parser.add_argument("-i", "--input", dest = "datasets", required = True, metavar = "DATASET", action = "append", help = ("CMS dataset name; supports wildcards; " "use multiple times for multiple datasets")) parser.add_argument("--dataset-filter", default = "", help = "regex to match within in the datasets matched," "in case the wildcard isn't flexible enough") parser.add_argument("-j", "--json", dest = "json", metavar = "PATH", help = "path to JSON file (optional)") parser.add_argument("-f", "--fraction", dest = "fraction", type = float, default = 1, help = "max. fraction of files used for alignment") parser.add_argument("--iov", dest = "iovs", metavar = "RUN", type = int, action = "append", default = [], help = ("define IOV by specifying first run; for " "multiple IOVs use this option multiple " "times; files from runs before the lowest " "IOV are discarded (default: 1)")) parser.add_argument("--miniiov", dest="miniiovs", metavar="RUN", type=int, action="append", default=[], help=("in addition to the standard IOVs, break up hippy jobs " "at these points, so that jobs from before and after " "these runs are not in the same job")) parser.add_argument("-r", "--random", action = "store_true", default = False, help = "select files randomly") parser.add_argument("-n", "--events-for-alignment", "--maxevents", dest = "events", type = int, metavar = "NUMBER", help = ("number of events needed for alignment; the" " remaining events in the dataset are used " "for validation; if n<=0, all events are " "used for validation")) parser.add_argument("--all-events", action = "store_true", help = "Use all events for alignment") parser.add_argument("--tracks-for-alignment", dest = "tracks", type = int, metavar = "NUMBER", help = "number of tracks needed for alignment") parser.add_argument("--track-rate", dest = "rate", type = float, metavar = "NUMBER", help = "number of tracks per event") parser.add_argument("--run-by-run", dest = "run_by_run", action = "store_true", default = False, help = "create validation file list for each run") parser.add_argument("--minimum-events-in-iov", dest = "minimum_events_in_iov", metavar = "NUMBER", type = int, default = 100000, help = ("minimum number of events for alignment per" " IOV; this option has a higher priority " "than '-f/--fraction' " "(default: %(default)s)")) parser.add_argument("--minimum-events-validation", dest = "minimum_events_validation", metavar = "NUMBER", type = int, default = 1, help = ("minimum number of events for validation; " "applies to IOVs; in case of --run-by-run " "it applies to runs runs " "(default: %(default)s)")) parser.add_argument("--use-cache", dest = "use_cache", action = "store_true", default = False, help = "use DAS-query results of previous run") parser.add_argument("-o", "--output-dir", dest = "output_dir", metavar = "PATH", default = os.getcwd(), help = "output base directory (default: %(default)s)") parser.add_argument("--create-ini", dest = "create_ini", action = "store_true", default = False, help = ("create dataset ini file based on the " "created file lists")) parser.add_argument("--force", action = "store_true", default = False, help = ("remove output directory from previous " "runs, if existing")) parser.add_argument("--hippy-events-per-job", type = int, default = 1, help = ("approximate number of events in each job for HipPy")) parser.add_argument("--test-mode", dest = "test_mode", action = "store_true", default = False, help = argparse.SUPPRESS) # hidden option return parser def _validate_input(self): """Validate command line arguments.""" if self._args.events is None: if self._args.all_events: self._args.events = float("inf") print_msg("Using all tracks for alignment") elif (self._args.tracks is None) and (self._args.rate is None): msg = ("either -n/--events-for-alignment, --all-events, or both of " "--tracks-for-alignment and --track-rate are required") self._parser.error(msg) elif (((self._args.tracks is not None) and (self._args.rate is None)) or ((self._args.rate is not None)and (self._args.tracks is None))): msg = ("--tracks-for-alignment and --track-rate must be used " "together") self._parser.error(msg) else: self._args.events = int(math.ceil(self._args.tracks / self._args.rate)) print_msg("Requested {0:d} tracks with {1:.2f} tracks/event " "-> {2:d} events for alignment." .format(self._args.tracks, self._args.rate, self._args.events)) else: if (self._args.tracks is not None) or (self._args.rate is not None) or self._args.all_events: msg = ("-n/--events-for-alignment must not be used with " "--tracks-for-alignment, --track-rate, or --all-events") self._parser.error(msg) print_msg("Requested {0:d} events for alignment." .format(self._args.events)) for dataset in self._args.datasets: if not re.match(self._dataset_regex, dataset): print_msg("Dataset pattern '"+dataset+"' is not in CMS format.") sys.exit(1) nonzero_events_per_iov = (self._args.minimum_events_in_iov > 0) if nonzero_events_per_iov and self._args.fraction <= 0: print_msg("Setting minimum number of events per IOV for alignment " "to 0 because a non-positive fraction of alignment events" " is chosen: {}".format(self._args.fraction)) nonzero_events_per_iov = False self._args.minimum_events_in_iov = 0 if nonzero_events_per_iov and self._args.events <= 0: print_msg("Setting minimum number of events per IOV for alignment " "to 0 because a non-positive number of alignment events" " is chosen: {}".format(self._args.events)) nonzero_events_per_iov = False self._args.minimum_events_in_iov = 0 def _prepare_iov_datastructures(self): """Create the needed objects for IOV handling.""" self._iovs = sorted(set(self._args.iovs)) if len(self._iovs) == 0: self._iovs.append(1) self._iov_info_alignment = {iov: {"events": 0, "files": []} for iov in self._iovs} self._iov_info_validation = {iov: {"events": 0, "files": []} for iov in self._iovs} self._miniiovs = sorted(set(self._iovs) | set(self._args.miniiovs)) def _get_iovs(self, runs, useminiiovs=False): """ Return the IOV start for `run`. Returns 'None' if the run is before any defined IOV. Arguments: - `runs`: run numbers """ iovlist = self._miniiovs if useminiiovs else self._iovs iovs = [] for run in runs: iov_index = bisect.bisect(iovlist, run) if iov_index > 0: iovs.append(iovlist[iov_index-1]) return iovs def _prepare_run_datastructures(self): """Create the needed objects for run-by-run validation file lists.""" self._run_info = {} def _add_file_info(self, container, keys, fileinfo): """Add file with `file_name` to `container` using `key`. Arguments: - `container`: dictionary holding information on files and event counts - `keys`: keys to which the info should be added; will be created if not existing - `file_name`: name of a dataset file """ for key in keys: if key not in container: container[key] = {"events": 0, "files": []} container[key]["events"] += fileinfo.nevents / len(keys) if fileinfo not in container[key]["files"]: container[key]["files"].append(fileinfo) def _remove_file_info(self, container, keys, fileinfo): """Remove file with `file_name` to `container` using `key`. Arguments: - `container`: dictionary holding information on files and event counts - `keys`: keys from which the info should be removed - `file_name`: name of a dataset file - `event_count`: number of events in `file_name` """ for key in keys: if key not in container: continue try: index = container[key]["files"].index(fileinfo) except ValueError: # file not found return del container[key]["files"][index] container[key]["events"] -= fileinfo.nevents / len(keys) def _request_dataset_information(self): """Retrieve general dataset information and create file list.""" if not self._cache.empty: print_msg("Using cached information.") (self._events_in_dataset, self._files, self._file_info, self._max_run) = self._cache.get() self.rereco = any(len(fileinfo.runs)>1 for fileinfo in self._file_info) if self._args.random: random.shuffle(self._files) return # workaround to deal with KeyboardInterrupts in the worker processes: # - ignore interrupt signals in workers (see initializer) # - use a timeout of size sys.maxsize to avoid a bug in multiprocessing number_of_processes = multiprocessing.cpu_count() - 1 number_of_processes = (number_of_processes if number_of_processes > 0 else 1) pool = multiprocessing.Pool( processes = number_of_processes, initializer = lambda: signal.signal(signal.SIGINT, signal.SIG_IGN)) print_msg("Requesting information for the following dataset(s):") for d in self._datasets: print_msg("\t"+d) print_msg("This may take a while...") result = pool.map_async(get_events_per_dataset, self._datasets).get(3600) self._events_in_dataset = sum(result) result = pool.map_async(get_max_run, self._datasets).get(3600) self._max_run = max(result) result = sum(pool.map_async(get_file_info, self._datasets).get(3600), []) files = pool.map_async(_make_file_info, result).get(3600) self._file_info = sorted(fileinfo for fileinfo in files) self.rereco = any(len(fileinfo.runs)>1 for fileinfo in self._file_info) if self._args.test_mode: self._file_info = self._file_info[-200:] # take only last chunk of files self._files = [fileinfo.name for fileinfo in self._file_info] # write information to cache self._cache.set(self._events_in_dataset, self._files, self._file_info, self._max_run) self._cache.dump() if self._args.random: random.shuffle(self._file_info) self._files = [fileinfo.name for fileinfo in self._file_info] def _create_file_lists(self): """Create file lists for alignment and validation.""" # collect files for alignment until minimal requirements are fulfilled self._files_alignment = [] self._files_validation = [] self._events_for_alignment = 0 self._events_for_validation = 0 max_range = (0 if self._args.events <= 0 else int(math.ceil(len(self._files)*self._args.fraction))) use_for_alignment = True for i, fileinfo in enumerate(self._file_info): enough_events = self._events_for_alignment >= self._args.events fraction_exceeded = i >= max_range if enough_events or fraction_exceeded: use_for_alignment = False dataset, f, number_of_events, runs = fileinfo iovs = self._get_iovs(runs) if use_for_alignment: if iovs: self._events_for_alignment += number_of_events self._files_alignment.append(fileinfo) self._add_file_info(self._iov_info_alignment, iovs, fileinfo) else: max_range += 1 # not used -> discard in fraction calculation else: if iovs: self._events_for_validation += number_of_events self._files_validation.append(fileinfo) self._add_file_info(self._iov_info_validation, iovs, fileinfo) if self._args.run_by_run: self._add_file_info(self._run_info, runs, fileinfo) self._fulfill_iov_eventcount() self._split_hippy_jobs() def _fulfill_iov_eventcount(self): """ Try to fulfill the requirement on the minimum number of events per IOV in the alignment file list by picking files from the validation list. """ for iov in self._iovs: if self._iov_info_alignment[iov]["events"] >= self._args.minimum_events_in_iov: continue for fileinfo in self._files_validation[:]: dataset, f, number_of_events, runs = fileinfo iovs = self._get_iovs(runs) if iov in iovs: self._files_alignment.append(fileinfo) self._events_for_alignment += number_of_events self._add_file_info(self._iov_info_alignment, iovs, fileinfo) self._events_for_validation -= number_of_events self._remove_file_info(self._iov_info_validation, iovs, fileinfo) if self._args.run_by_run: self._remove_file_info(self._run_info, runs, fileinfo) self._files_validation.remove(fileinfo) if (self._iov_info_alignment[iov]["events"] >= self._args.minimum_events_in_iov): break # break the file loop if already enough events def _split_hippy_jobs(self): hippyjobs = {} for dataset, miniiov in itertools.product(self._datasets, self._miniiovs): jobsforminiiov = [] hippyjobs[dataset,miniiov] = jobsforminiiov eventsinthisjob = float("inf") for fileinfo in self._files_alignment: if fileinfo.dataset != dataset: continue miniiovs = set(self._get_iovs(fileinfo.runs, useminiiovs=True)) if miniiov not in miniiovs: continue if len(miniiovs) > 1: hippyjobs[dataset,miniiov] = [] if eventsinthisjob >= self._args.hippy_events_per_job: currentjob = [] jobsforminiiov.append(currentjob) eventsinthisjob = 0 currentjob.append(fileinfo) currentjob.sort() eventsinthisjob += fileinfo.nevents self._hippy_jobs = { (dataset, iov): sum((hippyjobs[dataset, miniiov] for miniiov in self._miniiovs if iov == max(_ for _ in self._iovs if _ <= miniiov)), [] ) for dataset, iov in itertools.product(self._datasets, self._iovs) } def _print_eventcounts(self): """Print the event counts per file list and per IOV.""" log = os.path.join(self._output_dir, FileListCreator._event_count_log) print_msg("Using {0:d} events for alignment ({1:.2f}%)." .format(self._events_for_alignment, 100.0* self._events_for_alignment/self._events_in_dataset), log_file = log) for iov in sorted(self._iov_info_alignment): print_msg(("Approximate events" if self.rereco else "Events") + " for alignment in IOV since {0:f}: {1:f}" .format(iov, self._iov_info_alignment[iov]["events"]), log_file = log) print_msg("Using {0:d} events for validation ({1:.2f}%)." .format(self._events_for_validation, 100.0* self._events_for_validation/self._events_in_dataset), log_file = log) for iov in sorted(self._iov_info_validation): msg = ("Approximate events" if self.rereco else "Events") + " for validation in IOV since {0:f}: {1:f}".format( iov, self._iov_info_validation[iov]["events"]) if (self._iov_info_validation[iov]["events"] < self._args.minimum_events_validation): msg += " (not enough events -> no dataset file will be created)" print_msg(msg, log_file = log) for run in sorted(self._run_info): msg = ("Approximate events" if self.rereco else "Events") + " for validation in run {0:f}: {1:f}".format( run, self._run_info[run]["events"]) if (self._run_info[run]["events"] < self._args.minimum_events_validation): msg += " (not enough events -> no dataset file will be created)" print_msg(msg, log_file = log) unused_events = (self._events_in_dataset - self._events_for_validation - self._events_for_alignment) if unused_events > 0 != self._events_in_dataset: print_msg("Unused events: {0:d} ({1:.2f}%)" .format(unused_events, 100.0*unused_events/self._events_in_dataset), log_file = log) def _create_dataset_ini_section(self, name, collection, json_file = None): """Write dataset ini snippet. Arguments: - `name`: name of the dataset section - `collection`: track collection of this dataset - `json_file`: JSON file to be used for this dataset (optional) """ if json_file: splitted = name.split("_since") file_list = "_since".join(splitted[:-1] if len(splitted) > 1 else splitted) else: file_list = name output = "[dataset:{}]\n".format(name) output += "collection = {}\n".format(collection) output += "inputFileList = ${{datasetdir}}/{}.txt\n".format(file_list) output += "json = ${{datasetdir}}/{}\n".format(json_file) if json_file else "" if collection in ("ALCARECOTkAlCosmicsCTF0T", "ALCARECOTkAlCosmicsInCollisions"): if self._first_dataset_ini: print_msg("\tDetermined cosmics dataset, i.e. please replace " "'DUMMY_DECO_MODE_FLAG' and 'DUMMY_ZERO_TESLA_FLAG' " "with the correct values.") self._first_dataset_ini = False output += "cosmicsDecoMode = DUMMY_DECO_MODE_FLAG\n" output += "cosmicsZeroTesla = DUMMY_ZERO_TESLA_FLAG\n" output += "\n" return output def _create_json_file(self, name, first, last = None): """ Create JSON file with `name` covering runs from `first` to `last`. If a global JSON is provided, the resulting file is the intersection of the file created here and the global one. Returns the name of the created JSON file. Arguments: - `name`: name of the creted JSON file - `first`: first run covered by the JSON file - `last`: last run covered by the JSON file """ if last is None: last = self._max_run name += "_JSON.txt" print_msg("Creating JSON file: "+name) json_file = LumiList.LumiList(runs = range(first, last+1)) if self._args.json: global_json = LumiList.LumiList(filename = self._args.json) json_file = json_file & global_json json_file.writeJSON(os.path.join(self._output_dir, name)) return name def _get_track_collection(self, edm_file): """Extract track collection from given `edm_file`. Arguments: - `edm_file`: CMSSW dataset file """ # use global redirector to allow also files not yet at your site: cmd = ["edmDumpEventContent", r"root://cms-xrd-global.cern.ch/"+edm_file] try: event_content = subprocess.check_output(cmd).split("\n") except subprocess.CalledProcessError as e: splitted = edm_file.split("/") try: alcareco = splitted[splitted.index("ALCARECO")+1].split("-")[0] alcareco = alcareco.replace("TkAlCosmics0T", "TkAlCosmicsCTF0T") alcareco = "ALCARECO" + alcareco print_msg("\tDetermined track collection as '{}'.".format(alcareco)) return alcareco except ValueError: if "RECO" in splitted: print_msg("\tDetermined track collection as 'generalTracks'.") return "generalTracks" else: print_msg("\tCould not determine track collection " "automatically.") print_msg("\tPlease replace 'DUMMY_TRACK_COLLECTION' with " "the correct value.") return "DUMMY_TRACK_COLLECTION" track_collections = [] for line in event_content: splitted = line.split() if len(splitted) > 0 and splitted[0] == r"vector<reco::Track>": track_collections.append(splitted[1].strip().strip('"')) if len(track_collections) == 0: print_msg("No track collection found in file '{}'.".format(edm_file)) sys.exit(1) elif len(track_collections) == 1: print_msg("\tDetermined track collection as " "'{}'.".format(track_collections[0])) return track_collections[0] else: alcareco_tracks = filter(lambda x: x.startswith("ALCARECO"), track_collections) if len(alcareco_tracks) == 0 and "generalTracks" in track_collections: print_msg("\tDetermined track collection as 'generalTracks'.") return "generalTracks" elif len(alcareco_tracks) == 1: print_msg("\tDetermined track collection as " "'{}'.".format(alcareco_tracks[0])) return alcareco_tracks[0] print_msg("\tCould not unambiguously determine track collection in " "file '{}':".format(edm_file)) print_msg("\tPlease replace 'DUMMY_TRACK_COLLECTION' with " "the correct value from the following list.") for collection in track_collections: print_msg("\t - "+collection) return "DUMMY_TRACK_COLLECTION" def _write_file_lists(self): """Write file lists to disk.""" self._create_dataset_txt(self._formatted_dataset, self._files_alignment) self._create_hippy_txt(self._formatted_dataset, sum(self._hippy_jobs.values(), [])) self._create_dataset_cff( "_".join(["Alignment", self._formatted_dataset]), self._files_alignment) self._create_dataset_cff( "_".join(["Validation", self._formatted_dataset]), self._files_validation) if self._args.create_ini: dataset_ini_general = "[general]\n" dataset_ini_general += "datasetdir = {}\n".format(self._output_dir) dataset_ini_general += ("json = {}\n\n".format(self._args.json) if self._args.json else "\n") ini_path = self._formatted_dataset + ".ini" print_msg("Creating dataset ini file: " + ini_path) ini_path = os.path.join(self._output_dir, ini_path) collection = self._get_track_collection(self._files[0]) with open(ini_path, "w") as f: f.write(dataset_ini_general) f.write(self._create_dataset_ini_section( self._formatted_dataset, collection)) iov_wise_ini = dataset_ini_general for i,iov in enumerate(sorted(self._iovs)): iov_str = "since{0:d}".format(iov) iov_str = "_".join([self._formatted_dataset, iov_str]) if self.rereco: if i == len(self._iovs) - 1: last = None else: last = sorted(self._iovs)[i+1] - 1 local_json = self._create_json_file(iov_str, iov, last) else: local_json = None if self._args.create_ini: iov_wise_ini += self._create_dataset_ini_section(iov_str, collection, local_json) self._create_dataset_txt(iov_str, self._iov_info_alignment[iov]["files"]) self._create_hippy_txt(iov_str, sum((self._hippy_jobs[dataset,iov] for dataset in self._datasets), [])) self._create_dataset_cff( "_".join(["Alignment", iov_str]), self._iov_info_alignment[iov]["files"], json_file=local_json) if (self._iov_info_validation[iov]["events"] < self._args.minimum_events_validation): continue self._create_dataset_cff( "_".join(["Validation", iov_str]), self._iov_info_validation[iov]["files"], json_file=local_json) if self._args.create_ini and iov_wise_ini != dataset_ini_general: ini_path = self._formatted_dataset + "_IOVs.ini" print_msg("Creating dataset ini file: " + ini_path) ini_path = os.path.join(self._output_dir, ini_path) with open(ini_path, "w") as f: f.write(iov_wise_ini) for run in sorted(self._run_info): if args.rereco: continue #need to implement more jsons if (self._run_info[run]["events"] < self._args.minimum_events_validation): continue self._create_dataset_cff( "_".join(["Validation", self._formatted_dataset, str(run)]), self._run_info[run]["files"]) def _create_dataset_txt(self, name, file_list): """Write alignment file list to disk. Arguments: - `name`: name of the file list - `file_list`: list of files to write to `name` """ name += ".txt" print_msg("Creating dataset file list: "+name) with open(os.path.join(self._output_dir, name), "w") as f: f.write("\n".join(fileinfo.name for fileinfo in file_list)) def _create_hippy_txt(self, name, job_list): name += "_hippy.txt" print_msg("Creating dataset file list for HipPy: "+name) with open(os.path.join(self._output_dir, name), "w") as f: f.write("\n".join(",".join("'"+fileinfo.name+"'" for fileinfo in job) for job in job_list)+"\n") def _create_dataset_cff(self, name, file_list, json_file = None): """ Create configuration fragment to define a dataset. Arguments: - `name`: name of the configuration fragment - `file_list`: list of files to write to `name` - `json_file`: JSON file to be used for this dataset (optional) """ if json_file is None: json_file = self._args.json # might still be None if json_file is not None: json_file = os.path.join(self._output_dir, json_file) name = "_".join(["Dataset",name, "cff.py"]) print_msg("Creating dataset configuration fragment: "+name) file_list_str = "" for sub_list in get_chunks(file_list, 255): file_list_str += ("readFiles.extend([\n'"+ "',\n'".join(fileinfo.name for fileinfo in sub_list)+ "'\n])\n") fragment = FileListCreator._dataset_template.format( lumi_def = ("import FWCore.PythonUtilities.LumiList as LumiList\n\n" "lumiSecs = cms.untracked.VLuminosityBlockRange()\n" "goodLumiSecs = LumiList.LumiList(filename = " "'{0:s}').getCMSSWString().split(',')" .format(json_file) if json_file else ""), lumi_arg = ("lumisToProcess = lumiSecs,\n " if json_file else ""), lumi_extend = "lumiSecs.extend(goodLumiSecs)" if json_file else "", files = file_list_str) with open(os.path.join(self._output_dir, name), "w") as f: f.write(fragment) _dataset_template = """\ import FWCore.ParameterSet.Config as cms {lumi_def:s} readFiles = cms.untracked.vstring() source = cms.Source("PoolSource", {lumi_arg:s}fileNames = readFiles) {files:s}{lumi_extend:s} maxEvents = cms.untracked.PSet(input = cms.untracked.int32(-1)) """ class _DasCache(object): """Helper class to cache information from DAS requests.""" def __init__(self, file_list_id): """Constructor of the cache. Arguments: - `file_list_id`: ID of the cached file lists """ self._file_list_id = file_list_id self._cache_file_name = os.path.join(file_list_id, ".das_cache.pkl") self.reset() def reset(self): """Reset the cache contents and the 'empty' flag.""" self._empty = True self._events_in_dataset = 0 self._files = [] self._file_info = [] self._max_run = None def set(self, total_events, file_list, file_info, max_run): """Set the content of the cache. Arguments: - `total_events`: total number of events in dataset - `file_list`: list of files in dataset - `file_info`: dictionary with numbers of events per file - `max_run`: highest run number contained in the dataset """ self._events_in_dataset = total_events self._files = file_list self._file_info = file_info self._max_run = max_run self._empty = False def get(self): """ Get the content of the cache as tuple: result = (total number of events in dataset, list of files in dataset, dictionary with numbers of events and runs per file) """ return self._events_in_dataset, self._files, self._file_info, self._max_run def load(self): """Loads the cached contents.""" if not self.empty: print_msg("Overriding file information with cached information.") try: with open(self._cache_file_name, "rb") as f: tmp_dict = cPickle.load(f) self.__dict__.update(tmp_dict) except IOError as e: if e.args == (2, "No such file or directory"): msg = "Failed to load cache for '{}'.".format(self._file_list_id) if not self.empty: msg += " Keeping the previous file information." print_msg(msg) else: raise def dump(self): """Dumps the contents to the cache file.""" if self.empty: print_msg("Cache is empty. Not writing to file.") return with open(self._cache_file_name, "wb") as f: cPickle.dump(self.__dict__, f, 2) @property def empty(self): """ Flag indicating whether the cache is empty or has been filled (possibly with nothing). """ return self._empty ################################################################################ def das_client(query, check_key = None): """ Submit `query` to DAS client and handle possible errors. Further treatment of the output might be necessary. Arguments: - `query`: DAS query - `check_key`: optional key to be checked for; retriggers query if needed """ error = True das_data = {'status': 'error'} for i in range(5): # maximum of 5 tries try: das_data = cmssw_das_client.get_data(query, limit = 0) except IOError as e: if e.errno == 14: #https://stackoverflow.com/q/36397853/5228524 continue except ValueError as e: if str(e) == "No JSON object could be decoded": das_data['reason'] = str(e) continue if das_data["status"] == "ok": if das_data["nresults"] == 0 or check_key is None: error = False break result_count = 0 for d in find_key(das_data["data"], [check_key]): result_count += len(d) if result_count == 0: das_data["status"] = "error" das_data["reason"] = ("DAS did not return required data.") continue else: error = False break if das_data["status"] == "error": print_msg("DAS query '{}' failed 5 times. " "The last time for the the following reason:".format(query)) print(das_data.get("reason", "ERROR:UNKNOWN")) sys.exit(1) return das_data["data"] def find_key(collection, key_chain): """Searches for `key` in `collection` and returns first corresponding value. Arguments: - `collection`: list of dictionaries - `key_chain`: chain of keys to be searched for """ result = None for i,key in enumerate(key_chain): for item in collection: if key in item: if i == len(key_chain) - 1: result = item[key] else: try: result = find_key(item[key], key_chain[i+1:]) except LookupError: pass # continue with next `item` in `collection` else: pass # continue with next `item` in `collection` if result is not None: return result raise LookupError(key_chain, collection) # put def print_msg(text, line_break = True, log_file = None): """Formatted printing of `text`. Arguments: - `text`: string to be printed """ msg = " >>> " + str(text) if line_break: print(msg) else: print(msg, end=' ') sys.stdout.flush() if log_file: with open(log_file, "a") as f: f.write(msg+"\n") return msg def get_runs(file_name): """ Try to guess the run number from `file_name`. If run could not be determined, gets the run numbers from DAS (slow!) Arguments: - `file_name`: name of the considered file """ try: return [int("".join(file_name.split("/")[-4:-2]))] except ValueError: query = "run file="+file_name+" system=dbs3" return [int(_) for _ in find_key(das_client(query), ["run", "run_number"])] def get_max_run(dataset_name): """Retrieve the maximum run number in `dataset_name`. Arguments: - `dataset_name`: name of the dataset """ data = das_client("run dataset={0:s} system=dbs3".format(dataset_name)) runs = [f["run"][0]["run_number"] for f in data] return max(runs) def get_files(dataset_name): """Retrieve list of files in `dataset_name`. Arguments: - `dataset_name`: name of the dataset """ data = das_client(("file dataset={0:s} system=dbs3 detail=True | "+ "grep file.name, file.nevents > 0").format(dataset_name), "file") return [find_key(f["file"], ["name"]) for f in data] def get_datasets(dataset_pattern): """Retrieve list of dataset matching `dataset_pattern`. Arguments: - `dataset_pattern`: pattern of dataset names """ data = das_client("dataset dataset={0:s} system=dbs3 detail=True" "| grep dataset.name".format(dataset_pattern), "dataset") return sorted(set([find_key(f["dataset"], ["name"]) for f in data])) def get_events_per_dataset(dataset_name): """Retrieve the number of a events in `dataset_name`. Arguments: - `dataset_name`: name of a dataset """ return _get_events("dataset", dataset_name) def get_events_per_file(file_name): """Retrieve the number of a events in `file_name`. Arguments: - `file_name`: name of a dataset file """ return _get_events("file", file_name) def _get_events(entity, name): """Retrieve the number of events from `entity` called `name`. Arguments: - `entity`: type of entity - `name`: name of entity """ data = das_client("{0:s}={1:s} system=dbs3 detail=True | grep {0:s}.nevents" .format(entity, name), entity) return int(find_key(data, [entity, "nevents"])) def _get_properties(name, entity, properties, filters = None, sub_entity = None, aggregators = None): """Retrieve `properties` from `entity` called `name`. Arguments: - `name`: name of entity - `entity`: type of entity - `properties`: list of property names - `filters`: list of filters on properties - `sub_entity`: type of entity from which to extract the properties; defaults to `entity` - `aggregators`: additional aggregators/filters to amend to query """ if sub_entity is None: sub_entity = entity if filters is None: filters = [] props = ["{0:s}.{1:s}".format(sub_entity,prop.split()[0]) for prop in properties] conditions = ["{0:s}.{1:s}".format(sub_entity, filt) for filt in filters] add_ons = "" if aggregators is None else " | "+" | ".join(aggregators) data = das_client("{0:s} {1:s}={2:s} system=dbs3 detail=True | grep {3:s}{4:s}" .format(sub_entity, entity, name, ", ".join(props+conditions), add_ons), sub_entity) return [[find_key(f[sub_entity], [prop]) for prop in properties] for f in data] def get_file_info(dataset): result = _get_properties(name=dataset, properties = ["name", "nevents"], filters = ["nevents > 0"], entity = "dataset", sub_entity = "file") return [(dataset, name, nevents) for name, nevents in result] FileInfo = collections.namedtuple("FileInfo", "dataset name nevents runs") def _make_file_info(dataset_name_nevents): return FileInfo(*dataset_name_nevents, runs=get_runs(dataset_name_nevents[1])) def get_chunks(long_list, chunk_size): """ Generates list of sub-lists of `long_list` with a maximum size of `chunk_size`. Arguments: - `long_list`: original list - `chunk_size`: maximum size of created sub-lists """ for i in range(0, len(long_list), chunk_size): yield long_list[i:i+chunk_size] def merge_strings(strings): """Merge strings in `strings` into a common string. Arguments: - `strings`: list of strings """ if type(strings) == str: return strings elif len(strings) == 0: return "" elif len(strings) == 1: return strings[0] elif len(strings) == 2: first = strings[0] second = strings[1] else: first = merge_strings(strings[:-1]) second = strings[-1] merged_string = "" blocks = difflib.SequenceMatcher(None, first, second).get_matching_blocks() last_i, last_j, last_n = 0, 0, 0 for i, j, n in blocks: merged_string += first[last_i+last_n:i] merged_string += second[last_j+last_n:j] merged_string += first[i:i+n] last_i, last_j, last_n = i, j, n return str(merged_string) ################################################################################ if __name__ == "__main__": try: main() except KeyboardInterrupt: pass