neutralMesonTSSA/macro/CreateJobDSTLists.py

0001 import glob
0002 import re
0003 import os
0004
0005 if __name__ == "__main__":
0006     # break up the full list of DSTs into smaller lists for condor jobs
0007     # write multiple per-job file lists with 10 DST files per job
0008     dstlistdir = 'filelists/RD/2024/Nov7/'
0009     joblistdir = 'condor/job_filelists/NM/'
0010     maxjobs = 1000000 # in case we want a small test sample
0011     # maxjobs = 5000 # in case we want a small test sample
0012     filesperjob = 3 # (max) number of DSTs to process in one job
0013     jobnum = 0
0014     filesthisjob = 0
0015     filesthisrun = 0
0016     totaldsts = 0
0017
0018     dstlists = glob.glob(dstlistdir+'dst_calo*')
0019     print(f'Found {len(dstlists)} run-by-run DST lists')
0020
0021     lines = []
0022     for dstlist in dstlists:
0023         # print(f'Opening new dst list, jobnum = {jobnum}')
0024         if jobnum >= maxjobs:
0025             break
0026         with open(dstlist, 'r') as f:
0027             lines = f.readlines()
0028         numDSTsThisRun = len(lines)
0029         filesthisrun = 0
0030         totaldsts += numDSTsThisRun
0031         jobfile = joblistdir + 'job-' + str(jobnum) + '.list'
0032         if os.path.exists(jobfile):
0033             os.remove(jobfile)
0034         jf = open(jobfile, 'a')
0035         for line in lines:
0036             jf.write(line)
0037             filesthisjob += 1
0038             filesthisrun += 1
0039             # if we hit the last line, don't increment jobnum and open a new
0040             # job file -- this is done at the start of the next DST list
0041             if filesthisrun == numDSTsThisRun:
0042                 break
0043             # otherwise, once we hit filesperjob we increment jobnum and start
0044             # a new job file
0045             if filesthisjob >= filesperjob:
0046                 # move to the next job file
0047                 filesthisjob = 0
0048                 jobnum += 1
0049                 jf.close()
0050                 jobfile = joblistdir + 'job-' + str(jobnum) + '.list'
0051                 if os.path.exists(jobfile):
0052                     os.remove(jobfile)
0053                 jf = open(jobfile, 'a')
0054                 if jobnum >= maxjobs:
0055                     break
0056         jf.close()
0057         jobnum += 1
0058     print(f'Found {totaldsts} DST files')
0059     print(f'Wrote {jobnum} per-job DST lists with <= {filesperjob} DSTs each')
0060