submit_cafana.py

Go to the documentation of this file.

#!/usr/bin/env python

from __future__ import print_function
from builtins import str
from builtins import range
import argparse
import os
import sys
import string
import re
from glob import glob

# Various utility functions
import NovaGridUtils
from NovaGridUtils import warn, fail, sleep
import recommended_sites as rs

def remove_comments(src):
    if '#' not in src: return src
    return src[0:src.find('#')]

def test_not_dcache(l, warnOnly = False):
    loc = os.path.expandvars(l)
    for bad in ['/nova/app/', '/nova/ana/', '/nova/data/', '/grid']:
        if loc.startswith(bad):
            txt = "Location %s cannot be on %s it must be in dCache /pnfs/nova/" % (loc, bad)
            if loc.startswith('/nova/app'):
                txt = "Jobs can no longer access BlueArc directly. Test releases will be tarred up and sent to worker nodes, however input files should be moved to dCache."
            if warnOnly:
                warn(txt)
            else:
                fail(txt)


def find_file(paths, filename):
    if os.path.isfile(filename):
        return filename
    for path in paths:
        for root, dirs, files in os.walk(os.path.expandvars(path)):
            if filename in files:
                return os.path.join(root, filename)
    fail("Cannot find file "+filename)


def find_file_in_list(filepath, pathlist):
    for testpath in pathlist:
        if os.path.basename(filepath) == os.path.basename(testpath):
            return True
    return False


if __name__ == '__main__':

    input_files=[]

    NovaGridUtils.prog = 'submit_cafana.py' # for error messages etc

    # Start this parser up just to extract the -f or --file argument if it's
    # there, and push all of those args into the command line for the real
    # parser to deal with.
    if "-f" in sys.argv or "--file" in sys.argv:
        prelim_parser = argparse.ArgumentParser()
    
        prelim_parser.add_argument('-f', '--file', type=str, action='append')
        pre_args, unknown = prelim_parser.parse_known_args()

        if pre_args.file:
            for filepath in pre_args.file:
                NovaGridUtils.check_file(filepath)
                for line in open(filepath, 'r'):
                    sys.argv += remove_comments(line).split()



    parser = argparse.ArgumentParser(description = 'Submit a CAFAna macro with datasets split between N jobs')

    ###required options
    # So it'll show up in the help text, but it was already handled above
    parser.add_argument('-f', '--file', type=str, 
                        help='Text file containing any arguments to this utility',
                        action='append')

    parser.add_argument('-n', '--njobs', type=int, required=True, metavar='N', 
                        help='Number of grid processes')

    parser.add_argument('-r', '--rel', required=True, 
                        help='Release to use')

    parser.add_argument('-i', '--input_file', action="append", 
                        help="Copy this input file to work area on worker node")

    parser.add_argument('-o', '--outdir', required=True, metavar='DIR', 
                        help='Directory output files will go to')

    parser.add_argument('-ss', '--snapshot', action='store_true',
                        help='Use latest snapshot instead of requerying')

    parser.add_argument('-off', '--offsite', action='store_true',
                        help='Run this cafana job offsite')

    parser.add_argument('-d',  '--drain', action='store_true',
                        help='Recover files missing from your output directory')

    parser.add_argument('-x',  '--xrootdebug', action='store_true',
                        help='Add extra xrootd debugging information',
                        default=False)

    parser.add_argument('--numuccinc', action = 'store_true',
                        help = 'Load libraries for specific xsec analysis')
    parser.add_argument('--numubarccinc', action = 'store_true',
                        help = 'Load libraries for specific xsec analysis')
    parser.add_argument('--numucc2p2h', action = 'store_true',
                        help = 'Load libraries for specific xsec analysis')
    parser.add_argument('--numucc0pi', action = 'store_true',
                        help = 'Load libraries for specific xsec analysis')
    parser.add_argument('--nuebarccinc', action = 'store_true',
                        help = 'Load libraries for specific xsec analysis')
    parser.add_argument('--numubarccpi0', action = 'store_true',
                        help = 'Load libraries for specific xsec analysis')
    parser.add_argument('--no3flavor', action = 'store_true',
                        help = 'Do *not* load libraries for 3-flavor analysis')
    parser.add_argument('--nuxana', action="store_true",
                        help="Load libraries for NuX analysis")

    parser.add_argument('macro.C',
                        help='The CAFAna macro to run')

    parser.add_argument('args', nargs='*',
                        help='Arguments to the macro')

    ###Define job_control_args
    job_control_args = parser.add_argument_group("Job control options", "These optional arguments help control where and how your jobs land.") 
    ###tarball control options.
    tarball_gp = job_control_args.add_mutually_exclusive_group(required=False)
    tarball_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)
    
    tarball_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)

    ###general job control
    job_control_args.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)

    job_control_args.add_argument('--dedicated',
                        help='Only run on dedicated nodes on fermigrid (default is to run opportunistically)',
                        action='store_true',default=False)

    job_control_args.add_argument('--site',
                        help='Specify allowed offsite locations.  Omit to allow running at any offsite location',
                        type=str,action='append')

    job_control_args.add_argument('--exclude_site',
                        help='Specify an offsite location to exclude.',
                        metavar='SITE',
                        type=str,action='append')

    job_control_args.add_argument('--recommended_sites',
                        help='Specify known working offsite locations.',
                        action='store_true',default=False)

    job_control_args.add_argument('--disk',
                        help='Local disk space requirement for worker node in MB (default is 2000MB).',
                        type=int, default=2000)

    job_control_args.add_argument('--memory',
                        help='Local memory requirement for worker node in MB (default is 1900MB).',
                        type=int, default=1900)

    job_control_args.add_argument('--lifetime',
                        help='Expected job lifetime. Valid values are an integer number of seconds. (default is 10800=3h)',
                        type=int, default="10800")

    job_control_args.add_argument('--source',
                        help='Source script SOURCE:par1:par2:..',
                        type=str, action='append')

    job_control_args.add_argument('--ngu_test',
                                  help='Setup the test version of NovaGridUtils in the grid jobs.',
                                  action='store_true')

    job_control_args.add_argument('--ngu_version',
                                  help='Setup a specific NovaGridUtils version in the grid jobs.', metavar='VERSION',
                                  type=str)

    job_control_args.add_argument('--testrel_ngu',
                                  help="Must be used with --testrel, with NGU checked out. After unpacking tarball will setup the local version of NGU you are using on the work.",
                                  action='store_true')

    job_control_args.add_argument('-ep', '--extproduct', type=str, metavar='PRODUCT:VERSION',
                        help='Setup this external product on the worker node in format <product>:<version>',
                        default='NONE')

    job_control_args.add_argument("--mail_always",
                        help="Do you want an email whenever every jobs finishes?",
                        default=False, action="store_true")

    job_control_args.add_argument("--mail_on_error",
                        help="Do you want an email whenever a job fails on an error?",
                        default=False, action="store_true")

    # --singularity and --os are mutually exclusive
    sing_group = job_control_args.add_mutually_exclusive_group()
    sing_group.add_argument('--singularity',
                            help='Location in CVMFS of a singularity container to launch the job into. ' +\
                                 "If you're looking for a stock Scientific Linux image, try --os instead. " +\
                                 "(This option is mutually exclusive with --os.)",
                            type=str,
                            default=None)
    sing_group.add_argument("--os",
                            help="Run this job inside a Scientific Linux singularity image with given OS. " +\
                                 "(This option is mutually exclusive with --singularity.)",
                            choices=NovaGridUtils.KNOWN_OS_CONTAINERS)

    job_control_args.add_argument('--jobfactory',
                        help='Use the specified JobFactoryType. Only use with --singularity',
                        default="")

    job_control_args.add_argument("--gpu",
                        help="Request a node with a GPU",
                        default=False, action="store_true")

    job_control_args.add_argument("--export",
                        help='Export variable EXPORT to jobsub_submit',
                        type=str, action='append')

    job_control_args.add_argument('-c', 
                        help='Append Condor requirements',
                        type=str, action='append')

    job_control_args.add_argument('--cpu',
                        help='Request worker nodes to have at least NUMBER cpus',
                        type=int)

    NovaGridUtils.add_node_features_arg(job_control_args)

    ###debugging
    debugging_args = parser.add_argument_group("Debugging options", "These optional arguments can help debug your submission.")
    
    debugging_args.add_argument('--print_jobsub',
                        help='Print jobsub command',
                        action='store_true',default=False)

    debugging_args.add_argument('--test',
                        help='Do not actually do anything, just run tests and print jobsub cmd',
                        action='store_true',default=False)

    debugging_args.add_argument('--ifdh_debug', 
                        help='Verbose output for pinning down IFDH/dCache issues',
                        action='store_true', default=False)

    opts = parser.parse_args()

    macro = os.path.abspath(vars(opts)['macro.C'])

    # currently can't non-tarred files (such as macros) through new jobsub release. 
    # Until there's a patch, catch that case and setup old jobsub 
    if (opts.testrel == None and opts.user_tarball == None):
      import subprocess
      jobsub_client_ver = subprocess.check_output ("ups active | grep jobsub_client", shell=True)
      jobsub_client_ver = jobsub_client_ver.split()[-5]
      if (jobsub_client_ver == "v1_3_3"):
        print("WARNING: Can't pass macro with no testrel in jobsub_client v1_3_3, run") 
        print("           unsetup jobsub_client; setup jobsub_client v1_3_2_1")
        print("         and resubmit")
        sys.exit(1)

    # Some sanity checks
    NovaGridUtils.get_credentials('Analysis')
    rel = NovaGridUtils.check_tag(opts.rel)
    NovaGridUtils.check_file(macro)
    test_not_dcache(opts.outdir)
    NovaGridUtils.check_dir(opts.outdir)
    NovaGridUtils.check_is_group_writable(opts.outdir)
    if opts.testrel:
        test_not_dcache(opts.testrel, True)
        NovaGridUtils.check_dir(opts.testrel)
        if not os.path.isdir(opts.testrel+'/lib/'+os.getenv('SRT_ARCH')+'-GCC-maxopt'):
            NovaGridUtils.fail(opts.testrel+' has never been built maxopt')
    if opts.user_tarball:
        test_not_dcache(opts.user_tarball, True)

    # If draining, check what is missing in the output dir
    draincmd = ""
    if opts.drain:
        outfiles = os.listdir(opts.outdir)
        completejobs = set()
        missingjobs = []
        for filename in outfiles:
            # Log files aren't a sign of success of a job
            if filename.endswith('.log.txt'): continue
            # Nor are core files
            if filename.startswith('core.'): continue

            # Any other type of file output will be counted as a success
            regex = re.compile(".*?(\d+)_of_(\d+)\..*")
            reresult = regex.match(filename)
            if reresult:
                thisjob = int(reresult.group(1))
                thismax = int(reresult.group(2))
                completejobs.add(thisjob)
                if opts.njobs != thismax:
                    print("You specified {} total jobs, but output files have the form x_of_{}. These must be consistent! Bailing.".format(opts.njobs, thismax))
                    sys.exit(5)
        for jobno in range(1,opts.njobs+1):
            if jobno not in completejobs:
                missingjobs.append(jobno)
        if len(missingjobs) == 0:
            print("No jobs to drain. Exiting.")
            sys.exit(0)
        else:
            print("Found jobs to drain:")
            print("   ",missingjobs)
            if opts.njobs != len(missingjobs):
                print("    Adjusting njobs from %i to %i to match." % (opts.njobs, len(missingjobs)))
            draincmd = '      '
            for jobno in missingjobs:
                draincmd += '-d %i ' % jobno
            draincmd += '\\\n'

    # Singularity, offsite checks
    if opts.os:
        # just redirect to the known container path
        opts.singularity = NovaGridUtils.KNOWN_OS_CONTAINERS[opts.os]

    if opts.singularity:
        if not os.path.exists(opts.singularity):
            fail("Requested singularity image cannot be found: %s" % opts.singularity)
       
    if opts.gpu and not opts.singularity:
        warn("Requested GPU, but did not request singularity. This is not likely to succeed.")

    if opts.gpu and not opts.offsite:
        warn("GPUs are only available offsite, and you have not chosen --offsite")

    grid_script = os.getenv('NOVAGRIDUTILS_DIR')+'/bin/cafe_grid_script.sh'

    NovaGridUtils.check_file(grid_script)

    if opts.input_file:
        for inFile in opts.input_file:
            test_not_dcache(inFile)
        input_files += opts.input_file

    # grab user supplied source scripts
    source_scripts=[]
    if opts.source:
        source_scripts += opts.source

    for script in source_scripts:
        if ":" in script:
            script_path = script.split(":")[0]
        else:
            script_path = script
            
        if not find_file_in_list(os.path.expandvars(script_path), input_files):
            if not find_file(os.environ["PATH"].split(os.pathsep), os.path.expandvars(script_path)):
                fail("Script %s does not exist!" % script_path)

    # Process the job_control_args
    jobsub_opts = ""
    
    #usage model
    usage_models=["DEDICATED"]
    if not opts.dedicated:
        usage_models.append("OPPORTUNISTIC")
    if opts.offsite:
        usage_models.append("OFFSITE")
    jobsub_opts += "    --resource-provides=usage_model=%s \\\n" % (",".join(usage_models))

    #disk
    if opts.disk:
        disk_opt="    --disk=%sMB \\\n" % (opts.disk)
        jobsub_opts += disk_opt        
    #memory
    if opts.memory:
        mem_opt="    --memory=%sMB \\\n" % (opts.memory)
        jobsub_opts += mem_opt

    #cpu 
    if opts.cpu:
       cpu_opt="    --cpu=%s \\\n" %(opts.cpu)
       jobsub_opts += cpu_opt
 
    #extra condor requirements
    if opts.c:
       append_condor_opt = "    -c %s \\\n" % (",".join(opts.c))
       jobsub_opts += append_condor_opt        

    #emails about jobs?
    if opts.mail_always:
        jobsub_opts += "    --mail_always \\\n"
    elif opts.mail_on_error:
        jobsub_opts += "    --mail_on_error \\\n"
    else:
        jobsub_opts += "    --mail_never \\\n"

    if opts.node_features:
        jobsub_opts += "    %s \\\n" % (
            NovaGridUtils.make_jobsub_node_features_arg(opts.node_features)
        )

    # Never kill me for being over time. Find a node with enough time left in
    # its glidein for what I requested, but don't penalize me for going over
    # time.
    life_opt=("    --expected-lifetime=0s \\\n    --append_condor_requirements='(((TARGET.GLIDEIN_ToDie-CurrentTime)>%s)||isUndefined(TARGET.GLIDEIN_ToDie))' \\\n" % opts.lifetime)
    jobsub_opts += life_opt

    # xrootd debugging
    if opts.xrootdebug:
        jobsub_opts += "    -e XrdSecDEBUG=2 \\\n"

    # Singularity
    if opts.singularity:
        jobsub_opts += "    --line='+SingularityImage=\\\"%s\\\"' \\\n" % opts.singularity
        jobsub_opts += "    --append_condor_requirements='(TARGET.HAS_SINGULARITY=?=true)' \\\n"
    if opts.jobfactory:
        jobsub_opts += "    --line='+JobFactoryType=\\\"%s\\\"' \\\n" % opts.jobfactory

    # GPUs
    if opts.gpu:
        jobsub_opts += "    --line='+RequestGPUs=1' \\\n"

    if opts.recommended_sites or opts.site:
        site_opt="    --site="

        if opts.recommended_sites:
            istarball = opts.testrel or opts.user_tarball
            recommended_sites=rs.get_recommended_sites('default', istarball, False)
            for isite in recommended_sites:
                site_opt += isite + ","
        if opts.site:
            for isite in opts.site:
                if isite not in recommended_sites:
                    warn("Site "+isite+" is not known to work. Your jobs may fail at that site. Sleeping for 5 seconds")
                    sleep(5)
                site_opt += isite + ","

        site_opt=site_opt[:-1] + " \\\n"
        jobsub_opts += site_opt

    if opts.exclude_site:
        for isite in opts.exclude_site:
            jobsub_opts += "    --append_condor_requirements='(TARGET.GLIDEIN_Site\ isnt\ \\\"%s\\\")' \\\n" % isite

    #Process the debugging_args
    test=opts.test
    print_jobsub=opts.print_jobsub
    if test :
        print_jobsub=True
        print("")
        warn("--test was specified, so all we do is run checks and print jobsub cmd.")

    if opts.ifdh_debug:
        if os.path.isfile(".rootrc"):
            fail("--ifdh_debug creates a temporary .rootrc file, but there is already one in this directory. Submit from somewhere else.")
        else:
            rc=open(".rootrc","w+")
            rc.write("XNet.UseOldClient: yes\n") # root6 bug with new client, supresses other XNet configs like Debug
            rc.write("XNet.Debug: 2\n")
            rc.write("XNet.MaxRedirectCount: 255\n")
            rc.close()
        source_scripts += ['setup_test_product:NovaGridUtils']
        source_scripts += ['export_var:IFDH_SILENT=0']
        source_scripts += ['export_var:IFDH_DEBUG=1']
        print_jobsub=True

    # Export
    if opts.export:
        for export in opts.export:
            source_scripts += ['export_var:%s' % export]

    cmd = 'jobsub_submit '
    if opts.drain:
        cmd += '-N '+str(len(missingjobs))+' '
    else:
        cmd += '-N '+str(opts.njobs)+' '
    cmd += "\\\n"+jobsub_opts
    cmd += '    -f dropbox://%s \\\n' % (macro)
    for f in input_files:
        cmd += '    -f %s \\\n' % (f)
    if opts.ifdh_debug:
        cmd += '    -f dropbox://%s \\\n' % (os.path.abspath(".rootrc"))

    # Test release options.
    if opts.testrel:
        if opts.reuse_tarball: 
            cmd += '    --tar_file_name dropbox://'+os.path.basename(opts.testrel)+'.tar'
        else:
            cmd += '    --tar_file_name tardir://'+opts.testrel+'\\\n'
    elif opts.user_tarball:
        if not os.path.isfile(opts.user_tarball):
            print("Tarball filename passed to --user_tarball does not exit:", opts.user_tarball)
            sys.exit(5)
        cmd += '    --tar_file_name dropbox://' + opts.user_tarball + ' \\\n'

    # cafe_grid_script.sh and arguments
    cmd += '    file://%s \\\n' % (grid_script)
    for source_script in source_scripts:  
        cmd += '      --source %s \\\n' % (source_script)

    cmd += '      --release %s \\\n' % (rel)
    if opts.testrel or opts.user_tarball: 
        cmd += '      --tarball \\\n'

    # NovaGridUtils options.
    if opts.ngu_test:
        cmd += '      --source setup_test_product:NovaGridUtils \\\n'

    if opts.ngu_version:
        cmd += '      --source setup_product:NovaGridUtils:%s \\\n' %(opts.ngu_version)

    if opts.testrel_ngu:
        cmd += '      --testrel_ngu \\\n'
        
    cmd += '      --macro %s \\\n' % (os.path.basename(macro))
    cmd += '      --outdir %s \\\n' % (opts.outdir)
    cmd += '      --njobs %s \\\n' % (str(opts.njobs))

    if opts.snapshot: cmd += '      --snapshot \\\n'

    if opts.numuccinc:    cmd += ' --numuccinc \\\n'
    if opts.numubarccinc: cmd += ' --numubarccinc \\\n'
    if opts.numucc2p2h:   cmd += ' --numucc2p2h \\\n'
    if opts.numucc0pi:   cmd += ' --numucc0pi \\\n'
    if opts.nuebarccinc: cmd += ' --nuebarccinc \\\n'
    if opts.numubarccpi0: cmd += ' --numubarccpi0 \\\n'
    if opts.no3flavor:   cmd += ' --no3flavor \\\n'
    if opts.nuxana:      cmd += ' --nuxana \\\n'


    if opts.extproduct != 'NONE': cmd += '      --extproduct %s \\\n' % (opts.extproduct)
    if opts.drain:
        if not draincmd:
            print("Jobs to drain, but no draining command. This shouldn't happen.")
            sys.exit(5)
        cmd += draincmd
    for a in opts.args: cmd += ' '+a

    if print_jobsub:
        print(cmd)
        sys.stdout.flush()
        sys.stderr.flush()

    if not test:
        os.system(cmd)

    if opts.ifdh_debug:
        os.remove("./.rootrc")