Belle II Software  release-05-01-25
cli.py
1 #!/usr/bin/env python3
2 import argparse
3 import os
4 import tempfile
5 import shutil
6 import json
7 from concurrent.futures import ThreadPoolExecutor
8 import difflib
9 import pandas as pd
10 
11 import basf2
12 from conditions_db import cli_download, ConditionsDB, encode_name
13 from softwaretrigger import db_access
14 
15 
16 class HashableCut(dict):
17  """Small helper class as the difflib does not understand dicts directly (as they are not hashable)"""
18 
19  def __hash__(self):
20  """Create a hash for the object out of the json string"""
21  return hash(json.dumps(self))
22 
23 
25  """Helper class to translate the user-specified database(s) into parameters for basf2"""
26 
27  def __init__(self, command_argument):
28  """Init the stored databases and exp/run from the specified command argument"""
29 
30  self._database = []
31 
32  self._experiment = 99999
33 
34  self._run = 99999
35 
36  # If given, use the experiment run from the command_argument
37  split_argument = command_argument.split(":")
38  if len(split_argument) == 2:
39  command_argument, exp_run = split_argument
40 
41  if exp_run != "latest":
42  try:
43  self._experiment, self._run = map(int, exp_run.split("/"))
44  except BaseException:
45  raise argparse.ArgumentTypeError(
46  f"Do not understand the exp/run argument '{exp_run}'")
47 
48  elif len(split_argument) != 1:
49  raise argparse.ArgumentTypeError(
50  f"Do not understand the database argument '{command_argument}'")
51 
52  # Now split up the databases
53  self._database = command_argument.split(",")
54 
55  # However make sure we have them in the correct format
56  def normalize(database):
57  # In case a local file is specified we can just use it directly
58  if os.path.exists(database):
59  if os.path.basename(database) != "database.txt":
60  database = os.path.join(database, "database.txt")
61 
62  return database
63 
64  self._database = list(map(normalize, self._database))
65 
66  def set_database(self):
67  """
68  Set the basf2 database chain according to the specified databases.
69  Before that, clean up and invalidate everything from th database.
70 
71  The distinction between file databases and global databases is done
72  via the fact of a file/folder with this name exists or not.
73  """
74  from ROOT import Belle2
75  Belle2.DBStore.Instance().reset()
76 
77  basf2.reset_database(False)
78  basf2.conditions.override_globaltags()
79 
80  for database in self._database:
81  if os.path.exists(database):
82  basf2.conditions.prepend_testing_payloads(database)
83  else:
84  basf2.conditions.prepend_globaltag(database)
85 
86  db_access.set_event_number(evt_number=0, run_number=int(self._run),
87  exp_number=int(self._experiment))
88 
89  def get_all_cuts(self):
90  """
91  Get all cuts stored in the database(s)
92  and sort them according to base_identifier, cut_identifier.
93  """
94  self.set_database()
95 
96  all_cuts = db_access.get_all_cuts()
97  all_cuts = sorted(all_cuts,
98  key=lambda cut: (cut["Base Identifier"], cut["Cut Identifier"]))
99  all_cuts = list(map(HashableCut, all_cuts))
100  return all_cuts
101 
102 
103 def diff_function(args):
104  """
105  Show the diff between two specified databases.
106  """
107  first_database_cuts = args.first_database.get_all_cuts()
108  second_database_cuts = args.second_database.get_all_cuts()
109 
110  diff = difflib.SequenceMatcher(
111  a=list(map(str, first_database_cuts)), b=list(map(str, second_database_cuts)))
112 
113  def print_cut(cut, prefix=" "):
114  if prefix == "-":
115  print("\x1b[31m", end="")
116  elif prefix == "+":
117  print("\x1b[32m", end="")
118  print(prefix, cut)
119  print("\x1b[0m", end="")
120 
121  def print_cuts(prefix, cuts):
122  for c in cuts:
123  print_cut(c, prefix)
124 
125  for tag, i1, i2, j1, j2 in diff.get_opcodes():
126  if tag == "equal":
127  if args.only_changes:
128  continue
129  print_cuts(" ", diff.b[j1:j2])
130  if tag in ["delete", "replace"]:
131  print_cuts("-", diff.a[i1:i2])
132  if tag in ["insert", "replace"]:
133  print_cuts("+", diff.b[j1:j2])
134 
135 
136 def add_cut_function(args):
137  """
138  Add a cut with the given parameters and also add it to the trigger menu.
139  """
140  args.database.set_database()
141 
142  db_access.upload_cut_to_db(cut_string=args.cut_string, base_identifier=args.base_identifier,
143  cut_identifier=args.cut_identifier, prescale_factor=args.prescale_factor,
144  reject_cut=args.reject_cut.lower() == "true", iov=None)
145  trigger_menu = db_access.download_trigger_menu_from_db(args.base_identifier,
146  do_set_event_number=False)
147  cuts = list(trigger_menu.getCutIdentifiers())
148 
149  if args.cut_identifier not in cuts:
150  cuts.append(args.cut_identifier)
151 
152  db_access.upload_trigger_menu_to_db(args.base_identifier, cuts,
153  accept_mode=trigger_menu.isAcceptMode(), iov=None)
154 
155 
156 def remove_cut_function(args):
157  """
158  Remove a cut with the given name from the trigger menu.
159  """
160  args.database.set_database()
161 
162  trigger_menu = db_access.download_trigger_menu_from_db(
163  args.base_identifier, do_set_event_number=False)
164  cuts = trigger_menu.getCutIdentifiers()
165 
166  cuts = [cut for cut in cuts if cut != args.cut_identifier]
167  db_access.upload_trigger_menu_to_db(
168  args.base_identifier, cuts, accept_mode=trigger_menu.isAcceptMode(), iov=None)
169 
170 
171 def print_function(args):
172  """
173  Print the cuts stored in the database(s).
174  """
175  cuts = args.database.get_all_cuts()
176  df = pd.DataFrame(cuts)
177 
178  if args.format == "pandas":
179  pd.set_option("display.max_rows", 500)
180  pd.set_option("display.max_colwidth", 200)
181  pd.set_option('display.max_columns', 500)
182  pd.set_option('display.width', 1000)
183  print(df)
184  elif args.format == "jira":
185  from tabulate import tabulate
186  print(tabulate(df, tablefmt="jira", showindex=False, headers="keys"))
187  elif args.format == "grid":
188  from tabulate import tabulate
189  print(tabulate(df, tablefmt="grid", showindex=False, headers="keys"))
190  elif args.format == "json":
191  import json
192  print(json.dumps(df.to_dict("records"), indent=2))
193  elif args.format == "list":
194  for base_identifier, cuts in df.groupby("Base Identifier"):
195  for _, cut in cuts.iterrows():
196  print(cut["Base Identifier"], cut["Cut Identifier"])
197  elif args.format == "human-readable":
198  print("Currently, the following menus and triggers are in the database")
199  for base_identifier, cuts in df.groupby("Base Identifier"):
200  print(base_identifier)
201  print("")
202  print("\tUsed triggers:\n\t\t" +
203  ", ".join(list(cuts["Cut Identifier"])))
204  print("\tIs in accept mode:\n\t\t" +
205  str(cuts["Reject Menu"].iloc[0]))
206  for _, cut in cuts.iterrows():
207  print("\t\tCut Name:\n\t\t\t" + cut["Cut Identifier"])
208  print("\t\tCut condition:\n\t\t\t" + cut["Cut Condition"])
209  print("\t\tCut prescaling\n\t\t\t" +
210  str(cut["Cut Prescaling"]))
211  print("\t\tCut is a reject cut:\n\t\t\t" +
212  str(cut["Reject Cut"]))
213  print()
214  else:
215  raise AttributeError(f"Do not understand format {args.format}")
216 
217 
218 def create_script_function(args):
219  """
220  Print the b2hlt_trigger commands to create a lobal database copy.
221  """
222  cuts = args.database.get_all_cuts()
223  df = pd.DataFrame(cuts)
224 
225  sfmt = 'b2hlt_triggers add_cut \
226 "{Base Identifier}" "{Cut Identifier}" "{Cut Condition}" "{Cut Prescaling}" "{Reject Cut}"'.format
227  if args.filename is None:
228  df.apply(lambda x: print(sfmt(**x)), 1)
229  else:
230  with open(args.filename, 'w') as f:
231  df.apply(lambda x: f.write(sfmt(**x) + '\n'), 1)
232 
233 
234 def iov_includes(iov_list, exp, run):
235  """
236  Comparison function between two IoVs (start, end) stored in the database and
237  the given exp/run combination.
238  """
239  # Dirty hack: replace -1 by infinity to make the comparison easier
240  copied_iov_list = iov_list[2:]
241  copied_iov_list = list(map(lambda x: x if x != -1 else float("inf"), copied_iov_list))
242 
243  exp_start, run_start, exp_end, run_end = copied_iov_list
244 
245  return (exp_start, run_start) <= (exp, run) <= (exp_end, run_end)
246 
247 
248 def download_function(args):
249  """
250  Download the trigger cuts in the given database to disk and set their IoV to infinity.
251  """
252  if len(args.database._database) != 1:
253  raise AttributeError("Can only download from a single database! Please do not specify more than one.")
254 
255  global_tag = args.database._database[0]
256 
257  # The following is an adapted version of cli_download
258  os.makedirs(args.destination, exist_ok=True)
259 
260  db = ConditionsDB()
261  req = db.request("GET", f"/globalTag/{encode_name(global_tag)}/globalTagPayloads",
262  f"Downloading list of payloads for {global_tag} tag")
263 
264  download_list = {}
265  for payload in req.json():
266  name = payload["payloadId"]["basf2Module"]["name"]
267  if not name.startswith("software_trigger_cut"):
268  continue
269 
270  local_file, remote_file, checksum, iovlist = cli_download.check_payload(args.destination, payload)
271 
272  new_iovlist = list(filter(lambda iov: iov_includes(iov, args.database._experiment, args.database._run), iovlist))
273  if not new_iovlist:
274  continue
275 
276  if local_file in download_list:
277  download_list[local_file][-1] += iovlist
278  else:
279  download_list[local_file] = [local_file, remote_file, checksum, iovlist]
280 
281  # do the downloading
282  full_iovlist = []
283  failed = 0
284  with ThreadPoolExecutor(max_workers=20) as pool:
285  for iovlist in pool.map(lambda x: cli_download.download_file(db, *x), download_list.values()):
286  if iovlist is None:
287  failed += 1
288  continue
289 
290  full_iovlist += iovlist
291 
292  dbfile = []
293  for iov in sorted(full_iovlist):
294  # Set the IoV intentionally to 0, 0, -1, -1
295  iov = [iov[0], iov[1], 0, 0, -1, -1]
296  dbfile.append("dbstore/{} {} {},{},{},{}\n".format(*iov))
297  with open(os.path.join(args.destination, "database.txt"), "w") as txtfile:
298  txtfile.writelines(dbfile)
299 
300 
301 def main():
302  """
303  Main function to be called from b2hlt_triggers.
304  """
305  parser = argparse.ArgumentParser(
306  description="""
307 Execute different actions on stored trigger menus in the database.
308 
309 Call with `%(prog)s [command] --help` to get a description on each command.
310 Please also see the examples at the end of this help.
311 
312 Many commands require one (or many) specified databases. Different formats are possible.
313 All arguments need to be written in quotation marks.
314 * "online" Use the latest version in the "online" database
315  (or any other specified global tag).
316 * "online:latest" Same as just "online", makes things a bit clearer.
317 * "online:8/345" Use the version in the "online" database (or any other specified global tag)
318  which was present in exp 8 run 345.
319 * "localdb:4/42" Use the local database specified in the given folder for the given exp/run.
320 * "localdb/database.txt" It is also possible to specify a file directly.
321 * "online,localdb" First look in localdb, then in the online GT
322 * "online,localdb:9/1" Can also be combined with the exp/run (It is then valid for all database accesses)
323 
324 Examples:
325 
326 * Check what has changed between 8/1 and 9/1 in the online GT.
327 
328  %(prog)s diff --first-database "online:8/1" --second-database "online:9/1" --only-changes
329 
330 * Especially useful while editing trigger cuts and menus: check what has changed between the latest
331  version online and what is currently additionally in localdb
332 
333  %(prog)s diff --first-database "online:latest" --second-database "online,localdb:latest"
334 
335  This use case is so common, it is even the default
336 
337  %(prog)s diff
338 
339 * Print the latest version of the cuts in online (plus what is defined in the localdb) in a human-friendly way
340 
341  %(prog)s print
342 
343 * Print the version of the cuts which was present in 8/1 online in a format understandable by JIRA
344  (you need to have the tabulate package installed)
345 
346  %(prog)s print --database "online:8/1" --format jira
347 
348 * Add a new skim cut named "accept_b2bcluster_3D" with the specified parameters and upload it to localdb
349 
350  %(prog)s add_cut skim accept_b2bcluster_3D "[[nB2BCC3DLE >= 1] and [G1CMSBhabhaLE < 2]]" 1 False
351 
352 * Remove the cut "accept_bhabha" from the trigger menu "skim"
353 
354  %(prog)s remove_cut skim accept_bhabha
355 
356 * Download the latest state of the triggers into the folder "localdb", e.g. to be used for local studies
357 
358  %(prog)s download
359 
360  """,
361  formatter_class=argparse.RawDescriptionHelpFormatter,
362  usage="%(prog)s command"
363  )
364  parser.set_defaults(func=lambda *args: parser.print_help())
365  subparsers = parser.add_subparsers(title="command",
366  description="Choose the command to execute")
367 
368  # diff command
369  diff_parser = subparsers.add_parser("diff", help="Compare the trigger menu in two different databases.",
370  formatter_class=argparse.RawDescriptionHelpFormatter,
371  description="""
372 Compare the two trigger menus present in the two specified databases
373 (or database chains) for the given exp/run combination (or the latest
374 version).
375 Every line in the output is one trigger line. A "+" in front means the
376 trigger line is present in the second database, but not in the first.
377 A "-" means exactly the opposite. Updates trigger lines will show up
378 as both "-" and "+" (with different parameters).
379 
380 The two databases (or database chains) can be specified as describes
381 in the general help (check b2hlt_triggers --help).
382 By default, the latest version of the online database will be
383 compared with what is defined on top in the localdb.
384  """)
385  diff_parser.add_argument("--first-database", help="First database to compare. Defaults to 'online:latest'.",
386  type=DownloadableDatabase, default=DownloadableDatabase("online:latest"))
387  diff_parser.add_argument("--second-database", help="Second database to compare. Defaults to 'online,localdb:latest'.",
388  type=DownloadableDatabase, default=DownloadableDatabase("online,localdb:latest"))
389  diff_parser.add_argument(
390  "--only-changes", help="Do not show unchanged lines.", action="store_true")
391  diff_parser.set_defaults(func=diff_function)
392 
393  # print command
394  print_parser = subparsers.add_parser("print", help="Print the cuts stored in the given database.",
395  formatter_class=argparse.RawDescriptionHelpFormatter,
396  description="""
397 Print the defined trigger menu and trigger cuts in a human-friendly
398 (default) or machine-friendly way.
399 The database (or database chain) needs to be specified in the general
400 help (check b2hlt_triggers --help).
401 
402 For additional formatting options please install the tabulate package with
403 
404  pip3 install --user tabulate
405 
406 By default the latest version on the online database and what is defined on
407 top in the localdb will be shown.
408  """)
409  print_parser.add_argument("--database", help="Which database to print. Defaults to 'online,localdb:latest'.",
410  type=DownloadableDatabase, default=DownloadableDatabase("online,localdb:latest"))
411  choices = ["human-readable", "json", "list", "pandas"]
412  try:
413  from tabulate import tabulate
414  choices += ['jira', 'grid']
415  except ImportError:
416  pass
417  print_parser.add_argument("--format", help="Choose the format how to print the trigger cuts. "
418  "To get access to more options please install the tabulate package using pip",
419  choices=choices, default="human-readable")
420  print_parser.set_defaults(func=print_function)
421 
422  # create-script command
423  create_script_parser = subparsers.add_parser(
424  "create_script",
425  help="Create b2hlt_triggers command to create a online globaltag copy.",
426  formatter_class=argparse.RawDescriptionHelpFormatter,
427  description="""
428 Generate the required b2hlt_trigger commands to reproduce an online globaltag for a given exp/run
429 number to create a local database version of it.
430  """)
431  create_script_parser.add_argument("--database", help="Which database to print. Defaults to 'online:latest'.",
432  type=DownloadableDatabase, default=DownloadableDatabase("online:latest"))
433  create_script_parser.add_argument("--filename", default=None,
434  help="Write to given filename instead of stdout.")
435  create_script_parser.set_defaults(func=create_script_function)
436 
437  # add_cut command
438  add_cut_parser = subparsers.add_parser("add_cut", help="Add a new cut.",
439  formatter_class=argparse.RawDescriptionHelpFormatter,
440  description="""
441 Add a cut with the given properties and upload it into the localdb database.
442 After that, you can upload it to the central database, to e.g. staging_online.
443 
444 As a base line for editing, a database much be specified in the usual format
445 (check b2hlt_triggers --help).
446 It defaults to the latest version online and the already present changes in
447 localdb.
448 Please note that the IoV of the created trigger line and menu is set to infinite.
449  """)
450  add_cut_parser.add_argument("--database", help="Where to take the trigger menu from. Defaults to 'online,localdb:latest'.",
451  type=DownloadableDatabase, default=DownloadableDatabase("online,localdb:latest"))
452  add_cut_parser.add_argument("base_identifier",
453  help="base_identifier of the cut to add", choices=["filter", "skim"])
454  add_cut_parser.add_argument("cut_identifier",
455  help="cut_identifier of the cut to add")
456  add_cut_parser.add_argument("cut_string",
457  help="cut_string of the cut to add")
458  add_cut_parser.add_argument("prescale_factor", type=int,
459  help="prescale of the cut to add")
460  add_cut_parser.add_argument(
461  "reject_cut", help="Is the new cut a reject cut?")
462  add_cut_parser.set_defaults(func=add_cut_function)
463 
464  # remove_cut command
465  remove_cut_parser = subparsers.add_parser("remove_cut", help="Remove a cut of the given name.",
466  formatter_class=argparse.RawDescriptionHelpFormatter,
467  description="""
468 Remove a cut with the given base and cut identifier from the trigger menu
469 and upload the new trigger menu to the localdb.
470 After that, you can upload it to the central database, to e.g. staging_online.
471 
472 As a base line for editing, a database much be specified in the usual format
473 (check b2hlt_triggers --help).
474 It defaults to the latest version online and the already present changes in
475 localdb.
476 Please note that the IoV of the created trigger menu is set to infinite.
477 
478 The old cut payload will not be deleted from the database. This is not
479 needed as only cuts specified in a trigger menu are used.
480  """)
481  remove_cut_parser.add_argument("base_identifier",
482  help="base_identifier of the cut to delete", choices=["filter", "skim"])
483  remove_cut_parser.add_argument("cut_identifier",
484  help="cut_identifier of the cut to delete")
485  remove_cut_parser.add_argument("--database",
486  help="Where to take the trigger menu from. Defaults to 'online,localdb:latest'.",
487  type=DownloadableDatabase, default=DownloadableDatabase("online,localdb:latest"))
488  remove_cut_parser.set_defaults(func=remove_cut_function)
489 
490  # download command
491  download_parser = subparsers.add_parser("download", help="Download the trigger menu from the database.",
492  formatter_class=argparse.RawDescriptionHelpFormatter,
493  description="""
494 Download all software trigger related payloads from the specified database
495 into the folder localdb and create a localdb/database.txt. This is
496 especially useful when doing local trigger studies which should use the
497 latest version of the online triggers. By default, the latest
498 version of the online GT will be downloaded.
499 
500 Attention: this script will override a database defined in the destination
501 folder (default localdb)!
502 Attention 2: all IoVs of the downloaded triggers will be set to 0, 0, -1, -1
503 so you can use the payloads fro your local studies for whatever run you want.
504 This should not (never!) be used to upload or edit new triggers and
505 is purely a convenience function to synchronize your local studies
506 with the online database!
507 
508 Please note that for this command you can only specify a single database
509 (all others can work with multiple databases).
510  """)
511  download_parser.add_argument("--database",
512  help="Single database where to take the trigger menu from. Defaults to 'online:latest'.",
513  type=DownloadableDatabase, default=DownloadableDatabase("online:latest"))
514  download_parser.add_argument("--destination",
515  help="In which folder to store the output", default="localdb")
516  download_parser.set_defaults(func=download_function)
517 
518  args = parser.parse_args()
519  args.func(args)
softwaretrigger.cli.HashableCut
Definition: cli.py:16
softwaretrigger.cli.DownloadableDatabase.get_all_cuts
def get_all_cuts(self)
Definition: cli.py:89
Belle2::filter
std::map< ExpRun, std::pair< double, double > > filter(const std::map< ExpRun, std::pair< double, double >> &runs, double cut, std::map< ExpRun, std::pair< double, double >> &runsRemoved)
filter events to remove runs shorter than cut, it stores removed runs in runsRemoved
Definition: Splitter.cc:43
softwaretrigger.cli.DownloadableDatabase._run
_run
the run number, default (= latest) is 99999
Definition: cli.py:34
softwaretrigger.cli.DownloadableDatabase._experiment
_experiment
the experiment number, default (= latest) is 99999
Definition: cli.py:32
main
int main(int argc, char **argv)
Run all tests.
Definition: test_main.cc:77
conditions_db.ConditionsDB
Definition: __init__.py:144
Belle2::DBStore::Instance
static DBStore & Instance()
Instance of a singleton DBStore.
Definition: DBStore.cc:36
softwaretrigger.cli.DownloadableDatabase._database
_database
the specified databases
Definition: cli.py:30
softwaretrigger.cli.DownloadableDatabase.__init__
def __init__(self, command_argument)
Definition: cli.py:27
softwaretrigger.cli.DownloadableDatabase
Definition: cli.py:24
softwaretrigger.cli.HashableCut.__hash__
def __hash__(self)
Definition: cli.py:19
softwaretrigger.cli.DownloadableDatabase.set_database
def set_database(self)
Definition: cli.py:66