BTaggingCommission/checkProductionIntegrity.py at master · alphatsai/BTaggingCommission · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import os
import sys
import optparse
#from storeTools import getEOSlslist
from subprocess import Popen, PIPE

"""
steer the script
"""

def getEOSlslist(directory, mask='', prepend='root://eoscms//eos/cms/'):
    from subprocess import Popen, PIPE
    print 'looking into: '+directory+'...'

    eos_cmd = '/afs/cern.ch/project/eos/installation/0.3.84-aquamarine/bin/eos.select'
    data = Popen([eos_cmd, 'ls', '/eos/cms/'+directory],stdout=PIPE)
    out,err = data.communicate()

    full_list = []

    ## if input file was single root file:
    if directory.endswith('.root'):
        if len(out.split('\n')[0]) > 0:
            return [prepend + directory]

    ## instead of only the file name append the string to open the file in ROOT
    for line in out.split('\n'):
        if len(line.split()) == 0: continue
        full_list.append(prepend + directory + '/' + line)

    ## strip the list of files if required
    if mask != '':
        stripped_list = [x for x in full_list if mask in x]
        return stripped_list

    ## return
    return full_list

def main():

    #eos_cmd = '/afs/cern.ch/project/eos/installation/0.2.41/bin/eos.select'
    eos_cmd = '/afs/cern.ch/project/eos/installation/0.3.84-aquamarine/bin/eos.select'

    #configuration
    usage = 'usage: %prog [options]'
    parser = optparse.OptionParser(usage)
    parser.add_option('-i', '--inDir',       dest='inDir',       help='input directory with files',               default=None,   type='string')
    parser.add_option('-o', '--outDir',      dest='outDir',      help='output directory with files',              default=None,   type='string')
    parser.add_option('-c', '--cleanup',     dest='cleanup',     help='removes original crab directory',          default =False, action='store_true')
    parser.add_option(      '--nocheck',     dest='nocheck',     help='do not prompt user',                       default=False,  action='store_true')
    (opt, args) = parser.parse_args()

    Popen([eos_cmd, ' -b fuse mount', 'eos'],stdout=PIPE).communicate()

    if opt.outDir is None: opt.outDir=opt.inDir

    dset_list=getEOSlslist(directory=opt.inDir,prepend='')
    for dset in dset_list:
        dsetname=dset.split('/')[-1]

        pub_list=getEOSlslist(directory=dset,prepend='')
        for pubDir in pub_list:

            if not 'crab' in pubDir:
                print 'Ambiguity found @ <publication-name> for <primary-dataset>=%s , bailing out'%dsetname
                continue
            pub=pubDir.split('/crab_')[-1]

            time_list=getEOSlslist(directory=pubDir,prepend='')
            if len(time_list)!=1:
                print 'Ambiguity found @ <time-stamp> for <primary-dataset>=%s , bailing out'%dsetname
                continue
            time_stamp=time_list[0].split('/')[-1]

            out_list=[]
            count_list=getEOSlslist(directory=time_list[0],prepend='')
            for count in count_list: out_list += getEOSlslist(directory=count,prepend='')
            file_list=[x for x in out_list if '.root' in x]

            newDir='%s/%s' % (opt.outDir,pub)
            print '<primary-dataset>=%s <publication-name>=crab_%s <time-stamp>=%s has %d files' % (dsetname,pub,time_stamp,len(file_list) )
            if not opt.nocheck:
                choice = raw_input('Will move to %s current output directory. [y/n] ?' % newDir ).lower()
                if not 'y' in choice : continue

            Popen([eos_cmd, 'mkdir', '-p /eos/cms/'+newDir],stdout=PIPE).communicate()

            moveIndividualFiles=True
            if len(file_list)>0:
                subgroupMerge = int( raw_input('This set has %d files. Merge into groups? (enter 0 if no merging)' % len(file_list)) )
                if subgroupMerge>0:
                    moveIndividualFiles=False

                    splitFunc = lambda A, n=subgroupMerge: [A[i:i+n] for i in range(0, len(A), n)]
                    split_file_lists = splitFunc( file_list )

                    for ilist in xrange(0,len(split_file_lists)):
                        mergedFileName='/tmp/MergedJetTree_%d.root '%ilist
                        toAdd='%s ' % mergedFileName
                        for f in split_file_lists[ilist]:
                            toAdd += 'eos/cms/%s '%f

                        os.system('hadd -f %s'%toAdd)
                        os.system('xrdcp  -f %s root://eoscms//eos/cms/%s/' %(mergedFileName,newDir))

                #if still needed copy individual files
                if moveIndividualFiles:
                    for f in file_list : os.system('xrdcp  -f %s eos/cms/%s/' % (f, newDir) )

            if not opt.nocheck and opt.cleanup :
                choice = raw_input('Will remove output directory. [y/n] ?').lower()
                if 'y' in choice:
                    Popen([eos_cmd, 'rm', '-r /eos/cms/'+dset],stdout=PIPE).communicate()

            print 'Crab outputs may now be found in %s' % newDir

    Popen([eos_cmd, ' -b fuse umount', 'eos'],stdout=PIPE).communicate()
    print '-'*50
    print 'All done. In case errors were found check that the crab output structure is '
    print '<outLFNDirBase>/<primary-dataset>/<publication-name>/<time-stamp>/<counter>/<file-name>'
    print '-'*50


"""
for execution from another script
"""
if __name__ == "__main__":
    sys.exit(main())