benchmark.py

#!/usr/bin/env python3
 

 

 
import basf2 as b2
from generators import add_kkmc_generator
from simulation import add_simulation
from validation_gt import get_validation_globaltags
from rawdata import add_unpackers
from reconstruction import add_reconstruction
from argparse import ArgumentParser
import os
import glob
import sys
 
# parse command line options
parser = ArgumentParser(description="Measure the execution time.")
parser.add_argument(
    "-m",
    "--multiplicity",
    default="high",
    choices=["high", "low", "data"],
    help="Multiplicity or type of events",
)
parser.add_argument("-l", "--limits", help="Name of file containing limits")
parser.add_argument("-f", "--file", help="Name of benchmark output file")
parser.add_argument("-c", "--csv", help="Name of statistics csv file")
args = parser.parse_args()
 
# create path and reduce log level
main = b2.create_path()
b2.set_log_level(b2.LogLevel.ERROR)
 
if args.multiplicity == "data":
    # Global Tag needed as these are Raw data
    b2.conditions.override_globaltags(get_validation_globaltags())
 
    input_files = glob.glob(
        os.environ.get("BELLE2_VALIDATION_DATA_DIR", "")
        + "/rawdata/physics.0010.05095*.root"
    )
    main.add_module("RootInput", inputFileNames=input_files)
 
    main.add_module("EventInfoPrinter").set_log_level(b2.LogLevel.INFO)
 
    # gearbox and geometry
    main.add_module("Gearbox")
    main.add_module("Geometry", useDB=True)
 
    # unpacking
    add_unpackers(main)
 
else:
    # specify number of events to be generated
    main.add_module("EventInfoSetter", evtNumList=[1000])
    main.add_module("EventInfoPrinter").set_log_level(b2.LogLevel.INFO)
 
    if args.multiplicity == "high":
        # generate BBbar events if high multiplicity is selected
        main.add_module("EvtGenInput")
 
    elif args.multiplicity == "low":
        # generate mu pair events if low multiplicity is selected
        add_kkmc_generator(main, "mu+mu-")
 
    # detector and L1 trigger simulation
    add_simulation(
        main,
        bkgfiles=glob.glob(
            os.environ.get("BELLE2_BACKGROUND_DIR", "/sw/belle2/bkg") + "/*.root"
        ),
    )
 
# reconstruction
add_reconstruction(main)
 
# process events, print call statistics, and write statistics to a csv file
b2.process(main, calculateStatistics=True)
print(b2.statistics)
if args.csv is not None:
    b2.statistics.csv(args.csv)
 
# read limits
limits = {}
limits_file = args.limits
default_limits_file = args.multiplicity + ".limits"
if limits_file is None and os.path.isfile(default_limits_file):
    limits_file = default_limits_file
if limits_file is not None:
    for line in open(limits_file).readlines():
        entries = line.split()
        limits[entries[0]] = float(entries[1])
        if len(entries) > 2:
            limits[entries[0]] /= float(entries[2])
 
# get execution times
categories = []
times = {}
for module in b2.statistics.modules:
    if "Sum_" not in module.name:
        continue
    category = module.name[4:]
    categories.append(category)
    if category not in times.keys():
        times[category] = 0
    times[category] += module.time_mean(b2.statistics.EVENT) * 1e-6
 
# open output file
output = None
if args.file is not None:
    output = open(args.file, "w")
 
# print benchmark results and write them to the output file
b2.set_log_level(b2.LogLevel.INFO)
max_fraction = -1
for category in categories:
    if category not in times.keys():
        continue
    time = times[category]
    message = f"Execution time per event for {category} is {time:.0f} ms"
    fraction = -1
    if category in limits.keys():
        fraction = time / limits[category]
        if fraction > max_fraction:
            max_fraction = fraction
        message += f" = {fraction:.0%} of the limit."
        if fraction <= 0.9:
            b2.B2INFO(message)
        elif fraction <= 1:
            b2.B2WARNING(message)
        else:
            b2.B2ERROR(message)
    else:
        b2.B2INFO(message)
 
    if output is not None:
        output.write(f"{category} {time:.2f}")
        if fraction >= 0:
            output.write(f" {fraction:.4f}")
        output.write("\n")
 
if output is not None:
    output.close()
 
# fail if above limit
sys.exit(0 if max_fraction <= 1 else 1)