forked from cms-tau-pog/TauFW
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathpico.py
executable file
·1618 lines (1490 loc) · 71.4 KB
/
pico.py
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#! /usr/bin/env python
# Author: Izaak Neutelings (April 2020)
import os, sys, re, glob, json
from datetime import datetime
from collections import OrderedDict
import ROOT; ROOT.PyConfig.IgnoreCommandLineOptions = True
import TauFW.PicoProducer.tools.config as GLOB
from TauFW.common.tools.file import ensuredir, ensurefile, ensureinit, getline
from TauFW.common.tools.utils import execute, chunkify, repkey, alphanum_key, lreplace
from TauFW.common.tools.log import Logger, color, bold
from TauFW.PicoProducer.analysis.utils import getmodule, ensuremodule
from TauFW.PicoProducer.batch.utils import getbatch, getcfgsamples, chunkify_by_evts, evtsplitexp
from TauFW.PicoProducer.storage.utils import getstorage, getsamples, isvalid, print_no_samples
from argparse import ArgumentParser
os.chdir(GLOB.basedir)
CONFIG = GLOB.getconfig(verb=0)
LOG = Logger()
###############
# INSTALL #
###############
def main_install(args):
"""Install producer."""
# TODO:
# - guess location (lxplus/PSI/...)
# - set defaults of config file
# - outside CMSSW: create symlinks for standalone
if args.verbosity>=1:
print ">>> main_install", args
verbosity = args.verbosity
###########
# LIST #
###########
def main_list(args):
"""List contents of configuration for those lazy to do 'cat config/config.json'."""
if args.verbosity>=1:
print ">>> main_list", args
verbosity = args.verbosity
cfgname = CONFIG._path
if verbosity>=1:
print '-'*80
print ">>> %-14s = %s"%('cfgname',cfgname)
print ">>> %-14s = %s"%('config',CONFIG)
print '-'*80
print ">>> Configuration %s:"%(cfgname)
for variable, value in CONFIG.iteritems():
variable = "'"+color(variable)+"'"
if isinstance(value,dict):
print ">>> %s:"%(variable)
for key, item in value.iteritems():
if isinstance(item,basestring): item = str(item)
print ">>> %-12r -> %r"%(str(key),str(item))
else:
if isinstance(value,basestring): value = str(value)
print ">>> %-24s = %r"%(variable,value)
###########
# GET #
###########
def main_get(args):
"""Get information of given variable in configuration or samples."""
if args.verbosity>=1:
print ">>> main_get", args
variable = args.variable
eras = args.eras
channels = args.channels or [""]
dtypes = args.dtypes
filters = args.samples
vetoes = args.vetoes
inclurl = args.inclurl # include URL in filelist
checkdas = args.checkdas or args.dasfiles # check file list in DAS
checklocal = args.checklocal # check nevents in local files
limit = args.limit
writedir = args.write # write sample file list to text file
tag = args.tag
verbosity = args.verbosity
getnevts = variable in ['nevents','nevts']
cfgname = CONFIG._path
if verbosity>=1:
print '-'*80
print ">>> %-14s = %s"%('variable',variable)
print ">>> %-14s = %s"%('eras',eras)
print ">>> %-14s = %s"%('channels',channels)
print ">>> %-14s = %s"%('cfgname',cfgname)
print ">>> %-14s = %s"%('config',CONFIG)
print '-'*80
# LIST SAMPLES
if variable=='samples':
if not eras:
LOG.warning("Please specify an era to get a sample for.")
for era in eras:
for channel in channels:
if channel:
print ">>> Getting file list for era %r, channel %r"%(era,channel)
else:
print ">>> Getting file list for era %r"%(era)
samples = getsamples(era,channel=channel,dtype=dtypes,filter=filters,veto=vetoes,verb=verbosity)
if not samples:
LOG.warning("No samples found for era %r."%(era))
for sample in samples:
print ">>> %s"%(bold(sample.name))
for path in sample.paths:
print ">>> %s"%(path)
# LIST SAMPLE FILES
elif variable in ['files','nevents','nevts']:
# LOOP over ERAS & CHANNELS
if not eras:
LOG.warning("Please specify an era to get a sample for.")
for era in eras:
for channel in channels:
target = "file list" if variable=='files' else "nevents"
if channel:
print ">>> Getting %s for era %r, channel %r"%(target,era,channel)
else:
print ">>> Getting %s for era %r"%(target,era)
print ">>> "
# VERBOSE
if verbosity>=1:
print ">>> %-12s = %r"%('channel',channel)
# GET SAMPLES
LOG.insist(era in CONFIG.eras,"Era '%s' not found in the configuration file. Available: %s"%(era,CONFIG.eras))
samples = getsamples(era,channel=channel,dtype=dtypes,filter=filters,veto=vetoes,verb=verbosity)
# LOOP over SAMPLES
for sample in samples:
print ">>> %s"%(bold(sample.name))
for path in sample.paths:
print ">>> %s"%(bold(path))
if getnevts or checkdas or checklocal:
nevents = sample.getnevents(das=(not checklocal),verb=verbosity+1)
storage = sample.storage.__class__.__name__ if checklocal else "DAS"
print ">>> %-7s = %s (%s)"%('nevents',nevents,storage)
if variable=='files':
infiles = sample.getfiles(das=checkdas,url=inclurl,limit=limit,verb=verbosity+1)
print ">>> %-7s = %r"%('url',sample.url)
print ">>> %-7s = %r"%('postfix',sample.postfix)
print ">>> %-7s = %s"%('nfiles',len(infiles))
print ">>> %-7s = [ "%('infiles')
for file in infiles:
print ">>> %r"%file
print ">>> ]"
print ">>> "
if writedir: # write files to text files
flistname = repkey(writedir,ERA=era,GROUP=sample.group,SAMPLE=sample.name,TAG=tag)
print ">>> Write list to %r..."%(flistname)
sample.writefiles(flistname,nevts=getnevts)
# CONFIGURATION
else:
if variable in CONFIG:
print ">>> Configuration of %r: %s"%(variable,color(CONFIG[variable]))
else:
print ">>> Did not find %r in the configuration"%(variable)
#############
# WRITE #
#############
def main_write(args):
"""Get information of given variable in configuration or samples."""
if args.verbosity>=1:
print ">>> main_write", args
listname = args.listname # write sample file list to text file
eras = args.eras
channels = args.channels or [""]
dtypes = args.dtypes
filters = args.samples
vetoes = args.vetoes
checkdas = args.checkdas or args.dasfiles # check file list in DAS
getnevts = args.getnevts # check nevents in local files
verbosity = args.verbosity
cfgname = CONFIG._path
if verbosity>=1:
print '-'*80
print ">>> %-14s = %s"%('listname',listname)
print ">>> %-14s = %s"%('getnevts',getnevts)
print ">>> %-14s = %s"%('eras',eras)
print ">>> %-14s = %s"%('channels',channels)
print ">>> %-14s = %s"%('cfgname',cfgname)
print ">>> %-14s = %s"%('config',CONFIG)
print '-'*80
# LOOP over ERAS & CHANNELS
if not eras:
LOG.warning("Please specify an era to get a sample for.")
for era in eras:
for channel in channels:
info = ">>> Getting file list for era %r"%(era)
if channel:
info += ", channel %r"%(channel)
print info
print ">>> "
# VERBOSE
if verbosity>=1:
print ">>> %-12s = %r"%('channel',channel)
LOG.insist(era in CONFIG.eras,"Era '%s' not found in the configuration file. Available: %s"%(era,CONFIG.eras))
samples = getsamples(era,channel=channel,dtype=dtypes,filter=filters,veto=vetoes,verb=verbosity)
# LOOP over SAMPLES
for sample in samples:
print ">>> %s"%(bold(sample.name))
#infiles = sample.getfiles(das=checkdas,url=inclurl,limit=limit,verb=verbosity+1)
flistname = repkey(listname,ERA=era,GROUP=sample.group,SAMPLE=sample.name) #,TAG=tag
sample.writefiles(flistname,nevts=getnevts,das=checkdas)
print ">>> "
###########
# SET #
###########
def main_set(args):
"""Set variables in the config file."""
if args.verbosity>=1:
print ">>> main_set", args
variable = args.variable
key = args.key # 'channel' or 'era'
value = args.value
verbosity = args.verbosity
cfgname = CONFIG._path
if key: # redirect 'channel' and 'era' keys to main_link
args.subcommand = variable
return main_link(args)
elif variable in ['channel','era']:
LOG.throw(IOError,"Variable '%s' is reserved for dictionaries!"%(variable))
if verbosity>=1:
print '-'*80
print ">>> Setting variable '%s' to '%s' config"%(color(variable),value)
if verbosity>=1:
print ">>> %-14s = %s"%('cfgname',cfgname)
print ">>> %-14s = %s"%('config',CONFIG)
print '-'*80
if variable=='all':
if 'default' in value:
GLOB.setdefaultconfig(verb=verb)
else:
LOG.warning("Did not recognize value '%s'. Did you mean 'default'?"%(value))
elif variable in ['nfilesperjob','maxevtsperjob']:
CONFIG[variable] = int(value)
CONFIG.write()
else:
CONFIG[variable] = value
CONFIG.write()
############
# LINK #
############
def main_link(args):
"""Link channels or eras in the config file."""
if args.verbosity>=1:
print ">>> main_link", args
variable = args.subcommand
varkey = variable+'s'
key = args.key
value = args.value
verbosity = args.verbosity
cfgname = CONFIG._path
if verbosity>=1:
print '-'*80
print ">>> Linking %s '%s' to '%s' in the configuration..."%(variable,color(key),value)
if verbosity>=1:
print ">>> %-14s = %s"%('cfgname',cfgname)
print ">>> %-14s = %s"%('key',key)
print ">>> %-14s = %s"%('value',value)
print ">>> %-14s = %s"%('config',CONFIG)
print '-'*80
# SANITY CHECKS
if varkey not in CONFIG:
CONFIG[varkey] = { }
LOG.insist(isinstance(CONFIG[varkey],dict),"%s in %s has to be a dictionary"%(varkey,cfgname))
oldval = value
for char in '/\,:;!?\'" ':
if char in key:
LOG.throw(IOError,"Given key '%s', but keys cannot contain any of these characters: %s"%(key,char))
if varkey=='channels':
if 'skim' in key.lower(): #or 'test' in key:
parts = value.split(' ') # "PROCESSOR [--FLAG[=VALUE] ...]"
script = os.path.basename(parts[0]) # separate script from options
ensurefile("python/processors",script)
value = ' '.join([script]+parts[1:])
else:
parts = value.split(' ') # "MODULE [KEY=VALUE ...]"
module = parts[0]
LOG.insist(all('=' in o for o in parts[1:]),"All extra module options should be of format KEY=VALUE!")
if 'python/analysis/' in module: # useful for tab completion
module = module.split('python/analysis/')[-1].replace('/','.')
module = module.rstrip('.py')
path = os.path.join('python/analysis/','/'.join(module.split('.')[:-1]))
ensureinit(path,by="pico.py")
ensuremodule(module)
value = ' '.join([module]+parts[1:])
elif varkey=='eras':
if 'samples/' in value: # useful for tab completion
value = ''.join(value.split('samples/')[1:])
path = os.path.join("samples",repkey(value,ERA='*',CHANNEL='*',TAG='*'))
LOG.insist(glob.glob(path),"Did not find any sample lists '%s'"%(path))
ensureinit(os.path.dirname(path),by="pico.py")
if value!=oldval:
print ">>> Converted '%s' to '%s'"%(oldval,value)
CONFIG[varkey][key] = value
CONFIG.write()
##############
# REMOVE #
##############
def main_rm(args):
"""Remove variable from the config file."""
if args.verbosity>=1:
print ">>> main_rm", args
variable = args.variable
key = args.key # 'channel' or 'era'
verbosity = args.verbosity
cfgname = CONFIG._path
if verbosity>=1:
print '-'*80
if key:
print ">>> Removing %s '%s' from the configuration..."%(variable,color(key))
else:
print ">>> Removing variable '%s' from the configuration..."%(color(variable))
if verbosity>=1:
print ">>> %-14s = %s"%('variable',variable)
print ">>> %-14s = %s"%('key',key)
print ">>> %-14s = %s"%('cfgname',cfgname)
print ">>> %-14s = %s"%('config',CONFIG)
print '-'*80
if key: # redirect 'channel' and 'era' keys to main_link
variable = variable+'s'
if variable in CONFIG:
if key in CONFIG[variable]:
CONFIG[variable].pop(key,None)
CONFIG.write()
else:
print ">>> %s '%s' not in the configuration. Nothing to remove..."%(variable.capitalize(),key)
else:
print ">>> Variable '%s' not in the configuration. Nothing to remove..."%(variable)
else:
if variable in CONFIG:
CONFIG.pop(variable)
CONFIG.write()
else:
print ">>> Variable '%s' not in the configuration. Nothing to remove..."%(variable)
###########
# RUN #
###########
def main_run(args):
"""Run given module locally."""
if args.verbosity>=1:
print ">>> main_run", args
eras = args.eras
channels = args.channels
tag = args.tag
outdir = args.outdir
dtypes = args.dtypes
filters = args.samples
vetoes = args.vetoes
force = args.force
extraopts = args.extraopts # extra options for module (for all runs)
maxevts = args.maxevts
dasfiles = args.dasfiles
userfiles = args.infiles
nfiles = args.nfiles
nsamples = args.nsamples
prefetch = args.prefetch
dryrun = args.dryrun
verbosity = args.verbosity
preselect = args.preselect
# LOOP over ERAS
if not eras:
print ">>> Please specify a valid era (-y)."
if not channels:
print ">>> Please specify a valid channel (-c)."
for era in eras:
moddict = { } # save time by loading samples and get their files only once
# LOOP over CHANNELS
for channel in channels:
LOG.header("%s, %s"%(era,channel))
# MODULE & PROCESSOR
skim = 'skim' in channel.lower()
module, processor, procopts, extrachopts = getmodule(channel,extraopts)
# VERBOSE
if verbosity>=1:
print '-'*80
print ">>> Running %r"%(channel)
print ">>> %-12s = %r"%('channel',channel)
print ">>> %-12s = %r"%('module',module)
print ">>> %-12s = %r"%('processor',processor)
print ">>> %-12s = %r"%('procopts',procopts)
print ">>> %-12s = %r"%('extrachopts',extrachopts)
print ">>> %-12s = %r"%('prefetch',prefetch)
print ">>> %-12s = %r"%('preselect',preselect)
print ">>> %-12s = %s"%('filters',filters)
print ">>> %-12s = %s"%('vetoes',vetoes)
print ">>> %-12s = %r"%('dtypes',dtypes)
print ">>> %-12s = %r"%('userfiles',userfiles)
print ">>> %-12s = %r"%('outdir',outdir)
# GET SAMPLES
if not userfiles and (filters or vetoes or dtypes):
LOG.insist(era in CONFIG.eras,"Era '%s' not found in the configuration file. Available: %s"%(era,CONFIG.eras))
samples = getsamples(era,channel=channel,tag=tag,dtype=dtypes,filter=filters,veto=vetoes,moddict=moddict,verb=verbosity)
if nsamples>0:
samples = samples[:nsamples]
if not samples:
print_no_samples(dtypes,filters,vetoes)
else:
samples = [None]
if verbosity>=2:
print ">>> %-12s = %r"%('samples',samples)
if verbosity>=1:
print '-'*80
# LOOP over SAMPLES
for sample in samples:
if sample:
print ">>> %s"%(bold(sample.name))
if verbosity>=1:
for path in sample.paths:
print ">>> %s"%(bold(path))
# SETTINGS
dtype = None
extraopts_ = extrachopts[:] # extra options for module (for this channel & sample)
if sample:
filetag = "_%s_%s_%s%s"%(channel,era,sample.name,tag)
if sample.extraopts:
extraopts_.extend(sample.extraopts)
else:
filetag = "_%s_%s%s"%(channel,era,tag)
if verbosity>=1:
print ">>> %-12s = %s"%('sample',sample)
print ">>> %-12s = %r"%('filetag',filetag) # postfix
print ">>> %-12s = %s"%('extraopts',extraopts_)
# GET FILES
infiles = [ ]
if userfiles:
infiles = userfiles[:]
elif sample:
nevents = 0
infiles = sample.getfiles(das=dasfiles,verb=verbosity)
dtype = sample.dtype
if nfiles>0:
infiles = infiles[:nfiles]
if verbosity==1:
print ">>> %-12s = %r"%('dtype',dtype)
print ">>> %-12s = %s"%('nfiles',len(infiles))
print ">>> %-12s = [ "%('infiles')
for file in infiles:
print ">>> %r"%file
print ">>> ]"
if verbosity==1:
print '-'*80
# RUN
runcmd = processor
if procopts:
runcmd += " %s"%(procopts)
if skim:
runcmd += " -y %s -o %s"%(era,outdir)
if preselect:
runcmd += " --preselect '%s'"%(preselect)
###elif 'test' in channel:
### runcmd += " -o %s"%(outdir)
else: # analysis
runcmd += " -y %s -c %s -M %s -o %s"%(era,channel,module,outdir)
if dtype:
runcmd += " -d %r"%(dtype)
if filetag:
runcmd += " -t %r"%(filetag) # postfix
if maxevts:
runcmd += " -m %s"%(maxevts)
if infiles:
runcmd += " -i %s"%(' '.join(infiles))
if prefetch:
runcmd += " -p"
if extraopts_:
runcmd += " --opt '%s'"%("' '".join(extraopts_))
#elif nfiles:
# runcmd += " --nfiles %s"%(nfiles)
print ">>> Executing: "+bold(runcmd)
if not dryrun:
#execute(runcmd,dry=dryrun,verb=verbosity+1) # real-time print out does not work well with python script
os.system(runcmd)
print
##################
# GET MODULE #
##################
def getmodule(channel,extraopts):
"""Help function to get the module and processor."""
LOG.insist(channel in CONFIG.channels,"Channel '%s' not found in the configuration file. Available: %s"%(channel,CONFIG.channels))
module = CONFIG.channels[channel]
procopts = "" # extra options for processor
extrachopts = extraopts[:] # extra options for module (per channel)
if 'skim' in channel.lower():
parts = module.split(' ') # "PROCESSOR [--FLAG[=VALUE] ...]"
processor = parts[0]
procopts = ' '.join(parts[1:])
###elif channel=='test':
### processor = module
else:
parts = module.split(' ') # "MODULE [KEY=VALUE ...]"
processor = "picojob.py"
module = parts[0]
ensuremodule(module) # sanity check
extrachopts.extend(parts[1:])
procpath = os.path.join("python/processors",processor)
if not os.path.isfile(procpath):
LOG.throw(IOError,"Processor '%s' does not exist in '%s'..."%(processor,procpath))
processor = os.path.abspath(procpath)
return module, processor, procopts, extrachopts
####################
# PREPARE JOBS #
####################
def preparejobs(args):
"""Help function for (re)submission to iterate over samples per given channel and era
and prepare job config and list."""
if args.verbosity>=1:
print ">>> preparejobs", args
resubmit = args.subcommand=='resubmit'
eras = args.eras
channels = args.channels
tag = args.tag
dtypes = args.dtypes # filter (only include) these sample types ('data','mc','embed')
filters = args.samples # filter (only include) these samples (glob patterns)
vetoes = args.vetoes # exclude these sample (glob patterns)
dasfiles = args.dasfiles # explicitly process nanoAOD files stored on DAS (as opposed to local storage)
checkdas = args.checkdas # look up number of events in DAS and compare to processed events in job output
checkqueue = args.checkqueue # check job status to speed up if batch is slow: 0 (no check), 1 (check once), -1 (check every job)
extraopts = args.extraopts # extra options for module (for all runs)
prefetch = args.prefetch # copy input file first to local output directory
preselect = args.preselect # preselection string for post-processing
nfilesperjob = args.nfilesperjob # split jobs based on number of files
maxevts = args.maxevts # split jobs based on events
split_nfpj = args.split_nfpj # split failed (file-based) chunks into even smaller chunks
testrun = args.testrun # only run a few test jobs
queue = args.queue # queue option for the batch system (job flavor for HTCondor)
force = args.force # force submission, even if old job output exists
prompt = args.prompt # ask user for confirmation
tmpdir = args.tmpdir or CONFIG.get('tmpskimdir',None) # temporary dir for creating skimmed file before copying to outdir
verbosity = args.verbosity
jobs = [ ]
# LOOP over ERAS
for era in eras:
moddict = { } # save time by loading samples and get their file list only once
# LOOP over CHANNELS
for channel in channels:
LOG.header("%s, %s"%(era,channel))
# MODULE & PROCESSOR
skim = 'skim' in channel.lower()
module, processor, procopts, extrachopts = getmodule(channel,extraopts)
if verbosity>=1:
print '-'*80
print ">>> %-12s = %r"%('channel',channel)
print ">>> %-12s = %r"%('processor',processor)
print ">>> %-12s = %r"%('module',module)
print ">>> %-12s = %r"%('procopts',procopts)
print ">>> %-12s = %r"%('extrachopts',extrachopts)
print ">>> %-12s = %s"%('filters',filters)
print ">>> %-12s = %s"%('vetoes',vetoes)
print ">>> %-12s = %r"%('dtypes',dtypes)
# GET SAMPLES
jobdirformat = CONFIG.jobdir # for job config & log files
outdirformat = CONFIG.nanodir if skim else CONFIG.outdir # for job output
jobdir_ = ""
jobcfgs = ""
if resubmit:
# TODO: allow user to resubmit given config file
jobdir_ = repkey(jobdirformat,ERA=era,SAMPLE='*',CHANNEL=channel,TAG=tag)
jobcfgs = repkey(os.path.join(jobdir_,"config/jobconfig_$SAMPLE$TAG_try[0-9]*.json"),
ERA=era,SAMPLE='*',CHANNEL=channel,TAG=tag)
if verbosity>=2:
print ">>> %-12s = %s"%('cwd',os.getcwd())
print ">>> %-12s = %s"%('jobcfgs',jobcfgs)
samples = getcfgsamples(jobcfgs,filter=filters,veto=vetoes,dtype=dtypes,verb=verbosity)
else:
LOG.insist(era in CONFIG.eras,"Era '%s' not found in the configuration file. Available: %s"%(era,CONFIG.eras))
samples = getsamples(era,channel=channel,tag=tag,dtype=dtypes,filter=filters,veto=vetoes,moddict=moddict,verb=verbosity)
if verbosity>=2:
print ">>> Found samples: "+", ".join(repr(s.name) for s in samples)
if testrun:
samples = samples[:2] # run at most two samples
# SAMPLE over SAMPLES
found = len(samples)>=0
failed = [ ] # failed samples
for sample in samples:
print ">>> %s"%(bold(sample.name))
for path in sample.paths:
print ">>> %s"%(bold(path))
# DIRECTORIES
subtry = sample.subtry+1 if resubmit else 1
jobids = sample.jobcfg.get('jobids',[ ])
queue_ = queue or sample.jobcfg.get('queue',None)
prefetch_ = sample.jobcfg.get('prefetch',prefetch) or prefetch # if resubmit: reuse old setting, or override by user
dtype = sample.dtype
postfix = "_%s%s"%(channel,tag)
jobtag = "%s_try%d"%(postfix,subtry)
jobname = "%s%s_%s%s"%(sample.name,postfix,era,"_try%d"%subtry if subtry>1 else "")
extraopts_ = extrachopts[:] # extra options for module (for this channel & sample)
if sample.extraopts:
extraopts_.extend(sample.extraopts)
nfilesperjob_ = nfilesperjob if nfilesperjob>0 else sample.nfilesperjob if sample.nfilesperjob>0 else CONFIG.nfilesperjob # priority: USER > SAMPLE > CONFIG
maxevts_ = maxevts if maxevts>0 else sample.maxevts if sample.maxevts>0 else CONFIG.maxevtsperjob # priority: USER > SAMPLE > CONFIG
if split_nfpj>1: # divide nfilesperjob by split_nfpj
nfilesperjob_ = int(max(1,nfilesperjob_/float(split_nfpj)))
elif resubmit and maxevts<=0: # reuse previous maxevts settings if maxevts not set by user
maxevts_ = sample.jobcfg.get('maxevts',maxevts_)
if nfilesperjob<=0: # reuse previous nfilesperjob settings if nfilesperjob not set by user
nfilesperjob_ = sample.jobcfg.get('nfilesperjob',nfilesperjob_)
daspath = sample.paths[0].strip('/')
outdir = repkey(outdirformat,ERA=era,CHANNEL=channel,TAG=tag,SAMPLE=sample.name,
DAS=daspath,PATH=daspath,GROUP=sample.group)
jobdir = ensuredir(repkey(jobdirformat,ERA=era,CHANNEL=channel,TAG=tag,SAMPLE=sample.name,
DAS=daspath,PATH=daspath,GROUP=sample.group))
cfgdir = ensuredir(jobdir,"config")
logdir = ensuredir(jobdir,"log")
cfgname = "%s/jobconfig%s.json"%(cfgdir,jobtag)
joblist = '%s/jobarglist%s.txt'%(cfgdir,jobtag)
if verbosity==1:
print ">>> %-12s = %s"%('cfgname',cfgname)
print ">>> %-12s = %s"%('joblist',joblist)
elif verbosity>=2:
print '-'*80
print ">>> Preparing job %ssubmission for '%s'"%("re" if resubmit else "",sample.name)
print ">>> %-12s = %r"%('processor',processor)
print ">>> %-12s = %r"%('dtype',dtype)
print ">>> %-12s = %r"%('jobname',jobname)
print ">>> %-12s = %r"%('jobtag',jobtag)
print ">>> %-12s = %r"%('postfix',postfix)
print ">>> %-12s = %r"%('outdir',outdir)
print ">>> %-12s = %r"%('extraopts',extraopts_)
print ">>> %-12s = %r"%('prefetch',prefetch_)
print ">>> %-12s = %r"%('preselect',preselect)
print ">>> %-12s = %r"%('cfgdir',cfgdir)
print ">>> %-12s = %r"%('logdir',logdir)
print ">>> %-12s = %r"%('tmpdir',tmpdir)
print ">>> %-12s = %r"%('cfgname',cfgname)
print ">>> %-12s = %r"%('joblist',joblist)
print ">>> %-12s = %s"%('try',subtry)
print ">>> %-12s = %r"%('jobids',jobids)
print ">>> %-12s = %r"%('queue',queue_)
# CHECKS
if os.path.isfile(cfgname):
# TODO: check for running jobs ?
skip = False
if force:
LOG.warning("Job configuration %r already exists and will be overwritten! "%(cfgname)+
"Please beware of conflicting job output!")
elif args.prompt:
LOG.warning("Job configuration %r already exists and might cause conflicting job output!"%(cfgname))
while True:
submit = raw_input(">>> Submit anyway? [y/n] "%(nchunks))
if 'f' in submit.lower(): # submit this job, and stop asking
print ">>> Force all."
force = True; skip = True; break
elif 'y' in submit.lower(): # submit this job
print ">>> Continue submission..."
skip = True; break
elif 'n' in submit.lower(): # do not submit this job
print ">>> Not submitting."
break
else:
print ">>> '%s' is not a valid answer, please choose y/n."%submit
else:
skip = True
LOG.warning("Job configuration %r already exists and might cause conflicting job output! "%(cfgname)+
"To submit anyway, please use the --force flag")
if skip: # do not submit this job
failed.append(sample)
print ""
continue
if not resubmit: # check for existing jobss
cfgpattern = re.sub(r"(?<=try)\d+(?=.json$)",r"*",cfgname)
cfgnames = [f for f in glob.glob(cfgpattern) if not f.endswith("_try1.json")]
if cfgnames:
LOG.warning("Job configurations for resubmission already exists! This can cause conflicting job output!"+
"If you are sure you want to submit from scratch, please remove these files:\n>>> "+"\n>>> ".join(cfgnames))
storage = getstorage(outdir,verb=verbosity,ensure=True)
# GET FILES
nevents = 0
if resubmit: # resubmission
if checkqueue==0 and not jobs: # check jobs only once to speed up performance
batch = getbatch(CONFIG,verb=verbosity)
jobs = batch.jobs(verb=verbosity-1)
infiles, chunkdict = checkchunks(sample,channel=channel,tag=tag,jobs=jobs,
checkqueue=checkqueue,das=checkdas,verb=verbosity)[:2]
nevents = sample.jobcfg['nevents'] # updated in checkchunks
else: # first-time submission
infiles = sample.getfiles(das=dasfiles,verb=verbosity-1)
if checkdas:
nevents = sample.getnevents()
chunkdict = { }
if testrun:
infiles = infiles[:4] # only run two files per sample
if verbosity==1:
print ">>> %-12s = %s"%('maxevts',maxevts_)
print ">>> %-12s = %s"%('nfilesperjob',nfilesperjob_)
print ">>> %-12s = %s"%('nfiles',len(infiles))
elif verbosity>=2:
print ">>> %-12s = %s"%('maxevts',maxevts_)
print ">>> %-12s = %s"%('nfilesperjob',nfilesperjob_)
print ">>> %-12s = %s"%('nfiles',len(infiles))
print ">>> %-12s = [ "%('infiles')
for file in infiles:
print ">>> %r"%file
print ">>> ]"
print ">>> %-12s = %s"%('nevents',nevents)
# CHUNKS - partition/split list
infiles.sort() # to have consistent order with resubmission
chunks = [ ] # chunk indices
if maxevts_>1:
try:
fchunks = chunkify_by_evts(infiles,maxevts_,verb=verbosity) # list of file chunks split by events
except IOError as err: # capture if opening files fail
print "IOError: "+err.message
LOG.warning("Skipping submission...")
failed.append(sample)
print ""
continue # ignore this submission
if testrun:
fchunks = fchunks[:4]
else:
fchunks = chunkify(infiles,nfilesperjob_) # list of file chunks split by number of files
nfiles = len(infiles)
nchunks = len(fchunks)
if verbosity>=1:
print ">>> %-12s = %s"%('nchunks',nchunks)
if verbosity>=2:
print '-'*80
# WRITE JOB LIST with arguments per job
if args.verbosity>=1:
print ">>> Creating job list %s..."%(joblist)
if fchunks:
with open(joblist,'w') as listfile:
ichunk = 0
for fchunk in fchunks:
while ichunk in chunkdict:
ichunk += 1 # allows for different nfilesperjob on resubmission
continue
evtmatch = evtsplitexp.match(fchunk[0]) # $fname:$firstevt:$maxevts
if evtmatch:
LOG.insist(len(fchunk)==1,"Chunks of event-split files can only have one input file: %s"%(fchunk))
jobfiles = evtmatch.group(1) # input file
firstevt = int(evtmatch.group(2))
maxevts__ = int(evtmatch.group(3)) # limit number of events
else:
jobfiles = ' '.join(fchunk) # list of input files
firstevt = -1
maxevts__ = -1 # do not limit number of events
filetag = postfix
if not skim:
filetag += "_%d"%(ichunk)
elif firstevt>=0:
filetag += "_%d"%(firstevt/maxevts__)
jobcmd = processor
if procopts:
jobcmd += " %s"%(procopts)
if skim:
jobcmd += " -y %s -d '%s' -t %s --copydir %s"%(era,dtype,filetag,outdir)
if tmpdir:
jobcmd += " -o %s"%(tmpdir) # temporary file for job output before copying
###elif channel=='test':
### jobcmd += " -o %s -t %s -i %s"%(outdir,filetag)
else:
jobcmd += " -y %s -d %r -c %s -M %s --copydir %s -t %s"%(era,dtype,channel,module,outdir,filetag)
if prefetch_:
jobcmd += " -p"
if preselect and skim:
jobcmd += " --preselect '%s'"%(preselect)
if firstevt>=0:
jobcmd += " --firstevt %d"%(firstevt) # start at this entry (for event-based splitting)
if testrun: # override maxevts
jobcmd += " -m %d"%(testrun) # process a limited amount of events for test jobs
elif maxevts__>0:
jobcmd += " -m %d"%(maxevts__) # process a limited amount of events for event-based splitting
if extraopts_:
jobcmd += " --opt '%s'"%("' '".join(extraopts_))
jobcmd += " -i %s"%(jobfiles) # add last
if args.verbosity>=1:
print ">>> chunk=%d, jobcmd=%r"%(ichunk,jobcmd)
listfile.write(jobcmd+'\n')
chunkdict[ichunk] = fchunk
chunks.append(ichunk)
# JSON CONFIG
jobcfg = OrderedDict([
('time',str(datetime.now())),
('group',sample.group), ('paths',sample.paths), ('name',sample.name), ('nevents',nevents),
('dtype',dtype), ('channel',channel), ('module',module), ('extraopts',extraopts_),
('jobname',jobname), ('jobtag',jobtag), ('tag',tag), ('postfix',postfix),
('try',subtry), ('queue',queue_), ('jobids',jobids), ('prefetch',prefetch_),
('outdir',outdir), ('jobdir',jobdir), ('cfgdir',cfgdir), ('logdir',logdir),
('cfgname',cfgname), ('joblist',joblist), ('maxevts',maxevts_),
('nfiles',nfiles), ('files',infiles), ('nfilesperjob',nfilesperjob_), #('nchunks',nchunks),
('nchunks',nchunks), ('chunks',chunks), ('chunkdict',chunkdict),
])
# YIELD
yield jobcfg
print
if not found:
print_no_samples(dtypes,filters,vetoes,jobdir_,jobcfgs)
elif failed and len(failed)!=len(samples):
print ">>> %d/%d samples failed: %s\n"%(len(failed),len(samples),', '.join(s.name for s in failed))
##################
# CHECK JOBS #
##################
def checkchunks(sample,**kwargs):
"""Help function to check jobs status: success, pending, failed or missing.
Return list of files to be resubmitted, and a dictionary between chunk index and input files."""
outdir = kwargs.get('outdir', None )
channel = kwargs.get('channel', None )
tag = kwargs.get('tag', None )
checkqueue = kwargs.get('checkqueue', False ) # check queue of batch system for pending jobs
pendjobs = kwargs.get('jobs', [ ] )
checkdas = kwargs.get('das', True ) # check number of events from DAS
showlogs = kwargs.get('showlogs', False ) # print log files of failed jobs
verbosity = kwargs.get('verb', 0 )
oldjobcfg = sample.jobcfg # job config from last job
oldcfgname = oldjobcfg['config']
chunkdict = oldjobcfg['chunkdict'] # filenames
jobids = oldjobcfg['jobids']
joblist = oldjobcfg['joblist']
postfix = oldjobcfg['postfix']
logdir = oldjobcfg['logdir']
nfilesperjob = oldjobcfg['nfilesperjob']
if outdir==None:
outdir = oldjobcfg['outdir']
storage = getstorage(outdir,ensure=True) # StorageElement instance of output directory
if channel==None:
channel = oldjobcfg['channel']
if tag==None:
tag = oldjobcfg['tag']
evtsplit = any(any(evtsplitexp.match(f) for f in chunkdict[i]) for i in chunkdict)
noldchunks = len(chunkdict) # = number of jobs
goodchunks = [ ] # good job output
pendchunks = [ ] # pending or running jobs
badchunks = [ ] # corrupted job output
misschunks = [ ] # missing job output
resubfiles = [ ] # files to resubmit (if bad or missing)
# NUMBER OF EVENTS
nprocevents = 0 # total number of processed events
ndasevents = oldjobcfg['nevents'] # total number of available events
if checkdas and oldjobcfg['nevents']==0:
ndasevents = sample.getnevents()
oldjobcfg['nevents'] = ndasevents
if verbosity>=2:
print ">>> %-12s = %s"%('ndasevents',ndasevents)
if verbosity>=3:
print ">>> %-12s = %s"%('chunkdict',chunkdict)
# CHECK PENDING JOBS
if checkqueue<0 or pendjobs:
batch = getbatch(CONFIG,verb=verbosity)
if checkqueue!=1 or not pendjobs:
pendjobs = batch.jobs(jobids,verb=verbosity-1) # get refreshed job list
else:
pendjobs = [j for j in pendjobs if j.jobid in jobids] # get new job list with right job id
###########################################################################
# CHECK SKIMMED OUTPUT: nanoAOD format, one or more output files per job
if 'skim' in channel.lower(): # and nfilesperjob>1:
flagexp = re.compile(r"-i (.+\.root)") #r"-i ((?:(?<! -).)+\.root[, ])"
flagexp2 = re.compile(r"--firstevt (\d+) -m (\d+)")
chunkexp = re.compile(r"(.+)%s(?:_(\d+))?\.root"%(postfix))
fpatterns = ["*%s.root"%(postfix)]
if evtsplit:
fpatterns.append("*%s_[0-9]*.root"%(postfix))
if verbosity>=2:
print ">>> %-12s = %r"%('flagexp',flagexp.pattern)
print ">>> %-12s = %r"%('flagexp2',flagexp2.pattern)
print ">>> %-12s = %r"%('fpatterns',fpatterns)
print ">>> %-12s = %r"%('chunkexp',chunkexp.pattern)
print ">>> %-12s = %s"%('checkqueue',checkqueue)
print ">>> %-12s = %s"%('pendjobs',pendjobs)
print ">>> %-12s = %s"%('jobids',jobids)
# CHECK PENDING JOBS
pendfiles = [ ]
for job in pendjobs:
if verbosity>=3:
print ">>> Found job %r, status=%r, args=%r"%(job,job.getstatus(),job.args.rstrip())
if job.getstatus() in ['q','r']:
if 'HTCondor' in CONFIG.batch:
jobarg = str(job.args)
matches = flagexp.findall(jobarg)
matches2 = flagexp2.findall(jobarg)
else:
jobarg = getline(joblist,job.taskid-1)
matches = flagexp.findall(jobarg)
matches2 = flagexp2.findall(jobarg)
if verbosity>=3:
print ">>> jobarg =",jobarg.replace('\n','')
print ">>> matches =",matches
print ">>> matches2 =",matches2
if not matches:
continue
infiles = [ ]
for file in matches[0].split():
if not file.endswith('.root'):
break
if matches2:
file += ":%s"%(matches2[0][0]) #,matches2[0][1])
infiles.append(file)
LOG.insist(infiles,"Did not find any ROOT files in job arguments %r, matches=%r"%(jobarg,matches))
ichunk = -1
for i in chunkdict:
if all(any(f in c for c in chunkdict[i]) for f in infiles):
ichunk = i
break
LOG.insist(ichunk>=0,
"Did not find to which the input files of jobids %s belong! "%(jobids)+
"\nichunk=%s,\ninfiles=%s,\nchunkdict=%s"%(ichunk,infiles,chunkdict))
LOG.insist(len(chunkdict[i])==len(infiles),
"Mismatch between input files of jobids %s and chunkdict! "%(jobids)+
"\nichunk=%s,\ninfiles=%s,\nchunkdict[%s]=%s"%(ichunk,infiles,ichunk,chunkdict[ichunk]))
pendchunks.append(ichunk)
# CHECK OUTPUT FILES
badfiles = [ ]
goodfiles = [ ]
outfiles = storage.getfiles(filter=fpatterns,verb=verbosity-1) # get output files
if verbosity>=2:
print ">>> %-12s = %s"%('pendchunks',pendchunks)
print ">>> %-12s = %s"%('outfiles',outfiles)
for fname in outfiles:
if verbosity>=2:
print ">>> Checking job output '%s'..."%(fname)
basename = os.path.basename(fname)
infile = chunkexp.sub(r"\1.root",basename) # reconstruct input file without path or postfix
outmatch = chunkexp.match(basename)
ipart = int(outmatch.group(2) or -1) if outmatch else -1 # >0 if input file split by events
nevents = isvalid(fname) # check for corruption
ichunk = -1
for i in chunkdict:
if ichunk>-1: # found corresponding input file
break
for chunkfile in chunkdict[i]: # find chunk output file belongs to
if infile not in chunkfile: continue
inmatch = evtsplitexp.match(chunkfile)
if inmatch and int(inmatch.group(2))/int(inmatch.group(3))!=ipart: continue
ichunk = i