resource_analyzer.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
 
import ROOT
import math
import json
import csv
 
 
result = {'modules': {}, 'sums': {}}
 
# extract module statistics and memory usage from log file
log = open('resources.log', errors='ignore')
in_summary = False
for line in log:
    if in_summary and not line.startswith('================================================================================='):
        module = line.split('|')[0].strip()
        mem = max(int(line.split('|')[2]), 0)
        time_str = line.split('|')[4].split('+')
        time = float(time_str[0])
        time_error = float(time_str[1][1:])
        data = {'mem': mem, 'time': time, 'time_error': time_error, 'mem_init': 0, 'mem_begin_run': 0, 'mem_event': 0, 'mem_end_run': 0, 'mem_term': 0, 'mem_total': 0}
        if module.startswith('Sum_'):
            result['sums'][module[4:]] = data
            continue
        if module == 'Total':
            events = int(line.split('|')[1]) - 1
            factor = events / (events + 1.0)
            result['total'] = {'mem': mem, 'time': time * factor, 'time_error': time_error * factor}
            break
        if module in result['modules'].keys():
            result['modules'][module]['mem'] += data['mem']
            result['modules'][module]['time'] += data['time']
            result['modules'][module]['time_error'] = math.sqrt(result['modules'][module]['time_error']**2 + data['time_error']**2)
        else:
            result['modules'][module] = data
    if line.startswith('Name          ') and line.endswith(' |      Calls | Memory(MB) |    Time(s) |     Time(ms)/Call\n'):
        in_summary = True
    if line.startswith('[INFO] Virtual Memory usage at termination [MB]:'):
        result['vmem'] = line.split()[-1]
    if line.startswith('[INFO] Rss Memory usage at termination [MB]    :'):
        result['rss'] = line.split()[-1]
 
# extract size information from mdst and dst output
ROOT.gErrorIgnoreLevel = ROOT.kWarning + 1
for type in ['mdst', 'dst']:
    result[f'{type}_branches'] = {}
    file = ROOT.TFile(f'{type}.root')
    tree = file.Get('tree')
    root_events = tree.GetEntriesFast() * 1.0
    for branch in tree.GetListOfBranches():
        result[f'{type}_branches'][branch.GetName()] = branch.GetZipBytes('*') / root_events / 1024
    result[f'{type}_size'] = tree.GetZipBytes() / root_events / 1024
 
# extract memory usage from csv file
with open('memcheck.csv', 'r', errors='replace') as csvFile:
    reader = csv.DictReader(csvFile)
    for row in reader:
        module = row['name']
        if module in result['modules'].keys():
            result['modules'][module]['mem_init'] += float(row['memory init'])
            result['modules'][module]['mem_begin_run'] += float(row['memory begin_run'])
            result['modules'][module]['mem_event'] += float(row['memory event'])
            result['modules'][module]['mem_end_run'] += float(row['memory end_run'])
            result['modules'][module]['mem_term'] += float(row['memory term'])
            result['modules'][module]['mem_total'] += float(row['memory total'])
 
# write result file
with open('resources.json', 'w') as resources_file:
    json.dump(result, resources_file)