kcms-create-jobs

#!/usr/bin/env python

import sys, os
import stat, shutil
from math import ceil
from datetime import datetime

def usage():
    print sys.argv[0], " : create pbs jobs"
    print "  Mandatory options :"
    print "   --jobName  NAME                  Name of job"
    print "   --dataset  DATASET_NAME          Dataset name, it should be present in UOS T3"
    print "   --maxFiles N                     Maximum number of files per job"
    print "   --cfg      CONFIG_FILE_cfg.py    Configuration file"
    print "  Optional :"
    print "   -n                               Do not submit jobs to batch"
    sys.exit()

# Parse arguments
if len(sys.argv) < 2: usage()
opts = {}
args = []
i = 0
while i < len(sys.argv)-1:
    i += 1
    arg = sys.argv[i]

    arg1, arg2 = None, None
    if len(arg) > 2 and arg[:2] == '--':
        if '=' in arg:
            arg1 = arg.split('=')[0]
            arg2 = '='.join(arg.split('=')[1:])
        elif i+1 < len(sys.argv):
            arg1 = arg
            arg2 = sys.argv[i+1]
            i += 1
        else:
            print "ERROR: Value must be followed after the option", arg
            sys.exit()
        opts[arg1] = arg2
    elif arg[0] == '-':
        opts[arg] = True
    else:
        args.append(arg)

## Mandatory options
if '--jobName' not in opts: usage()
uniqueID = datetime.strftime(datetime.now(), "%Y%m%d-%H%M%S")
jobName = opts['--jobName']+"-"+uniqueID
if '--dataset' not in opts: usage()
datasetName = opts['--dataset']
if '--maxFiles' not in opts: usage()
maxFiles = int(opts['--maxFiles'])
if '--cfg' not in opts: usage()
cfgFileName = opts['--cfg']
if '_cfg.py' != cfgFileName[-min(7,len(cfgFileName)):]: usage()
if not os.path.exists(cfgFileName):
    print "ERROR: Cannot find config file", cfgFileName
    sys.exit()
if os.path.isdir(jobName):
    print "ERROR: Output directory already exists."
    sys.exit()
doSubmit = ('-n' not in opts)

# Check necessary envvars
if 'CMSSW_VERSION' not in os.environ:
    print "ERROR: You have to set up CMSSW. Run cmsenv in your working directory."
    sys.exit()

## Prepare working directory
print "@@ Preparing batch jobs in", jobName, "..."
os.makedirs(jobName)
scratchPath = "/pnfs/user/%s/scratch/%s" % (os.environ["USER"], jobName)
print "@@ Preparing output scratch directory in", scratchPath, "..."
os.makedirs(scratchPath)

## Get file list from DBS
print "@@ Retriving file list of dataset..."
dbsURL = 'http://cmsdbsprod.cern.ch/cms_dbs_ph_analysis_01/servlet/DBSServlet'
query = "find file where dataset = %s" % datasetName
os.system("dbs search --url='%s' --query='%s' --noheader > %s/files.txt" % (dbsURL, query, jobName))
## Collect root files
files = []
for f in open("%s/files.txt" % jobName).readlines():
    f = f.strip()
    if len(f) < 5: continue
    if '#' == f[0] or '.root' != f[-5:]: continue
    files.append(f)
nFiles = len(files)
if nFiles == 0:
    print "ERROR: Empty dataset."
    sys.exit()
## Sort files by its key "i", filename_[i]_[j]_hash.root
files.sort(key=lambda f: int(f.split('/')[-1].split('_')[1]))

## Load cfg file
print "@@ Loading python cfg..."
sys.path.append('.')
cout = sys.stdout
sys.stdout = open("%s/log.txt" % jobName, "w")
process = __import__('.'.join(cfgFileName.split('.')[:-1])).process
sys.stdout = cout

## Memorise to modify output file names
print "@@ Setting output modules..."
outFileModes = {}
if hasattr(process, 'TFileService'):
    outFileModes['TFileService'] = process.TFileService.fileName.value()
for modName in process.outputModules_():
    outFileModes[modName] = getattr(process, modName).fileName.value()

## Split files into jobs and write python cfg
print "@@ Splitting jobs..."
nSection = int(ceil(1.0*nFiles/maxFiles))
for section in range(nSection):
    begin = section*maxFiles
    end = min(begin+maxFiles, nFiles)

    process.source.fileNames = files[begin:end]
    for modName in outFileModes:
        getattr(process, modName).fileName = "%s/%s_%03d.root" % (scratchPath, outFileModes[modName][:-5], section)

    cfgFileName = "%s/ntuple_%03d_cfg.py" % (jobName, section)
    open(cfgFileName, "w").write(process.dumpPython())

## Checking voms proxy
print "@@ Checking grid certificate to access files..."
if os.system("voms-proxy-info -exists --hours 8") != 0:
    os.system("voms-proxy-init -voms cms --valid 144:00")
uid = os.getuid()
shutil.copy2("/tmp/x509up_u%d" % uid, "%s/x509up_u%d" % (jobName, uid))

## Write run script
print "@@ Writing run script..."
runFileName = "%s/run.sh" % jobName
fout = open(runFileName, "w")
print>>fout, """#!/bin/bash
SECTION=`printf %%03d $PBS_ARRAYID`
cd %s
eval `scram runtime -sh`
cd %s
cp -f x509up_u%d /tmp/
echo BEGIN `date` cmsRun ntuple_${SECTION}_cfg.py >> %s/submit.log
time cmsRun ntuple_${SECTION}_cfg.py &> %s/ntuple_${SECTION}.log
if [ $? == 0 ]; then
    echo ENDED `date` cmsRun ntuple_${SECTION}_cfg.py >> %s/submit.log
fi""" % (
    os.environ["CMSSW_BASE"], os.path.abspath(jobName), uid,
    os.path.abspath(jobName), scratchPath, os.path.abspath(jobName))
fout = None
os.chmod(runFileName, os.stat(runFileName).st_mode | stat.S_IEXEC)

print "@@ Writing submit script..."
submitFileName = "%s/submit.sh" % jobName
fout = open(submitFileName, "w")
print>>fout, """#!/bin/bash
qsub -N %s -d `pwd` -q batch -t 0-%d run.sh""" % (jobName, nSection-1)
fout = None
os.chmod(submitFileName, os.stat(submitFileName).st_mode | stat.S_IEXEC)
if doSubmit:
    print "@@ Submitting jobs..."
    os.system("ssh uosaf0008 'cd %s;./submit.sh'" % os.path.abspath(jobName))
else:
    print "@@ Jobs are prepared. You can submit jobs with following command:"
    print "ssh uosaf0008 'cd %s;./submit.sh'" % os.path.abspath(jobName)

print "@@ Done."