runNovaSAM_cmake.py

Go to the documentation of this file.

#!/bin/env python

import os, sys
import shutil 
import samweb_client, ifdh
import pprint
import argparse
import subprocess
import ROOT
import re
import md5, json

from samweb_client.utility import fileEnstoreChecksum

#  Stop ROOT from complaining about not having dictionaries, 3000 is below kError but above kWarning. 
ROOT.gErrorIgnoreLevel=3000

detectors = dict()
detectors["fd"]   = "fardet"
detectors["nd"]   = "neardet"
detectors["ndos"] = "ndos"

#Necessary for createMetadata function
_lineregex = re.compile(r'\d+: ([^ \t]+)[ \t]+(.+)')

#files 
_skip_pattern = re.compile(r"^.*(hist|RootOutput).*\.root", re.IGNORECASE)
#_skip_pattern = re.compile(r"^.*(RootOutput).*\.root", re.IGNORECASE)

def unCamelCase(s):
    """ Convert CamelCase to camel_case """
    # \B matches an empty string which is NOT at the beginning of a word
    # so requiring this means no _ will be inserted at the start
    return re.sub( r'\B([A-Z])', r'_\1', s).lower()
                

def createMetadata(inFile):
    filename = os.path.basename(inFile)
    filesize = os.path.getsize(inFile)
    
    try:
        data = subprocess.check_output(["sam_metadata_dumper", str(inFile)])
    except subprocess.CalledProcessError:
        print "sam_metadata_dumper failed!"
        return None

    md = {'file_name': filename, 'group': 'nova', 'file_size': filesize, 'file_format': 'root'}

    # An ad-hoc parser for the output
    mode = 'new'
    for line in data.split('\n'):
        m = _lineregex.match(line)
        if mode == 'json':
            if not m and not line.startswith('-'):
                data += line
            else:
                md[key] = json.loads(data)
                mode = 'new'

        if mode == 'new':
            if m:
                key = m.group(1)
                value = m.group(2)

                if '.' not in key:
                    key = unCamelCase(key)
                    
                    # ART does not appear to set the application name
                    # but does provide a process name
                if key == 'process_name' and md.get('application',[]).get('name') is None:
                    key = 'application_name'
                elif key == 'stream_name':
                    key = 'data_stream'
                    
                if key in ('start_date','end_date'):
                    value = long(value)
                    
                if key == 'runs' or key == 'parents':
                    mode = 'json'
                    data = value
                elif key.startswith('application_'):
                    # special handling for applications
                    if 'application' not in md:
                        md['application'] = {}
                    md['application'][key[12:]] = value
                else:
                    md[key] = value
                    
    md['crc'] = fileEnstoreChecksum(inFile)
    
    return md

def getOutDir(pathname, hashDirs=False):
    if "DEST" in os.environ:
        dest = os.environ["DEST"]
    else:
        print "No destination set!"
        exit(1)
    dirs = [dest]
    if hashDirs:
      head, tail = os.path.split(pathname)
      hash = md5.new(tail)
      dirs += list(hash.hexdigest()[:3])
    return os.path.join(*dirs)

def checkAndMoveFiles(inFile, declareFiles):
    """Checks all root files in the current directories for zombie or recovered status. Bad files are deleted while good files are moved to the results subdirectory for copy out ease"""
    inFileBase = os.path.basename(inFile)
    baseDir = "."
    for root, dirs, filenames in os.walk(baseDir):
        if root == baseDir:
            for file in filenames:
                if file.endswith (".root") and file != inFileBase:
                    fileWPath = os.path.join(root, file)
                    rootFile = ROOT.TFile(fileWPath)
                    if rootFile.IsZombie() or rootFile.TestBit(ROOT.TFile.kRecovered):
                        #do not copy over, just delete 
                        print "File", fileWPath, "is Zombie - remove it"
                        os.remove(fileWPath)
                    else:
                        if declareFiles == True:
                            declareFile(fileWPath)
                        newFilePath = os.path.join(root, "results", file)
                        os.renames(fileWPath, newFilePath)
                    rootFile.Close()
    return

def declareFile(fileWPath):
    """Checks file for a TKey of RootFileDB.  If it exists, run sam_metadata_dumper and construct appropriate metadata for the file. Use that metadata to declare the file to SAM"""
    samweb = samweb_client.SAMWebClient(experiment='nova')
    rootFile = ROOT.TFile(fileWPath)
    if rootFile.FindKey("RootFileDB"):
        print "Found key (RootFileDB): constructing metadata for", fileWPath
        md = createMetadata(fileWPath)
        if (md == None):
            print "No metadata found!"
        else:
            pprint.pprint(md)
            print ''
            print "Declaring", fileWPath, "to SAM"
            try:
                samweb.declareFile(md)
            except:
                print fileWPath, "already exists in SAM"
    else:
        print fileWPath, "does not contain RootFileDB, do not try to declare"
    rootFile.Close()
    return

def copyOutFiles(hashDirs=False):
    """Builtin facility to copy out art files (automatically excludes hist files). This adds in a subdirectories with a single hex digit each corresponding to the first three digits in the hash of the output file name. This splits up files into 4096 separate subdirectories, preventing overfull directories. Copy out does not happen if the file already exists in the output"""
    print "Copying out files"
    dh = ifdh.ifdh("http://samweb.fnal.gov:8480/sam/nova/api")
    baseDir = "./results"
    for root, dirs, filenames in os.walk(baseDir):
        if root == baseDir:
            for file in filenames:
                if file.endswith (".root") and file != inFileBase:
                    fileWPath = os.path.join(root, file)
                    dest = getOutDir(file, hashDirs)
                    skip_match = _skip_pattern.match(file)
                    if skip_match == None:
                        try:
                            print "Checking if", dest + "/" + file, "exists"
                            dh.ls(dest + "/" + file, 1, "")
                        except RuntimeError:
                            print "It doesn't - copy", fileWPath, "to", dest
                            dh.cp(["-D", fileWPath, dest])
                    print "Removing", fileWPath
                    os.remove(fileWPath)
    dh.cleanup()
    return

def getOutputFileName(detector, runNum, subNum, streamEntry, release, subversion, dataFlag, tier, nevt):
  if dataFlag == "sim":
    pattern = re.compile(r"^.*?_(.*?)_\d+.*_v\d+_(.*?)(\.|\_)sim\..+")
    match = pattern.match(inFileBase)
    if match != None:
      simLabel = match.groups()[0]
      uniq     = match.groups()[1]
      outName = '{0}_{1}_{2}_r{3:08d}_s{4:02d}_{5}_v{6}_{7}_{8}.{9}.root'.format(detector, simLabel, nevt, runNum, subNum, release, subversion, uniq, dataFlag, tier)
    else:
      print "regex couldn't find the sim label and the uniqueifier"
      exit(1)
  else:
    outName = '{0}_r{1:08d}_s{2:02d}_{3}_{4}_v{5}_{6}.{7}.root'.format(detector, runNum, subNum, streamEntry, release, subversion, dataFlag, tier)
  return outName
  
def resolveFclPath(fcl):
   # Check if we have an absolute path name, return it if so
  if fcl[0] == "/":
    return fcl 

  # SRT sets up FHICL_FILE_PATH with extra / at end of directories;
  # cmake does not. Need to add + "/" + below to account for this.
  fclPaths = os.environ["FHICL_FILE_PATH"].split(":")
  for path in fclPaths:
    if os.path.isfile(path + "/" + fcl):
      return path + "/" + fcl
      
  # If we haven't found it, we have a problem.
  raise IOError(sys.argv[0] + ": config  file "+ fcl+" not found in FHICL_FILE_PATH: $FHICL_FILE_PATH")
    

if __name__=='__main__':
    parser = argparse.ArgumentParser(description='Run the nova command using SAM metadata')
    parser.add_argument('inFile', help='The input file to run over', type=str)
    parser.add_argument('--config', '-c', help='FHiCL file to use as configuration for nova executable', type=str)
    parser.add_argument('--outTier', help='Data tier of the output file, multiple allowed, formatted as <name_in_fcl_outputs>:<data_tier>', type=str, action='append')
    parser.add_argument('--cafTier', help='Module label for CAF output, multiple allowed.  Format as <cafmaker_module_label>:<data_tier>', type=str, action='append')
    parser.add_argument('-n', help='Number of events to run over', type=int)
    parser.add_argument('--copyOut', help='Use the built in copy out mechanism', action='store_true')
    parser.add_argument('--hashDirs', help='Use hash directory structure in destination directory.', action='store_true')
    parser.add_argument('--declareFiles', help='Declare files with metadata on worker node', action='store_true')
    parser.add_argument('--noCleanup', help='Skip working directory cleanup step, good for interactive debugging', action='store_true')
    
    args = parser.parse_args()

    inFile = args.inFile
    
    samweb = samweb_client.SAMWebClient(experiment='nova')

    inFileBase = os.path.basename(inFile)

    metadata = samweb.getMetadata(inFileBase)

    #make all lower case keys
    metadata = {k.lower():v for k,v in metadata.items()}

    print "Input File = ", inFile, "with metadata:"
    pprint.pprint(metadata)
    print ''

    #if new format "runs" metadata, use that, otherwise, get from file name
    if "runs" in metadata and len(metadata["runs"][0]) == 3:
        runNum = int(metadata["runs"][0][0])
        subNum = int(metadata["runs"][0][1])
    else:
        run_pattern = re.compile(r"^.*?_r([0-9]+)_s([0-9]+).+")
        run_match = run_pattern.match(inFileBase)
        if run_match != None:
            runNum = int(run_match.groups()[0])
            subNum = int(run_match.groups()[1])
        else:
            print "No run number/subrun number found!"
            exit(1)
            
    print "Run Number:    ", runNum
    print "Subrun Number: ", subNum

    if "data_stream" in metadata:
        try:
            stream = int(metadata["data_stream"])
            streamEntry = 't{0:02d}'.format(stream)
        except ValueError:
            stream = metadata["data_stream"]
            streamEntry = stream
    else:
        print "No data stream found!"
        exit(1)

    #Check for a non-standard data stream in the file name

    stream_pattern = re.compile(r"^.*?_(dd.*?)[_,.].+", re.IGNORECASE)
    stream_match = stream_pattern.match(inFileBase)
    if stream_match != None:
        stream = stream_match.groups()[0]
        streamEntry = stream
        print "Found match!", stream
    else:
        print "No match, using", stream

    if "online.detector" in metadata:
        detector = metadata["online.detector"].lower()
    elif "nova.detectorid" in metadata:
        detId = metadata["nova.detectorid"].lower()
        if detId in detectors.keys():
            detector = detectors[detId] 
        else:
            print "No detector name found for detectorid:", detId
            exit(1)
    else:
        print "No detector name found!"
        exit(1)

    # Changed for cmake. Assumed to have no effect beyond 
    # release entry in metadata
    if "NOVASOFT_VERSION" in os.environ:
        release = os.environ["NOVASOFT_VERSION"]
    else:
        print "No release set!"
        exit(1)

    if "file_type" in metadata:
        file_type = metadata["file_type"]
    else:
        print "No file_type set!"
        exit(1)

    if file_type == "importedDetector":
        dataFlag = "data"
    elif file_type == "importedSimulated":
        dataFlag = "sim"
    else:
        print "Unrecognized file_type:", file_type
        exit(1)

    if "NPASS" in os.environ:
        subversion = os.environ["NPASS"]
    elif "nova.subversion" in metadata:
        subversion = metadata["nova.subversion"]
    else:
        subversion = 1

    nevt = None
    if "simulated.number_of_spills" in metadata:
        nevt = metadata["simulated.number_of_spills"]

    fclFile = args.config

    # Get a copy of the fcl file and bring it here
    # First step, find the fcl
    fclPath = resolveFclPath(fclFile)
    print "Found fcl file here: ", fclPath
    # Next step, make a temporary name for the fcl in the local directory and copy it here
    tmpFclName = os.path.basename(fclPath).replace(".fcl", "_" + os.path.splitext(inFileBase)[0] + ".fcl") 
    print "Creating local copy : ", tmpFclName
    shutil.copy(fclPath, tmpFclName)

    # Open the fcl file so that we can append output filenames to it 
    fclFileObj = open(tmpFclName, 'a')

    # Start setting up the nova command, add SAM parameters 
    cmdList = ['nova']
    cmdList.append('-c')
    cmdList.append(tmpFclName)
    cmdList.append('--sam-application-family=nova')
    cmdList.append('--sam-application-version=' + release)
    cmdList.append('--sam-file-type=' + file_type)

    # if not (args.outTier or args.cafTier):
    #     raise ValueError("No output tiers or CAF module labels specified.")
    if not args.outTier:
      args.outTier = []
    if not args.cafTier:
      args.cafTier = []

    outList = []
    # Loop over output tiers 
    for outTier in args.outTier:
        try:
            output = outTier.split(":")[0]
            tier   = outTier.split(":")[1]
        except:
            raise ValueError("Output data tier: " + outTier + "not formatted corectly, should be <output_name>:<data_tier>")
        cmdList.append('--sam-data-tier=' + outTier)
        cmdList.append('--sam-stream-name=' +output + ':' + str(stream))
        outName = getOutputFileName(detector, runNum, subNum, streamEntry, release, subversion, dataFlag, tier, nevt)
        outList.append(outName);
        print "Output file name: ", outName, " for tier ", tier, " and output ", output  
        fclFileObj.write("outputs." + output + '.fileName: "'+ outName + '"\n')


    for cafTier in args.cafTier:
      try:
        cafLabel = cafTier.split(":")[0]
        tier   = cafTier.split(":")[1]
      except:
        raise ValueError("Output data tier: " + outTier + "not formatted corectly, should be <output_name>:<data_tier>")
      cafName = getOutputFileName(detector, runNum, subNum, streamEntry, release, subversion, dataFlag, tier, nevt)
      print cafLabel, tier, cafName
      fclFileObj.write("physics.producers." + cafLabel + '.CAFFilename: "' + cafName + '" \n')
      fclFileObj.write("physics.producers." + cafLabel + '.DataTier: "' + tier + '" \n')


    fclFileObj.close()

    print "Config: "
    for line in open(tmpFclName):
        print line.strip()

    if args.n != None:
        cmdList.append('-n')
        cmdList.append(str(args.n))
    cmdList.append(inFile)

    cmd =  ' '.join(cmdList)

    print 'Running:', cmd

    retCode = subprocess.call(cmdList)

    if retCode == 0:
        if args.declareFiles == True:
            checkAndMoveFiles(inFile, True)
        else:
            checkAndMoveFiles(inFile, False)

        if args.copyOut == True:
            copyOutFiles(args.hashDirs)
    else:
        #job didn't succeed, remove output files if they exist
        for file in outList:
            try:
                os.remove("./" + file)
            except OSError:
                pass

    #clean up section
    dirList = os.listdir(".")
    for file in dirList:
        skip_match = _skip_pattern.match(file)
        if skip_match != None:
            print file, "contains *hist*.root or RootOutput*.root: clean up"
            os.remove("./" + file)
    if not args.noCleanup:
      os.remove(tmpFclName)
      os.remove(inFile)
    dh = ifdh.ifdh("http://samweb.fnal.gov:8480/sam/nova/api")
    dh.cleanup()
            
    exit(retCode)