submit_hadd.py

Go to the documentation of this file.

#!/usr/bin/env python

from __future__ import print_function
import argparse
import subprocess
import datetime
import os
import ifdh
import sys

def split_list(flist):
    return

if __name__ == "__main__":

    if "-f" in sys.argv or "--file" in sys.argv:
        prelim_parser = argparse.ArgumentParser(prog='submit_hadd.py',
        description="""Submit grid jobs to hadd CAF trees or ROOT ntuples. """)
    
        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:
                if not os.path.exists(filepath):
                    print("Config file is not found, exit now.")
                    os.exit(1)
                for line in open(filepath, 'r'):
                    sys.argv += remove_comments(line).split()

    parser = argparse.ArgumentParser(
        prog='sub_hadd',
        description="""Submit grid jobs to hadd CAF trees or ROOT ntuples. """,
        epilog="Questions and Comments to dingpf@fnal.gov")

    # required arguments
    parser.add_argument('--job-name',  help="job name", required=True)
    parser.add_argument('--njobs', help='njobs', required=True)
    parser.add_argument('--outdir', help='output directory', required=True)
    parser.add_argument('--flist', help='list of files to add')
    parser.add_argument('--indir', help='input directory')
    parser.add_argument('--pattern', help='find file pattern')
    parser.add_argument('--release', help='release version', required=True)
    parser.add_argument('--method', default='hadd_cafana',
            choices=['hadd','hadd_caf', 'hadd_cafana'],
            help='sample type')
    parser.add_argument('--job-script', default='NONE',
            help='worker node job script')

    # Optional arguments
    parser.add_argument('--offsite', action='store_true',
            help='run jobs on offsite locations')
    parser.add_argument('--just-say', action='store_true',
            help="do not run command, but just print out what will be done.")
    parser.add_argument('--user', default='NONE', help="user name")
    args = parser.parse_args()

    if args.user == "NONE":
        args.user = os.environ["USER"]
    if args.job_script == "NONE":
        args.job_script = os.environ["NOVAGRIDUTILS_DIR"] + \
                '/bin/hadd_grid_script.sh'

    # determine job name
    ct_stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    job_name = "{}_{}_{}".format(args.user, args.job_name, ct_stamp)

    # make sure paths get translated to absolute path
    outdir = os.path.abspath(args.outdir)
    job_outdir = '{}/{}'.format(outdir, job_name)
    job_script = os.path.abspath(args.job_script)

    # split filelist into small file lists
    list_file_name = "{}.list".format(ct_stamp)
    list_file_url = "http://home.fnal.gov/~{}/".format(args.user) + \
            list_file_name

    if args.indir and args.pattern:
        # make flist on the fly
        cmd = 'find {} -name {}>{};'.format(
                args.indir, args.pattern, list_file_name)
        cmd += ' wc -l {i}; '
        cmd += 'scp {i} {u}@fnalu.fnal.gov:/publicweb/{c}/{u}/;'.format(
                u=args.user, i=list_file_name, c=args.user[0])
        cmd += ' rm -f {}'.format(list_file_name)
        print(cmd)
        print("using find to make flist, and scp file to sever")
    elif args.flist:
        cmd = 'scp {f} {u}@fnalu.fnal.gov:/publicweb/{c}/{u}/{l}'.format(
                f=args.flist, u=args.user, l=list_file_name, c=args.user[0])
        print("scp list file to fermilab central web sever")
    else:
        print("either file list option or a combination of indir and")
        print("patter should be specified!")
        exit(1)

    if not args.just_say:
        subprocess.call(cmd, shell=True)

    # setup ifdh
    if "IFDH_BASE_URI" not in os.environ:
        os.environ["IFDH_BASE_URI"] = "http://samweb.fnal.gov:8480/sam/nova/api"
    i = ifdh.ifdh(os.environ["IFDH_BASE_URI"])

    # setup envs used by jobs script
    envs = ["OUT_DIR", "RELEASE_NAME", "NJOBS", "LIST_FILE_URL", \
            "HADD_METHOD", "GRIDUSER"]
    os.environ["RELEASE_NAME"] = args.release
    os.environ["OUT_DIR"] = job_outdir
    os.environ["NJOBS"] = args.njobs
    os.environ["HADD_METHOD"] = args.method
    os.environ["LIST_FILE_URL"] = list_file_url
    os.environ["GRIDUSER"] = args.user

    env_pass = ""
    for i in envs:
        print("INFO:  {:<15} = {}".format(i, os.environ[i]))
        env_pass += " -e {}".format(i)

    # job submission command
    cmd = 'jobsub_submit -N {}'.format(args.njobs)
    if not args.offsite:
        cmd += ' --resource-provides=usage_model=DEDICATED,OPPORTUNISTIC'
    else:
        cmd += ' --resource-provides=usage_model=DEDICATED,OFFSITE,'
        cmd += 'OPPORTUNISTIC --OS=SL6 --site=BNL,Caltech,SMU_HPC,UCSD,'
        cmd += 'Wisconsin,Cornell -M'
    # cmd += ' -M -f {}'.format(input_tar)
    cmd += env_pass
    cmd += ' -G nova file://{}'.format(job_script)
    print(80*'*')
    print("EXECUTING: {}".format(cmd))

    # run jobs submission command
    if not args.just_say:
        subprocess.call(cmd, shell=True, env=os.environ)