benchmark

#! /usr/bin/env python3

import sys
import os
import copy
import subprocess

# check that CMSSW_BASE is set
if not 'CMSSW_BASE' in os.environ:
    sys.stderr.write('Error: the CMS environment is not set up, please run "cmsenv" or "eval `scram runtime -sh`".\n')
    sys.exit(1)

# CMSSW python language
import FWCore.ParameterSet.Config as cms

# local packages
from common import *
from options import OptionParser, printCommonArgs
from multirun import *
from slot import Slot

if __name__ == "__main__":
  parser = OptionParser()
  opts = parser.parse(sys.argv[1:])

  # common options for multirun
  options = {
    'verbose'             : opts.verbose,
    'plumbing'            : opts.plumbing,
    'warmup'              : opts.warmup,
    'events'              : opts.events,
    'resolution'          : opts.event_resolution,
    'skipevents'          : opts.event_skip,
    'repeats'             : opts.repeats,
    'wait'                : opts.wait,
    'jobs'                : opts.jobs,
    'threads'             : opts.threads,
    'streams'             : opts.streams,
    'gpus_per_job'        : opts.gpus_per_job,
    'allow_hyperthreading': opts.allow_hyperthreading,
    'set_numa_affinity'   : opts.numa_affinity,
    'set_cpu_affinity'    : opts.cpu_affinity,
    'set_gpu_affinity'    : opts.gpu_affinity,
    'slots'               : opts.slots,
    'executable'          : opts.executable,
    'data'                : open(opts.csv, 'wt', buffering=1, encoding='utf-8') if opts.csv else None,
    'header'              : opts.csvheader,
    'logdir'              : opts.logdir if opts.logdir else None,
    'tmpdir'              : opts.tmpdir,
    'keep'                : opts.keep,
    'automerge'           : opts.automerge,
    'autodelete'          : opts.autodelete,
    'autodelete_delay'    : opts.autodelete_delay,
    'debug_cpu_usage'     : opts.debug_cpu_usage,
    'debug_affinity'      : opts.debug_affinity,
  }

  # print a system overview
  info()

  # check the available cpus
  cpus = get_cpu_info()
  if options['allow_hyperthreading']:
    count = sum(len(cpu.hardware_threads) for cpu in cpus.values())
  else:
    count = sum(len(cpu.physical_processors) for cpu in cpus.values())

  # autodetermine either the number of jobs ot the nuber of threads per job
  if options['threads'] is None and options['jobs'] is None:
    sys.stderr.write('%s: error: either the number of jobs ot the nuber of threads per job must be specified\n' % sys.argv[0])
  elif options['threads'] is None:
    options['threads'] = count // options['jobs']
  elif options['jobs'] is None:
    options['jobs'] = count // options['threads']

  # the number of streams defaults to the number of threads per job
  if options['streams'] is None:
    options['streams'] = options['threads']

  process = parseProcess(opts.config)

  # print configuration
  if opts.verbose:
    printCommonArgs(options)
    
  if opts.input_benchmark:
    collections = opts.input_collections.split(',')

    # read the list of input collections from the ReadBranches of a framework job report
    if opts.input_xml:
        collections = get_read_branch_names(opts.input_xml)

    # prepare a trimmed down configuration for benchmarking only reading the input data
    io_process = copy.deepcopy(process)
    io_process.prefetch = cms.EDAnalyzer("GenericConsumer",
        eventProducts = cms.untracked.vstring(collections)
    )
    io_process.path = cms.Path(io_process.prefetch)
    io_process.schedule = cms.Schedule(io_process.path)
    if 'PrescaleService' in io_process.__dict__:
      del io_process.PrescaleService

    # benchmark reading the input data
    print('Benchmarking only the input')
    io_options = dict(options, repeats = opts.input_benchmark, logdir = None, keep = [], data = None)
    multiCmsRun(io_process, **io_options)
    print()
    # wait the required number of seconds between the I/O benchmark and the actual measurement
    if opts.wait > 0:
      time.sleep(opts.wait)

  if opts.reference_benchmark:
    # benchmark the same configuration using the reference CMSSW release
    if 'CMSSW_RELEASE_BASE' in os.environ:
      # query scram for the reference release environment
      release_base = os.environ['CMSSW_RELEASE_BASE']
      variables = subprocess.Popen(['scram', 'runtime', '-sh'], cwd=release_base, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True).communicate()[0].splitlines()
      # make a copy of the environment and update it
      environ = copy.deepcopy(os.environ)
      export = re.compile(r'export +(\w+)="(.*)";?')
      unset = re.compile(r'unset +((\w| )+);?')
      for line in variables:
        # set or update environment variables
        matches = export.match(line)
        if matches:
          name = matches[1]
          value = matches[2]
          environ[name] = value
        # unset environment variables
        matches = unset.match(line)
        if matches:
          names = matches[1].split()
          for name in names:
            if name in environ:
              del environ[name]
      # run the benchmark with the modified environment
      print('Benchmarking %s with the reference release at %s' % (opts.config, release_base))
      if opts.logdir:
          ref_logdir = opts.logdir + '_reference'
      else:
          ref_logdir = None
      if opts.csv:
          name, ext = os.path.splitext(opts.csv)
          ref_csv = name + '_reference' + ext
          ref_data = open(ref_csv, 'wt', buffering=1, encoding='utf-8')
      else:
          ref_data = None
      ref_options = dict(options, logdir = ref_logdir, data = ref_data, environ = environ)
      multiCmsRun(process, **ref_options)
      print()
    else:
      print('$CMSSW_RELEASE_BASE is not set, cannot benchmark the reference release')
    print()

  print('Benchmarking %s' % opts.config)
  multiCmsRun(process, **options)

  # close the CSV file
  if options['data']:
    options['data'].close()