Belle II Software  release-05-01-25
submit_create_mcruns.py
1 #!/usr/bin/env python3
2 # -*- coding: utf-8 -*-
3 
4 # Script using CAF backends to submit jobs for generating of mc files with PXDSimHits for
5 # PXD gain calibration.
6 #
7 # basf2 submit_create_mcruns.py -- --backend='local' --outputdir='pxd_mc_phase2' --runLow=4000 --runHigh=6522 --expNo=3
8 #
9 # author: benjamin.schwenker@phys.uni-goettingen.de
10 
11 import sys
12 import os
13 import pickle
14 from caf import backends
15 from pathlib import Path
16 from caf.utils import IoV
17 
18 
19 _tempdir_name = 'tmp_create_mc'
20 
21 
22 _number_of_sets = 10
23 
24 
25 _tag = 'Calibration_Offline_Development'
26 
27 
28 _bg = '/group/belle2/BGFile/OfficialBKG/15thCampaign/phase2'
29 
30 
31 def run_job_submission(backend, mc_iov_list):
32 
33  # List of jobs to be submitted
34  jobs = []
35 
36  for mc_iov in mc_iov_list:
37  for setNo in range(0, _number_of_sets):
38  expNo = mc_iov.exp_low
39  runNo = mc_iov.run_low
40  job = backends.Job("Create_Exp_{}_Run_{}_Set_{}".format(expNo, runNo, setNo))
41  job.output_dir = str(Path(_tempdir_name, job.name).absolute())
42  job.working_dir = str(Path(_tempdir_name, job.name).absolute())
43  job.cmd = [
44  'basf2',
45  'create_mcrun.py',
46  '--',
47  '--tag={}'.format(_tag),
48  '--bg={}'.format(_bg),
49  '--expNo={}'.format(expNo),
50  '--runNo={}'.format(runNo),
51  '--setNo={}'.format(setNo)]
52 
53  # Any files appended to this list will be place in the Job's working directory
54  job.input_sandbox_files.append("create_mcrun.py")
55 
56  # Sometimes you need to configure how the backend will treat your Job but in a backend specific way
57  # Any arguments to deal with this should be created in the backend_args dictionary
58  # Here we set the batch queue used to submit the job
59  if isinstance(backend, backends.LSF):
60  job.backend_args['queue'] = 's'
61  elif isinstance(backend, backends.PBS):
62  job.backend_args['queue'] = 'short'
63 
64  # Append to job to jobs
65  jobs.append(job)
66 
67  # Submit alls jobs at once
68  backend.submit(jobs)
69 
70  from basf2 import B2INFO
71  import time
72 
73  # Jobs provide a job.ready() function which returns True once the Job (or all its SubJobs) have finished their running
74  ready = False
75  while not ready:
76  ready = all(job.ready() for job in jobs)
77  B2INFO("Not finished all jobs yet")
78  time.sleep(5)
79  B2INFO("Finished")
80 
81 
82 def main(backend_choice, mc_iov_list, output_dir):
83  # Create a backend of our choosing
84  if backend_choice == "qsub":
85  backend = backends.PBS()
86  elif backend_choice == "bsub":
87  backend = backends.LSF()
88  elif backend_choice == "local":
89  backend = backends.Local(max_processes=20)
90  else:
91  print("How did I get here?")
92  sys.exit(1)
93  run_job_submission(backend, mc_iov_list)
94  if backend_choice == 'local':
95  # Should really be closing down the processing pool once we're done submitting.
96  backend.join()
97 
98  # Find all output files
99  import glob
100  output_files = glob.glob(str(Path(_tempdir_name).absolute()) + '/**/beam*.root')
101 
102  # Copy all outputs to output_dir
103  import shutil
104  for file_name in output_files:
105  shutil.move(file_name, output_dir)
106 
107 
108 if __name__ == "__main__":
109 
110  import argparse
111  parser = argparse.ArgumentParser(description="Submit jobs for creattion of SimHits for a run with user specified ExpRuns")
112  parser.add_argument('--outputdir', default='pxd_mc_phase2', type=str,
113  help='Add all simulated files to this folder. Create it if it does not exist.')
114  parser.add_argument('--backend', default='local', type=str, help='Where <backend> is one of: local, qsub, bsub.')
115  parser.add_argument('--runLow', default=0, type=int, help='Compute mask for specific IoV')
116  parser.add_argument('--runHigh', default=-1, type=int, help='Compute mask for specific IoV')
117  parser.add_argument('--expNo', default=3, type=int, help='Compute mask for specific IoV')
118  args = parser.parse_args()
119 
120  # Directory for collecting generated files for later usage in PXD gain calibration
121  output_dir = str(Path(args.outputdir).absolute())
122 
123  # Make sure output_dir exists
124  if not os.path.exists(output_dir):
125  os.makedirs(output_dir)
126 
127  # Set the IoV range for mc calibration
128  iov_to_calibrate = IoV(exp_low=args.expNo, run_low=args.runLow, exp_high=args.expNo, run_high=args.runHigh)
129 
130  # Access files_to_iovs for beam runs
131  with open("file_iov_map.pkl", 'br') as map_file:
132  files_to_iovs = pickle.load(map_file)
133 
134  # Set of all IoVs for beam runs (data). Potentially very, very large.
135  data_iov_set = set(files_to_iovs.values())
136 
137  # List of IoV for mc runs to be submitted to backends
138  mc_iov_list = []
139  for data_iov in data_iov_set:
140  if iov_to_calibrate.contains(data_iov):
141  mc_iov_list.append(data_iov)
142 
143  print('Number selected iovs for mc generation : {}'.format(len(mc_iov_list)))
144 
145  main(args.backend, mc_iov_list, output_dir)
main
int main(int argc, char **argv)
Run all tests.
Definition: test_main.cc:77
backends.LSF
Definition: backends.py:1567
backends.Job
Definition: backends.py:310
backends.PBS
Definition: backends.py:1304
backends.Local
Definition: backends.py:878