submit_concat_project.py

Go to the documentation of this file.

#!/usr/bin/env python

"""
  submit_concat_project.py:
      Submit jobs for producing concats from a dataset.
    
    Original author: J. Wolcott <jwolcott@fnal.gov>, rewritten from submit_concat_project.sh
               Date: August 2017
"""

import argparse
import os, os.path
import shlex
import stat
import subprocess
import sys

import samweb_client

try:
    search_path = os.environ["FW_SEARCH_PATH"]
except KeyError:
    print >> sys.stderr, "Release not yet set up.  Please set up novasoft before using this script!"
    sys.exit(1)
cafana_dir = ""
for path in search_path.split(":"):
    if os.path.isdir(os.path.join(path, "CAFAna")):
        cafana_dir = path
        break
if cafana_dir == "":
    print >> sys.stderr, "Couldn't find CAFAna directory!  Abort."
    sys.exit(1)

CVMFS_RELEASE_DIR = "/cvmfs/nova.opensciencegrid.org/novasoft/slf6/novasoft/releases"

DECAF_MACROS={
    "validation"         : cafana_dir + "/CAFAna/nus/reduce_nue_or_numu_or_nus.C",
    "concat"             : "$NOVAGRIDUTILS_DIR/bin/concat_dataset.C",
    "nueSA"              : cafana_dir + "/CAFAna/nue/reduce_nue_sa.C",
    "nue2017"            : cafana_dir + "/CAFAna/nue/reducer/reduce_bendecomp2017.C",
    "nue2018"            : cafana_dir + "/CAFAna/nue/reducer/reduce_nue_2018.C",
    "nue2019"            : cafana_dir + "/CAFAna/nue/reducer/reduce_nue_2018.C",
    "numuSA"             : cafana_dir + "/CAFAna/numu/FirstAnalysis/reduce_numu_fa.C",
    "numu2017"           : cafana_dir + "/CAFAna/numu/Analysis2017/reduce_numu_ana2017.C",
    "numu2018"           : cafana_dir + "/CAFAna/numu/Analysis2018/reduce_numu_ana2018.C",
    "numu2019"           : cafana_dir + "/CAFAna/numu/Analysis2018/reduce_numu_ana2018.C",
    "nus"                : cafana_dir + "/CAFAna/nus/reduce_nus.C",
    "nus2019"            : cafana_dir + "/CAFAna/nus/reduce_nus_ana2019.C",
    "numu2020"           : cafana_dir + "/3FlavorAna/Ana2020/reducer/reduce_prod5_numu.C",
    "nue2020"            : cafana_dir + "/3FlavorAna/Ana2020/reducer/reduce_prod5_nue.C",
    "nus2020"            : cafana_dir + "/NuXAna/macros/reduce_nus_ana2020.C",
    "nue_or_numu_SA"     : cafana_dir + "/CAFAna/nue/reduce_nue_or_numu_sa.C",
    "nu_or_numu_or_nus"  : cafana_dir + "/CAFAna/nus/reduce_nue_or_numu_or_nus.C",
}

# There's a hard cap on the total number of SAM projects
# allowed to be running experiment-wide (=1000).
# Don't allow a single concatting pass to eat more than
# 10% of that.
MAX_N_JOBS = 100

# this will be the SAM client,
# but we don't want to load it until we've checked
# command-line arguments etc.
SAM = None

def BuildOutputNameStub(dataset_name, decaf_macro, rename_cafs=True):
    # make the output file name stub.
    # this is a concat script, so we want to ensure "caf" -> "sumcaf",
    # and "decaf" -> "sumdecaf"...

    out_name = dataset_name
    replacements = {"_caf": "_sumcaf", "_decaf": "_sumdecaf", "_restrictedcaf": "_sumrestrictedcaf"}
    for repl in replacements:
        out_name = out_name.replace(repl, replacements[repl])
    
    # ... and if we're actually decaffing, that "caf" -> "decaf"
    if rename_cafs and os.path.basename(decaf_macro) != "concat_dataset.C":
        out_name = out_name.replace("_sumcaf", "_sumdecaf")
        out_name = out_name.replace("_sumrestrictedcaf", "_sumrestricteddecaf")
        if decaf_macro in DECAF_MACROS:
          out_name = out_name + "_" + decaf_macro
        
    return out_name

"""
  Compute number of parts dataset should be split into.
  
    snapshot_id = id of snapshot of dataset
    chunk_size = size of one file (in GB)
    max_singlefile_size = if dataset size is less than this size in GB, only a single file will be made
"""
def CalculateNParts(defname, snapshot_id, chunk_size=2, max_singlefile_size=20, requested_num_parts=None):
    if snapshot_id:
        dimensions = "snapshot_id=%d" % snapshot_id
    else:
        dimensions = "dataset_def_name_newest_snapshot %s" % defname
    summary = SAM.listFilesSummary(dimensions=dimensions)
    size = summary["total_file_size"]
    n_files = summary["file_count"]
    
    if requested_num_parts:
        if requested_num_parts <= n_files:
            return requested_num_parts
        else:
            print >> sys.stderr, "You requested %d parts, but there are only %d files in the dataset.  Can't proceed." % (requested_num_parts, n_files)
            sys.exit(1)
    
    GBs = size / 1024**3  # bytes to gigabytes
    if GBs < max_singlefile_size:
        return 1
    else:
        return GBs/chunk_size

def CheckProjects(dataset):
    already_running_projects = SAM.listProjects(name='%%_proj_concat_%s%%' % dataset, state='running')
    if len(already_running_projects) + nparts > MAX_N_JOBS:
        print >> sys.stderr, "Current concat request would require %d jobs, but" % nparts
        if len(already_running_projects) > 0:
            print >> sys.stderr, "there are already %d SAM projects running for this definition," %  len(already_running_projects)
            print >> sys.stderr, "and ",
        print >> sys.stderr, "there is a maximum of %d running concat projects for the same definition." % MAX_N_JOBS
        print >> sys.stderr, "Please retry using --parts to restrict how many jobs are submitted."
        sys.exit(1)


def CheckProxy(role):
    # now check we have a usable proxy
    arg = "-c"
    if role == "Production":
        arg += "p"
    try:
        subprocess.check_call(["setup_fnal_security", arg])
    except subprocess.CalledProcessError:
        print >> sys.stderr, "Please run 'setup_fnal_security' with the correct options to get a valid proxy and try again."
        sys.exit(2)

def MakeOutputDir(dirname):
    try:
        if not os.path.isdir(dirname):
            os.makedirs(dirname)
        os.chmod(dirname, 0775)
    except OSError:
        print >> sys.stderr, "Output directory '%s' cannot be created!" % dirname
        sys.exit(1)

def ParseArgs():    
    parser = argparse.ArgumentParser(description="Submit jobs for concatentation of (de)CAFs.")
    
    parser.add_argument("dataset",
                        metavar="DATASET",
                        help="Dataset containing (de)CAFs to concatenate",
    )
    parser.add_argument("outputdir",
                        metavar="OUTPUTDIR",
                        help="Where to write output",
    )
    parser.add_argument("release",
                        metavar="RELEASE",
                        help="Version of novasoft to use",
    )
    
    # support old (positional) syntax as well as nicer (option-based) one
    gp = parser.add_mutually_exclusive_group()
    gp.add_argument("decaf_positional",
                    metavar="DECAF",
                    nargs="?",
                    default="$NOVAGRIDUTILS_DIR/bin/concat_dataset.C",
                    help="Decaffing macro to run.  Either specify full path here or one of the following shortcut options: %s.  Can also specify as --decaf (below)." % DECAF_MACROS.keys(),
    )
    gp.add_argument("--decaf", "--decaf_macro",
                    metavar="DECAF",
                    help="Alternate form of positional DECAF (above).",
    )
    
    # etc.
    gp = parser.add_mutually_exclusive_group()
    gp.add_argument("role_positional",
                    metavar="ROLE",
                    nargs="?",
                    default="Analysis",
                    choices=["Analysis", "Production"],
                    help="VOMS role to submit with.  (You must have Production-eligible proxy to use the Production role.)  Can also specify as --role or --voms_role (below).",
    )
    gp.add_argument("--role", "--voms_role",
                    choices=["Analysis", "Production"],
                    help="Alternate form of positional ROLE (above).",
    )
     
    gp = parser.add_mutually_exclusive_group()
    gp.add_argument("nparts_positional",
                    nargs="?",
                    type=int,
                    metavar="NPARTS",
                    help="Divide concats into this many files.  If PARTS is also specified, NPARTS must correspond exactly to the number of parts.  Can also specify as --nparts (below).",
    )
    gp.add_argument("--nparts",
                    metavar="NPARTS",
                    type=int,
                    help="Alternate form of positional NPARTS (above).",
    )
    
    gp = parser.add_mutually_exclusive_group()
    gp.add_argument("parts_positional",
                    nargs="?",
                    metavar="PARTS",
                    help="The indices of which concat files to process.  Optional form '--parts' (see description below) is preferred. ",
    )
    gp.add_argument("--parts",
                    metavar="PARTS",
                    action="append",
                    help="The indices of which concat files to process.  Can separate multiple items with commas, or specify --parts multiple times." \
                         +"  Ranges can be specified using a hyphen.  So, '--parts 3,5,7-9 --parts 11-15 --parts 27' is legal.",
    )
    
    # allow user to opt out of 'caf' -> 'decaf' renaming in case we don't automatically get it right
    parser.add_argument("--no_caf_rename",
                    action="store_true",
                    default=False,
                    help="Do not substitute 'decaf' for 'caf' in the output name.  Useful if you're only concatting CAFs with some script other than concat_dataset.C.",
    )
    
    parser.add_argument("--use_last_snapshot",
                    action="store_true",
                    default=False,
                    help="Use the last snapshot of the definition rather than taking a new one.",
    )

    parser.add_argument("--flatten",
                    action="store_true",
                    default=True,
                    help="Flatten concat file",
    )
    
    parser.add_argument("--flat_only",
                    action="store_true",
                    default=False,
                    help="Produce only flattened concat file",
    )

    gp = parser.add_mutually_exclusive_group(required=False)
    gp.add_argument('-t', '--testrel', metavar='DIR', 
                            help='Use a test release at location TESTREL. It will be tarred up, and sent to the worker node. (Conflicts with --user_tarball)',
                            default=None)
    
    gp.add_argument("--user_tarball",
                            help="Use existing test release tarball in specified location rather than having jobsub make one for you (conflicts with --testrel, and is redunant with --reuse_tarball)",
                            type=str)

    gp.add_argument('--reuse_tarball',
                            help='Do you want to reuse a tarball that is already in resilient space? If using this option avoid trailing slash in --testrel option.  (redundant with --user_tarball)',
                            action='store_true',default=False)
    
    parser.add_argument("-n", "--dry-run", "--test-only",
                    action="store_true",
                    default=False,
                    help="Test only; don't actually submit any jobs.  Print jobsub command and exit."
    )
    
    
    
    return parser.parse_args()

def ParseParts(parts_string):
    parts = set()
    for tok in parts_string.split(","):
        if "-" in tok:
            start, end = (int(t) for t in tok.split("-"))
            parts |= set(range(start, end+1))
        else:
            parts.add(int(tok))
    
    return parts

def Submit(args):
    # can be used for interactive testing if needed
#     cmd = "$NOVAGRIDUTILS_DIR/bin/concat_job_project.sh" \
   

    decafscript = args["decaf_macro"]

    cmd = "jobsub_submit -g -N%(nparts)d --role=%(role)s " 

    # increase resource requirements for flattening
    if args["flatten"]:
        cmd += " --memory=1999MB --disk=5GB --expected-lifetime=23h"
    else:
        cmd += " --memory=2400MB --expected-lifetime=23h"

    # files on blue arc are not accessible directly, we need to copy them to the grid nodes
    if "/nova/" in os.path.expandvars(args["decaf_macro"]):
        print "decaf macro has a blue arc location, copying it to grid nodes"
        # copy it at submission time
        cmd += " -f dropbox://%(decaf_macro)s"
        # remove path string containing /nova/ for concat_job_project.sh
        decafscript = args["decaf_macro"].split('/')[-1]
    
    if args["testrel"]:
        cmd += " --tar_file_name " + args["testrel"]

    cmd += " file://$NOVAGRIDUTILS_DIR/bin/concat_job_project.sh" \
           + " --dataset %(dataset)s --outname %(outname)s --snapid %(snapshot_id)s" \
           + " --nparts %(nparts)d --release %(release)s --outdir %(outputdir)s"
    
    cmd %= args
    cmd += " --decafscript %s " % decafscript
    
    if args["testrel"]:
        cmd += " --testrel"
    if len(args["parts"]):
        cmd += " --parts " + ",".join(str(a) for a in args["parts"])
    if args["flatten"]:
        cmd += "--flatten"
    if args["flat_only"]:
        cmd += "--flat_only"

    cmd = os.path.expandvars(cmd)
    print cmd
    if not args["dry_run"]:
        subprocess.check_call(shlex.split(cmd))

#####################################

if __name__ == "__main__":
    args = ParseArgs()

    # get path to test release
    testrel = ""
    if args.testrel:
        if args.reuse_tarball: 
            testrel += 'dropbox://'+os.path.basename(args.testrel)+'.tar'
        else:
            testrel += 'tardir://'+args.testrel+'\\\n'
    elif args.user_tarball:
        if not os.path.isfile(args.user_tarball):
            print "Tarball filename passed to --user_tarball does not exit:", args.user_tarball
            sys.exit(1)
        testrel += 'dropbox://' + args.user_tarball + ' \\\n'

    # sanity checking
    if args.flat_only and not args.flatten:
        print "Need to provide a --flatten option if you need a flattened concat!"
        sys.exit(1)

    release_path = os.path.expandvars(os.path.join(CVMFS_RELEASE_DIR, args.release))
    if not os.path.isdir(release_path):
#        print release_path
        print >> sys.stderr, "Invalid release specified:", args.release
        sys.exit(1) 
    
    out_name = BuildOutputNameStub(args.dataset, args.decaf or args.decaf_positional, rename_cafs=not args.no_caf_rename)
    
    # compute the PARTS if commas or ranges
    parts = set()
    parts_args = args.parts or args.parts_positional
    if parts_args:
        parts_joined = "".join(parts_args)
        try:
            parts = ParseParts(parts_joined)
        except ValueError as e:
#            raise
            print e
            print "Invalid value specified for --parts:", parts
            sys.exit(1)

    CheckProxy(args.role or args.role_positional)

    SAM = samweb_client.SAMWebClient(experiment='nova')
    if args.use_last_snapshot:
        snapshot_id = None
    else:
        try:
            snapshot_id = int(SAM.takeSnapshot(args.dataset))
            print "Took snapshot %d for defn %s" % (snapshot_id, args.dataset)
        except samweb_client.exceptions.Error as e:
            print >> sys.stderr, e
            sys.exit(1)

    nparts = CalculateNParts(args.dataset, snapshot_id, requested_num_parts=(args.nparts or args.nparts_positional))
    if any(p > nparts-1 for p in parts):
        print >> sys.stderr, "One or more requested parts (you requested: %s) is larger than the number of parts (%d)!" % (list(parts), nparts)
        sys.exit(1)
    print "Requires submitting", nparts, "jobs."

    # don't think we need this anymore
    #  CheckProjects(args.dataset)

    # check if this is a shortcut string
    concat = args.decaf or args.decaf_positional
    if concat in DECAF_MACROS:
        concat = DECAF_MACROS[concat]
    print concat

    MakeOutputDir(args.outputdir)
    print "Output sent to:", args.outputdir

    print
    if args.dry_run:
        print "Invoked with --dry-run, so not submitting.  But I would have used the following jobsub command:"
    else:
        print "Submitting:"
    Submit({
        "nparts":      nparts,
        "role":        args.role or args.role_positional,
        "dataset":     args.dataset,
        "outname":     out_name,
        "snapshot_id": snapshot_id or "latest",
        "release":     args.release,
        "testrel":     testrel,
        "outputdir":   args.outputdir,
        "decaf_macro": concat,
        "parts":       parts,
        "dry_run":     args.dry_run,
        "flatten":     args.flatten,
        "flat_only":   args.flat_only
    })