memnesia-report

#!/usr/bin/env python3

#
# Copyright (c) 2017-2021 Triad National Security, LLC
#                         All rights reserved.
#
# This file is part of the mpimemu project. See the LICENSE file at the
# top-level directory of this distribution.
#

import os
import sys
import operator
import shutil
import collections
import matplotlib.cm as cm
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import numpy as np


###############################################################################
class Util:

    @staticmethod
    def btomb(b):
        return float(b) / 1024.0 / 1024.0

    @staticmethod
    def mean(num_lst):
        return float(sum(num_lst)) / max(len(num_lst), 1)

    @staticmethod
    def sort_human(lst):
        return sorted(lst, key=lambda x: '{0:0>8}'.format(x).lower())


###############################################################################
class RunMetadata:
    def __init__(self, data):

        self.lines_consumed = 0

        self.data = {
            'Report Date Time': '',
            'Application Name': '',
            'Hostname': '',
            'MPI_COMM_WORLD Rank': 0,
            'MPI_COMM_WORLD Size': 0,
            'MPI Init Time (s)': 0.,
            'Number of smaps Captures Performed': 0,
            'High Memory Usage Watermark (MPI) (MB)': 0.,
            'High Memory Usage Watermark (Application + MPI) (MB)': 0.
        }
        assert(data[0] == '# [Run Info Begin]')
        # Skip header
        self.lines_consumed = 1
        meta_lines = []
        while True:
            if data[self.lines_consumed] == '# [Run Info End]':
                break
            meta_lines.append(data[self.lines_consumed].strip('# '))
            self.lines_consumed += 1

        for ml in meta_lines:
            kv = [x.strip(' ') for x in ml.split(':')]
            assert(len(kv) == 2)
            key = kv[0]
            val_str = kv[1]
            if key not in self.data:
                print("ERROR: '{}' not a recognized key...".format(key))
                assert(False)
            self.data[key] = self.cast(type(self.data[key]), val_str)

    def get_lines_consumed(self):
        return self.lines_consumed

    def cast(self, val_type, string):
        return val_type(string)

    def get_numpe(self):
        return self.data['MPI_COMM_WORLD Size']

    @staticmethod
    def emit_stats(meta_list, statf):

        header = '# Run Statistics '
        statf.write('{}{}\n'.format(header, '#' * (80 - len(header))))

        statf.write('# Report Date Time: {}\n'.format(
            meta_list[0].data['Report Date Time'])
        )

        statf.write('# Application: {}\n'.format(
            meta_list[0].data['Application Name'])
        )

        for kprefix in ['MPI_COMM_WORLD', 'MPI Init Time']:
            mpi_keys = [k for k in meta_list[0].data.keys()
                        if k.startswith(kprefix)]

            for k in mpi_keys:
                statf.write('# {}: {}\n'.format(k, meta_list[0].data[k]))

        statf.write(
            '# Number of Process Data Analyzed: {}\n'.format(len(meta_list))
        )

        for kprefix in ['Number of', 'High Memory Usage Watermark']:
            stat_keys = [k for k in meta_list[0].data.keys()
                         if k.startswith(kprefix)]
            RunMetadata.emit_min_max_aves(meta_list, stat_keys, statf)

    @staticmethod
    def emit_min_max_aves(meta_list, keys, statf):
        for k in keys:
            statf.write('# {}\n'.format(k))

            vals = [meta.data[k] for meta in meta_list]

            minidx, minval = min(enumerate(vals), key=operator.itemgetter(1))
            maxidx, maxval = max(enumerate(vals), key=operator.itemgetter(1))

            min_where = meta_list[minidx].data['Hostname']
            max_where = meta_list[maxidx].data['Hostname']
            ave = Util.mean(vals)

            statf.write('- Min: {}, Who: {}, Where: {}\n'.format(
                minval, minidx, min_where)
            )
            statf.write('- Max: {}, Who: {}, Where: {}\n'.format(
                maxval, maxidx, max_where)
            )
            statf.write('- Ave: {0:0.3f}\n'.format(ave))


###############################################################################
class TimeSeries:
    def __init__(self):
        self.data = {
            'times': [],
            'svals': []
        }

    def push(self, time, samp_val):
        self.data['times'].append(time)
        self.data['svals'].append(samp_val)

    def get(self, key):
        assert(key in self.data)
        return self.data[key]


###############################################################################
class TimeSeriesAccumulator:

    @staticmethod
    def deltas(tva):
        res = []
        for i in range(0, len(tva)):
            if i == 0:
                res.append(tva[i])
                continue
            t = tva[i][0]
            v = tva[i][1]
            v_prev = tva[i - 1][1]
            res.append((t, v - v_prev))

        return res

    @staticmethod
    def accumulate(e):
        res = {
            'MPI_MEM_USAGE': TimeSeries(),
            'ALL_MEM_USAGE': TimeSeries()
        }
        ut_map = {
            'MPI_MEM_USAGE': {},
            'ALL_MEM_USAGE': {}
        }
        delta_map = {
            'MPI_MEM_USAGE': [],
            'ALL_MEM_USAGE': []
        }
        for _, ts in e.rank_to_time_series.items():
            for tsk, tsi in ts.items():
                tva = list(zip(tsi.data['times'], tsi.data['svals']))
                delta_map[tsk] += TimeSeriesAccumulator.deltas(tva)

        for k, v in delta_map.items():
            ut = ut_map[k]
            for i in v:
                t = i[0]
                u = i[1]
                if t in ut:
                    ut[t] += u
                else:
                    ut[t] = u

        tus_map = {
            'MPI_MEM_USAGE': [],
            'ALL_MEM_USAGE': []
        }
        for mk, mv in ut_map.items():
            tu = ut_map[mk]
            tus = tus_map[mk]
            first = True
            for t, u in sorted(tu.items()):
                if first:
                    tus.append((t, u))
                    first = False
                    continue
                prev = tus[-1]
                tui = (t, prev[1] + u)
                tus.append(tui)

        for k, v in sorted(tus_map.items()):
            ts = res[k]
            for i in v:
                ts.push(i[0], i[1])

        return res


###############################################################################
class Experiment:
    def __init__(self, log_file):
        self.log_path = os.path.dirname(os.path.abspath(log_file))
        self.log_file_name = os.path.basename(log_file)
        self.log_file = '{}/{}'.format(self.log_path, self.log_file_name)
        self.run_meta = []
        '''
        Two level dict maps rank to usage type to time series.
        '''
        self.rank_to_time_series = collections.defaultdict(dict)
        self.agg_ts = None

    def get_num_species(self):
        return len(self.rank_to_time_series)

    def get_run_meta(self):
        meta = []
        meta.append(RunMetadata(self.log_file))
        return meta

    def get_app_name(self):
        return self.run_meta[0].data['Application Name']

    def crunch(self):
        with open(self.log_file, 'r') as f:
            content = [x.rstrip() for x in f.readlines()]
            num_lines = len(content)
            line_num = 0

            while (line_num < num_lines):
                rmeta = RunMetadata(content[line_num:])
                self.run_meta.append(rmeta)
                line_num += rmeta.get_lines_consumed()

                rank = rmeta.data['MPI_COMM_WORLD Rank']

                ts = {
                    'MPI_MEM_USAGE': TimeSeries(),
                    'ALL_MEM_USAGE': TimeSeries()
                }

                while (line_num < num_lines and
                       content[line_num] != '# [Run Info Begin]'):
                    ln = content[line_num]
                    # Skip comments.
                    if ln.startswith('#'):
                        line_num += 1
                        continue

                    ldata = ln.split(' ')
                    dtype = ldata[0]
                    # dfunc = ldata[1]
                    dtime = float(ldata[2])
                    dmem = float(ldata[3])
                    ts[dtype].push(dtime, dmem)

                    line_num += 1

                self.rank_to_time_series[rank] = ts

        self.agg_ts = TimeSeriesAccumulator.accumulate(self)

        RunMetadata.emit_stats(self.run_meta, sys.stdout)
        print('')


###############################################################################
class Colorer:
    def __init__(self, n_colors):
        self.n_colors = n_colors
        self.color_idx = 0
        self.pvalues = np.logspace(-0.1, 0, n_colors)
        self.pmin = self.pvalues[0]
        self.pmax = self.pvalues[-1]

    def norm(self, pval):
        denom = float(self.pmax - self.pmin)
        if (self.pmin == self.pmax):
            denom = 1.0
        return (pval - self.pmin) / denom

    def get_color(self):
        cm = Colorer.get_cm()
        c = cm(
            self.norm(self.pvalues[self.color_idx])
        )
        self.color_idx += 1
        return c

    @staticmethod
    def get_cm():
        # return plt.cm.copper
        # return plt.cm.viridis
        return plt.cm.inferno


###############################################################################
class Reporter:
    def __init__(self, e_to_plot):
        for e in e_to_plot:
            self.experiment = e
            self.e_numpe = e.run_meta[0].get_numpe()
            self.app_name = e.get_app_name()
            self.dname = self.get_usable_dir_name(
                             '{}-{}PE-{}'.format(
                                 self.app_name,
                                 self.e_numpe,
                                 'memnesia-report')
                         )
            # Create report directory.
            os.makedirs(self.dname)
            # Copy input data into directory.
            shutil.copy(e.log_file, self.dname)

            # Crunch and capture experiment data.
            with open('{}/{}.txt'.format(
                          self.dname,
                          'statistics'
                      ), 'w') as statf:
                RunMetadata.emit_stats(self.experiment.run_meta, statf)

            self.numpes = set()

            self.key2plot = {
                'ALL_MEM_USAGE': 'total-mem-usage-timeline.pdf',
                'MPI_MEM_USAGE': 'mpi-mem-usage-timeline.pdf',
                'AGG_ALL_MEM_USAGE': 'aggregate-total-mem-usage-timeline.pdf',
                'AGG_MPI_MEM_USAGE': 'aggregate-mpi-mem-usage-timeline.pdf'
            }
            self.ts2title = {
                'ALL_MEM_USAGE': 'Total Memory Usage',
                'MPI_MEM_USAGE': 'MPI Library Memory Usage',
                'AGG_ALL_MEM_USAGE': 'Aggregate Total Memory Usage',
                'AGG_MPI_MEM_USAGE': 'Aggregate MPI Library Memory Usage',
            }

            self.fig = None
            self.ax = None

            self.max_x = sys.float_info.min
            self.max_y = sys.float_info.min
            self.min_y = sys.float_info.max
            self.line_style = 'solid'
            self.line_width = 1.0
            self.title_font_size = 12
            self.axis_font_size = 12
            self.fig_width = 6
            self.fig_height = 3
            self.figsize = (self.fig_width, self.fig_height)

            for plotk in self.key2plot.keys():
                self.plot(plotk)

            print('# Report written to {}'.format(self.dname))

    def plot(self, plot_key):
        self.fig, self.ax = plt.subplots(figsize=self.figsize)
        ''' Color maps setup '''
        num_colors = self.experiment.get_num_species()
        if (plot_key in ['AGG_ALL_MEM_USAGE', 'AGG_MPI_MEM_USAGE']):
            num_colors = 1

        colorer = Colorer(num_colors)
        if (plot_key in ['AGG_ALL_MEM_USAGE', 'AGG_MPI_MEM_USAGE']):
            colorer = Colorer(4)
            for burn in range(0, 2):
                # burn a color
                colorer.get_color()
            real_key = 'ALL_MEM_USAGE'
            if plot_key == 'AGG_MPI_MEM_USAGE':
                real_key = 'MPI_MEM_USAGE'
            tsi = self.experiment.agg_ts[real_key]
            self.add_plot(
                self.ax, colorer, tsi
            )
        else:
            for rank, ts in self.experiment.rank_to_time_series.items():
                tsi = ts[plot_key]
                self.add_plot(
                    self.ax, colorer, tsi
                )

        scalf = 1.05
        xmin = 0
        xmax = self.max_x * scalf
        '''
        ymin = self.min_y / scalf
        ymax = self.max_y * scalf
        '''
        self.ax.set_xlim(xmin, xmax)
        # Show grids.
        self.ax.grid(linestyle='dotted')
        self.ax.grid(True)
        # Set y labels.
        self.ax.set_ylabel(
            'Memory Usage (MB)', fontsize=self.axis_font_size
        )
        # Set y-tick format.
        # self.ax.yaxis.set_major_formatter(FormatStrFormatter('%.1f'))
        plt.ticklabel_format(style='sci', axis='both', scilimits=(0, 0))
        # For 'x10^i'
        self.ax.xaxis.major.formatter._useMathText = True
        self.ax.yaxis.major.formatter._useMathText = True

        self.ax.set_title(
            self.ts2title[plot_key],
            fontsize=self.title_font_size
        )

        if (plot_key not in ['AGG_ALL_MEM_USAGE', 'AGG_MPI_MEM_USAGE']):
            self.add_colorbar_to_plot(num_colors)
        # self.add_mpi_init_shading_to_subplots()
        '''
        Set common y label.
        self.fig.text(
            0.06, 0.5,
            'Memory Usage (MB)',
            fontsize=self.axis_font_size,
            ha='center', va='center', rotation='vertical'
        )
        '''

        '''
        Only bottom-most axis will have x label.
        '''
        self.ax.set_xlabel(
            'Wall Time Since Initialization (s)', fontsize=self.axis_font_size
        )

        '''
        plt.legend(loc=2)
        '''

        self.fig.savefig(
            '{}/{}'.format(self.dname, self.key2plot[plot_key]),
            format='pdf',
            bbox_inches='tight',
            pad_inches=0.03
        )

    def add_mpi_init_shading_to_subplots(self):
        max_time = sys.float_info.min
        for meta in self.experiment.run_meta:
            max_time = max(max_time, meta.data['MPI Init Time (s)'])
        self.ax.axvspan(0, max_time, alpha=0.5, color='#f1f1f1')

    def add_colorbar_to_plot(self, num_colors):
        # To number of ranks.
        normalize = mcolors.Normalize(vmin=0, vmax=num_colors)
        smappaple = cm.ScalarMappable(norm=normalize, cmap=Colorer.get_cm())
        smappaple.set_array(num_colors)

        # Create an axis on the right side of ax. The width of cax will be 2%
        # of ax and the padding between cax and ax will be fixed at 0.05 inch.
        divider = make_axes_locatable(self.ax)
        cax = divider.append_axes(
            "right",
            size="2%",
            pad=0.05
        )

        tick_format = '%d'
        if (num_colors > 99):
            tick_format = '%.0e'

        cbar = plt.colorbar(
            smappaple,
            cax=cax,
            format=tick_format
        )
        cbar.set_ticks([0, num_colors])
        tick_labels = [0, num_colors - 1]
        if (num_colors == 1):
            cbar.set_ticklabels([tick_labels[0]])
        else:
            cbar.set_ticklabels(tick_labels)
        cbar.set_label('Global Process Identifier', rotation='vertical')

    def add_plot(self, target_ax, colorer, time_series):
        x = time_series.get('times')
        self.max_x = max(self.max_x, max(x))

        y = time_series.get('svals')
        self.min_y = min(self.min_y, min(y))
        self.max_y = max(self.max_y, max(y))

        #
        clr = colorer.get_color()
        #
        target_ax.plot(
            x, y,
            linestyle=self.line_style,
            linewidth=self.line_width,
            color=clr,
            alpha=0.9
        )
        '''
        target_ax.fill_between(
            range(len(y)), 0, y,
            interpolate=True, color=clr, alpha=0.2
        )
        '''

    def get_usable_dir_name(self, name):
        fidx = 1
        cwd = os.getcwd()
        target_dname = '{}/{}'.format(cwd, name)
        while (os.path.exists(target_dname)):
            target_dname = '{}/{}-{}'.format(cwd, name, fidx)
            fidx += 1

        return target_dname


###############################################################################
def usage():
    print('usage: memnesia-report out.memnesia [another.memnesia ...]')


###############################################################################
def check_args(argv):
    if len(argv) < 2:
        usage()
        exit(os.EX_USAGE)

    for arg in argv[1:]:
        if not os.path.isfile(arg):
            print("'{}' is not a file. Cannot continue.".format(arg))
            usage()
            exit(os.EX_USAGE)


###############################################################################
def process_experiment_data(data_files):
    experiments = []

    for data_file in data_files:
        print('# Processing {}'.format(data_file))
        e = Experiment(data_file)
        e.crunch()
        experiments.append(e)

    return experiments


###############################################################################
def main(argv=None):
    if argv is None:
        argv = sys.argv

    check_args(argv)

    experiments = process_experiment_data(argv[1:])

    Reporter(experiments)

    return os.EX_OK


###############################################################################
###############################################################################
###############################################################################
if __name__ == "__main__":
    sys.exit(main())