check_test_output.py

#!/usr/bin/env python3

import sys
import csv
from io import StringIO

def read_data():
    lines = sys.stdin.readlines()
    lines = list(filter(lambda line: 'RSS' in line or 'RAM' in line, lines))
    for line in lines:
        print(line.rstrip('\n'))

    data = {'serial': [], 'node': []}
    f = StringIO(''.join(lines))
    reader = csv.reader(f, delimiter=' ')
    for row in reader:
        line = {}
        if row[0] == 'MPI':
            assert(row[5] == 'RSS:')
            line['RSS'] = convert_number(row[6])
            assert(row[9] == 'HWM:')
            line['HWM'] = convert_number(row[10])
            assert(row[13] == 'OOM:')
            line['OOM'] = convert_number(row[14])
            assert(row[17] == 'dt:')
            line['dt'] = convert_number(row[18])
            rank = int(row[2])
            if not rank in data:
                data[rank] = [line]
            else:
                data[rank].append(line)
        elif row[0] == 'Node:':
            assert(row[4] == 'RAM:')
            line['RAM'] = convert_number(row[5])
            assert(row[6] == 'of')
            line['RAM_tot'] = convert_number(row[7])
            if len(row) > 9:
                assert(row[10] == 'GPU:')
                line['GPU'] = convert_number(row[11])
                assert(row[12] == 'of')
                line['GPU_tot'] = convert_number(row[13])
            data['node'].append(line)
        else:
            assert(row[0] == 'RSS:')
            line['RSS'] = convert_number(row[1])
            assert(row[4] == 'HWM:')
            line['HWM'] = convert_number(row[5])
            assert(row[8] == 'OOM:')
            line['OOM'] = convert_number(row[9])
            assert(row[12] == 'dt:')
            line['dt'] = convert_number(row[13])
            data['serial'].append(line)
    return data

def convert_number(value):
    if value.endswith('us'):
        return int(value[:-2])
    elif value.endswith('ms'):
        return int(value[:-2]) * 1e3
    elif value.endswith('s'):
        return int(value[:-1]) * 1e6
    elif value.endswith('K'):
        return int(value[:-1])
    elif value.endswith('M'):
        return int(value[:-1]) * 1024
    elif value.endswith('G'):
        return int(value[:-1]) * 1024**2
    return int(value)

def assert_node(line):
    assert(line.startswith('Node:'))

def assert_approx_equal(value1, value2):
    if abs(value2 - value1) / value1 > 0.05:
        print("Values differ by more than 5%:", value1, value2, abs(value2 - value1) / value1)
        assert(False)


mpi_ranks = int(sys.argv[1])
has_node_hwm = 'monitor' in sys.argv
has_gpu = 'gpu' in sys.argv
strict = 'strict' in sys.argv

data = read_data()

assert(len(data['serial']) == mpi_ranks)
assert(len(data[0]) == 3)
for rank in range(1,mpi_ranks):
    assert(rank in data)
    assert(len(data[rank]) == 2)

for d in data['serial']:
    assert(d['dt'] == 0)
for rank in range(mpi_ranks):
    for d in data[rank]:
        assert(d['dt'] >= 0)
        if strict:
            assert(d['dt'] > 0)

assert_approx_equal(data[0][1]['RSS'], data[0][1]['HWM'])
for rank in range(1,mpi_ranks):
    assert_approx_equal(data[rank][0]['RSS'], data[rank][0]['HWM'])

for rank in range(mpi_ranks):
    assert_approx_equal(data[rank][-2]['HWM'], data[rank][-1]['HWM'])

if strict:
    assert_approx_equal(data[0][0]['RSS'], data[0][0]['HWM'])
    assert_approx_equal(data[0][1]['RSS'], data[0][0]['RSS'] + 102400)
    assert_approx_equal(data[0][2]['RSS'], data[0][0]['RSS'])
    for rank in range(mpi_ranks):
        assert_approx_equal(data[rank][-1]['RSS'], data[rank][-2]['RSS'] - 102400)

if has_node_hwm:
    assert(len(data['node']) == 1)

    assert(data['node'][0]['RAM'] > 0)
    assert(data['node'][0]['RAM_tot'] > data['node'][0]['RAM'])

    if has_gpu:
        assert(data['node'][0]['GPU'] == 53 * 1024**2)
        assert(data['node'][0]['GPU_tot'] == 158 * 1024**2)