MillePedeAlignmentAlgorithm/scripts/mps_setup_new_align.py

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