MillePedeAlignmentAlgorithm/scripts/mps_setup_new_align.py

0001 #!/usr/bin/env python3
0002
0003 import os
0004 import sys
0005 import re
0006 import fcntl
0007 import glob
0008 import shutil
0009 import datetime
0010 import subprocess
0011
0012 if "CMSSW_BASE" not in os.environ:
0013     print("You need to source the CMSSW environment first.")
0014     sys.exit(1)
0015
0016 from Alignment.MillePedeAlignmentAlgorithm.alignmentsetup.helper \
0017     import checked_out_MPS
0018
0019 required_version = (2,7)
0020 if sys.version_info < required_version:
0021     print("Your Python interpreter is too old. Need version 2.7 or higher.")
0022     sys.exit(1)
0023
0024 import argparse
0025
0026
0027 ################################################################################
0028 def main(argv = None):
0029     """Main routine of the script.
0030
0031     Arguments:
0032     - `argv`: arguments passed to the main routine
0033     """
0034
0035     if argv == None:
0036         argv = sys.argv[1:]
0037
0038     parser = argparse.ArgumentParser(
0039         description="Setup a new alignment campaign in the MPproduction area.")
0040     parser.add_argument("-d", "--description", dest="description", required=True,
0041                         help="comment to describe the purpose of the campaign")
0042     parser.add_argument("-t", "--data-type", dest="type", required=True,
0043                         metavar="TYPE", choices=["MC", "data"],
0044                         help="type of the input data (choices: %(choices)s)")
0045     parser.add_argument("-c", "--copy", dest="copy", metavar="CAMPAIGN",
0046                         help="input campaign (optional)")
0047     args = parser.parse_args(argv)
0048
0049
0050     if os.path.basename(os.path.normpath(os.getcwd())) != "MPproduction":
0051         print(">>> Cannot create a campaign outside of the 'MPproduction' area.")
0052         print(">>> Please change to the 'MPproduction' directory first.")
0053         sys.exit(1)
0054
0055     if len(args.description.strip()) == 0:
0056         print(">>> Please provide a non-empty description of the campaign")
0057         sys.exit(1)
0058
0059     MPS_dir = os.path.join("src", "Alignment", "MillePedeAlignmentAlgorithm")
0060     args.checked_out = checked_out_MPS()
0061     if args.checked_out[0]:
0062         MPS_dir = os.path.join(os.environ["CMSSW_BASE"], MPS_dir)
0063     else:
0064         MPS_dir = os.path.join(os.environ["CMSSW_RELEASE_BASE"], MPS_dir)
0065     args.MPS_dir = MPS_dir
0066
0067     mp_regex = re.compile(r"mp([0-9]+).*")
0068     all_campaign_numbers = sorted(map(lambda x: get_first_match(mp_regex, x),
0069                                       os.listdir(".")))
0070     next_number = (0
0071                    if len(all_campaign_numbers) == 0
0072                    else sorted(all_campaign_numbers)[-1] + 1)
0073
0074     while True:
0075         try:
0076             number_of_digits = len(str(next_number))
0077             number_of_digits = 4 if number_of_digits <= 4 else number_of_digits
0078             next_campaign = "mp{{0:0{0}d}}".format(number_of_digits)
0079             next_campaign = next_campaign.format(next_number)
0080             os.makedirs(next_campaign)
0081             print(">>> Created new campaign:", next_campaign)
0082
0083             campaign_list = "MP_ali_list.txt"
0084             with open(campaign_list, "a") as f:
0085                 campaign_info = add_campaign(f, next_campaign, args)
0086             backup_dir = ".MP_ali_list"
0087             try:
0088                 os.makedirs(backup_dir)
0089             except OSError as e:
0090                 if e.args == (17, 'File exists'):
0091                     pass
0092                 else:
0093                     raise
0094             with open(os.path.join(backup_dir, campaign_list), "a") as f:
0095                 fcntl.flock(f, fcntl.LOCK_EX)
0096                 f.write(campaign_info)
0097                 fcntl.flock(f, fcntl.LOCK_UN)
0098             print("    - updated campaign list '"+campaign_list+"'")
0099
0100             if args.copy is None:
0101                 copy_default_templates(args, next_campaign)
0102             else:
0103                 copied_files = []
0104                 for ext in ("py", "ini", "txt"):
0105                     for config_file in glob.glob(args.copy+"/*."+ext):
0106                         copied_files.append(os.path.basename(config_file))
0107                         shutil.copy(config_file, next_campaign)
0108                 if len(copied_files) == 0:
0109                     print("    - no configuration files for '"+args.copy+"'")
0110                     copy_default_templates(args, next_campaign)
0111                 else:
0112                     alignment_config_ini = os.path.join(next_campaign,
0113                                                         "alignment_config.ini")
0114                     regex_input = (r"^(jobname\s*[=:])(\s*)"+
0115                                    os.path.basename(args.copy.strip("/"))+r"\s*$",
0116                                    r"\1 "+next_campaign+r"\n")
0117                     if os.path.isfile(alignment_config_ini):
0118                         customize_default_template(alignment_config_ini,
0119                                                    regex_input)
0120                     print("    - copied configuration files from", end=' ')
0121                     print("'"+args.copy+"':", ", ".join(copied_files))
0122
0123         except OSError as e:
0124             if e.args == (17, 'File exists'):
0125                 next_number += 1 # someone created a campaign ~at the same time
0126                 continue
0127             else:
0128                 raise
0129         break
0130
0131
0132 ################################################################################
0133 def get_first_match(regex, directory):
0134     """
0135     Checks if `directory` matches `regex` and returns the first match converted
0136     to an integer. If it does not match -1 is returned.
0137
0138     Arguments:
0139     - `regex`: Regular expression to be tested against
0140     - `directory`: name of the directory under test
0141     """
0142
0143     result = regex.search(directory)
0144     if result is None:
0145         return -1
0146     else:
0147         return int(result.group(1))
0148
0149
0150 def add_campaign(campaign_file, campaign, args):
0151     """Adds a line with campaign information from `args` to `campaign_file`.
0152
0153     Arguments:
0154     - `campaign_file`: output file
0155     - `campaign`: name of the campaign
0156     - `args`: command line arguments for this campaign
0157     """
0158
0159     campaign_info = campaign.ljust(10)
0160     campaign_info += os.environ["USER"].ljust(12)
0161     campaign_info += datetime.date.today().isoformat().ljust(11)
0162
0163     version = os.environ["CMSSW_VERSION"]
0164     if args.checked_out[1]:
0165         local_area = os.path.join(os.environ["CMSSW_BASE"], "src")
0166         with open(os.devnull, 'w') as devnull:
0167             # check which tags (-> e.g. release names) point at current commit
0168             p = subprocess.Popen(["git", "tag", "--points-at", "HEAD"],
0169                                  cwd = local_area, stdout=subprocess.PIPE,
0170                                  stderr=devnull)
0171             tags = p.communicate()[0].split()
0172             # check for deleted, untracked, modified files respecting .gitignore:
0173             p = subprocess.Popen(["git", "ls-files", "-d", "-o", "-m",
0174                                   "--exclude-standard"],
0175                                  cwd = local_area, stdout=subprocess.PIPE,
0176                                  stderr=devnull)
0177             files = p.communicate()[0].split()
0178             # check for staged tracked files:
0179             p = subprocess.Popen(["git", "diff", "--name-only", "--staged"],
0180                                  cwd = local_area, stdout=subprocess.PIPE,
0181                                  stderr=devnull)
0182             files.extend(p.communicate()[0].split())
0183         if version not in tags or len(files) != 0:
0184             version += " (mod.)"
0185
0186     campaign_info += version.ljust(34)
0187     campaign_info += args.type.ljust(17)
0188     campaign_info += args.description.strip() + "\n"
0189
0190     fcntl.flock(campaign_file, fcntl.LOCK_EX)
0191     campaign_file.write(campaign_info)
0192     fcntl.flock(campaign_file, fcntl.LOCK_UN)
0193
0194     return campaign_info
0195
0196
0197 def copy_default_templates(args, next_campaign):
0198     """Copies the default configuration templates.
0199
0200     Arguments:
0201     - `args`: container with the needed information
0202     - `next_campaign`: destination for the copy operation
0203     """
0204
0205     default_conf_dir = os.path.join(args.MPS_dir, "templates")
0206     template_files = ("universalConfigTemplate.py", "alignment_config.ini")
0207     for f in template_files:
0208         shutil.copy(os.path.join(default_conf_dir, f), next_campaign)
0209
0210     # customize alignment_config.ini
0211     # - replace job name with campaign ID as initial value
0212     # - replace global tag with the corresponding auto GT depending on data type
0213     auto_gt = args.type.replace("MC", "phase1_2017_realistic")
0214     auto_gt = auto_gt.replace("data", "run2_data")
0215     customize_default_template(os.path.join(next_campaign, "alignment_config.ini"),
0216                                (r"(jobname\s*[=:])(.*)", r"\1 "+next_campaign),
0217                                (r"(globaltag\s*[=:])(.*)",
0218                                 r"\1 auto:"+auto_gt))
0219
0220     print("    - copied default configuration templates from", end=' ')
0221     print("'"+default_conf_dir+"'")
0222     print("    - please modify these template files according to your needs:", end=' ')
0223     print(", ".join(template_files))
0224
0225
0226 def customize_default_template(file_name, *regex_replace_pairs):
0227     """
0228     Replace all lines in `file_name` matching `regex_string` with
0229     `replace_string`.
0230     Lines starting with '#' or ';' are ignored.
0231
0232     Arguments:
0233     - `file_name`: path to the file to be customized
0234     - `regex_replace_pairs`: tuples containing a regex string and its
0235                              replacement
0236     """
0237
0238     comment = re.compile(r"^\s*[;#]")
0239     replacements = []
0240     for regex_string, replace_string in regex_replace_pairs:
0241         replacements.append((re.compile(regex_string), replace_string))
0242
0243     customized = ""
0244     with open(file_name, "r") as f:
0245         for line in f:
0246             custom_line = line
0247             if not re.match(comment, custom_line):
0248                 for regex,replace_string in replacements:
0249                     custom_line = re.sub(regex, replace_string, custom_line)
0250             customized += custom_line
0251     with open(file_name, "w") as f:
0252         f.write(customized)
0253
0254
0255 ################################################################################
0256 if __name__ == "__main__":
0257     main()