data-do-biglittle

#!/usr/bin/env python3
'''
  data-do-biglittle GNU/Linux implementation
'''
# pylint: disable=logging-format-interpolation
from argparse import ArgumentParser
from copy import deepcopy
from doctest import testmod
from logging import basicConfig, getLogger, DEBUG, INFO
from hashlib import sha256
from sys import platform, stderr
from sys import exit as _exit
from platform import node
from string import ascii_lowercase, digits
import subprocess

try:
    from sys import implementation  # pylint: disable=ungrouped-imports
except ImportError:
    implementation = False
    from os import uname
try:
    if FileNotFoundError:  # pylint: disable=using-constant-test
        pass
except NameError:
    FileNotFoundError = IOError  # pylint: disable=redefined-builtin
_LOG_FMT = '%(asctime)s: %(message)s'
_DATE_FMT = '%Y-%m-%d-%H:%M:%s'
_CPUINFO = '/proc/cpuinfo'
_SIBLINGS_FMT = '/sys/devices/system/cpu/cpu{}/topology/thread_siblings_list'
_MAX_FREQ_FMT = '/sys/devices/system/cpu/cpu{}/cpufreq/cpuinfo_max_freq'
_MODEL_KEYS = [
    'CPU variant',
    'model name',
    'isa',
    'CPU part',
    'uarch',
    'CPU implementer',
    'CPU architecture',
    'CPU part',
    'CPU revision',
]


def shorthexhash(item):
    '''
    *shorthexhash* returns eight bytes of a sha256 hash for *item*.

    >>> shorthexhash('supercop'.encode())
    '12975e47'
    '''
    return sha256(item).hexdigest()[:8]


def _machine2threadslist(_file=_CPUINFO):
    '''
    *_machine2threadslist* counts CPU threads; GNU/Linux shows them as
    individual processor cores but this is a noble lie.

    This returns a list of operating system 'processor' numbers which is a
    number of a running thread that much later needs to be translated to a real
    core.

    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> thread_list = _machine2threadslist(_file=info)
    >>> thread_list
    ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11']
    '''
    log = getLogger()
    with open(_file, 'r') as o_f:
        log.debug('Opened %s', _file)
        cpus = []
        for line in o_f:
            pair = line.split(':')
            if pair[0].strip() == 'processor':
                p_id = pair[1].strip()
                cpus.append(p_id)
        return cpus


def _threadslist2coreslist(cpu_ids, _file=_SIBLINGS_FMT):
    '''
    *_threadslist2coreslist* takes a list of *cpu_ids* and returns a set of
    *cores* which removes duplicate (eg: hyperthreads) items from the *cpu_ids*
    list.

    >>> _SIBLINGS_FMT = 'samples/topology/cpu{}'
    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> thread_list = _machine2threadslist(_file=info)
    >>> thread_list
    ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', '10', '11']
    >>> _threadslist2coreslist(thread_list, _file=_SIBLINGS_FMT)
    ['0', '1', '2', '3', '4', '5']
    '''
    candidates = []
    unique_cores = []
    for cpu in cpu_ids:
        _f = _file.format(cpu)
        with open(_f, 'r') as o_f:
            for line in o_f:
                if not line.strip() in candidates:
                    candidates.append(line.strip())
                    unique_cores.append(cpu)
    return unique_cores


def _extract_max_freq(core_id, _fmt=_MAX_FREQ_FMT):
    '''
    *_extract_max_freq* extracts the maximum core frequency for core *core_id*.

    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> _extract_max_freq(0, _fmt=_MAX_FREQ_FMT)
    '4500000'
    '''
    log = getLogger()
    _file = _fmt.format(core_id)
    try:
        with open(_file, 'r') as o_f:
            return o_f.readline().strip()
    except FileNotFoundError:
        log.debug(
            'Unable to extract maximum CPU frequency for core: %s', core_id
        )
        return 0


def _machine2coremodels(_f=_CPUINFO):
    '''
    *_machine2coremodels* tells us the unique types of cores.

    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> core_models = _machine2coremodels(_f=info)
    >>> len(core_models.keys())
    12
    '''
    modelkeys = _MODEL_KEYS
    with open(_f, 'r') as o_f:
        cores = {}
        seen_cores = []
        for line in o_f:
            pair = line.split(':')
            if len(pair) < 2:
                continue
            if pair[0].strip() == 'processor':
                p_id = pair[1].strip()
                cores[p_id] = {'processor': p_id}
                seen_cores.append(p_id)
            else:
                key = pair[0].strip()
                if key in modelkeys:
                    for core in seen_cores:
                        if key not in cores[core]:
                            cores[core][key] = pair[1].strip()
        return cores


def _cores2detailedcores(cores, _fmt=_MAX_FREQ_FMT):
    '''
    *_cores2detailedcores* extracts frequency and midr information for each
    *core* in *cores* and returns a modified *cores* that contains additional
    metadata.

    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> core_models = _machine2coremodels(_f=info)
    >>> cores = _cores2detailedcores(core_models, _fmt=_MAX_FREQ_FMT)
    >>> for core in cores:
    ...  assert cores[core]['cpuinfo_max_freq'] == '4500000'
    ...
    >>> len(cores)
    12
    '''
    enhanced_cores = deepcopy(cores)
    for core in cores:
        enhanced_cores[core]['cpuinfo_max_freq'] = _extract_max_freq(
            core, _fmt=_fmt
        )
        enhanced_cores[core]['midr'] = _extract_midr(cores[core])
    return enhanced_cores


def _extract_midr(core):
    '''
    *_extract_midr* parses the core data structure and calculates the MIDR for
    the core in hex. For non-arm systems, it returns 00000000.

    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> info = 'samples/proc_cpuinfo_armv8_64_big_little_bigger'
    >>> core_models = _machine2coremodels(_f=info)
    >>> cores = _cores2detailedcores(core_models, _fmt=_MAX_FREQ_FMT)
    >>> for core in cores:
    ...   x = _extract_midr(cores[core])
    ...   assert x is not None
    ...
    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> core_models = _machine2coremodels(_f=info)
    >>> cores = _cores2detailedcores(core_models, _fmt=_MAX_FREQ_FMT)
    >>> for core in cores:
    ...   assert _extract_midr(cores[core]) == '00000000'
    ...
    '''
    log = getLogger()
    midr = 0
    keys = [
        ('CPU implementer', 24),
        ('CPU variant', 20),
        ('CPU architecture', 16),
        ('CPU part', 4),
        ('CPU revision', 0),
    ]
    for key, startbit in keys:
        midrpart = 0
        if key in core:
            if key == 'CPU architecture':
                assert core[key] in ('7', '8', 'AArch64')
                midrpart = 15
            else:
                midrpart = int(core[key], 16)
        midr += midrpart << startbit
    midr = '%08x' % midr
    log.debug('core: {} midr: 0x{}'.format(core['processor'], midr))
    return midr


def _get_arch():
    '''
    return arm, armv7, mipsel, mips64el, aarch64, i686, x86_64, etc.

    >>> arches = ['arm', 'armv7', 'mipsel', 'mips64el', 'aarch64',
    ...          'i686', 'x86_64', 'riscv32', 'riscv64']
    >>> _get_arch() in arches
    True
    '''
    # pylint: disable=protected-access
    if implementation:
        arch = implementation._multiarch.split('-')[0]
    else:
        arch = uname()[-1]
    return arch


def _make_hostname(args, core):
    '''
    *_make_hostname* takes a core dict and returns a unique string for that
    core which is architecture specific.
    With one cpu type: hostname
    With several cpu types:
    arm:
        hostname fmt = hostname,midr,cpuinfo_max_freq
    x86_64/amd64:
        hostname fmt = hostname,h(model name)[5],cpuinfo_max_freq
    riscv/mips/mips64:
        hostname fmt = hostname,h(isa)[5],cpuinfo_max_freq
    ppc:
        hostname fmt = hostname,h(revision),cpuinfo_max_freq

    >>> args = parse_args()
    >>> args.hostname = 'host'
    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> thread_ids = _machine2threadslist(_file=info)
    >>> core_models = _machine2coremodels(_f=info)
    >>> cores = _cores2detailedcores(core_models, _fmt=_MAX_FREQ_FMT)
    >>> for core in cores:
    ...  x = _make_hostname(args, cores[core])
    ...
    >>> print(x)
    host,ff0f5016,4500000
    '''
    fmt = '{},{},{}'
    arch = _get_arch().lower()
    basename = ""
    if arch in ('aarch64') or arch.startswith('arm'):
        basename = _extract_midr(core)
    if arch == 'x86_64':
        basename = shorthexhash(core['model name'].encode())
    if arch == 'riscv':
        basename = shorthexhash(core['isa'].encode())
    if arch == 'ppc':
        basename = shorthexhash(core['revision'].encode())
    if arch in ('mips', 'mips64'):
        basename = shorthexhash(core['isa'].encode())
    if not basename:
        basename = 'unknown-arch-' + arch
    # This should be more careful as some systems will not have frequency
    # information available and it will be set to 0 which is not very
    # informative.
    hostname = fmt.format(args.hostname, basename, core['cpuinfo_max_freq'])
    return _clean_hostname(hostname)


def _siblings_cores2init_run(args, siblings, cores):
    '''
    *_siblings_cores2init_run* currently split by CPU variations except
    processor id which is expected to vary from core to core as it's a counter.

    >>> args = parse_args()
    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> _SIBLINGS_FMT = 'samples/topology/cpu{}'
    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> thread_ids = _machine2threadslist(_file=info)
    >>> siblings = _threadslist2coreslist(thread_ids, _file=_SIBLINGS_FMT)
    >>> cores = _cores2detailedcores(_machine2coremodels(_f=info), _fmt=_MAX_FREQ_FMT)
    >>> cores_data_init, cores_data_run = _siblings_cores2init_run(
    ...    args, siblings, cores
    ... )
    ...
    '''
    data_init, data_run = [], []
    cpus = set()
    cores = deepcopy(cores)
    for core in siblings:
        variant = tuple(
            (key, value)
            for key, value in sorted(cores[core].items())
            if key != 'processor'
        )
        if variant not in cpus:
            cpus.add(variant)
            data_init.append(cores[core])
        data_run.append(cores[core])
    data_init = _cores2cores_w_shorthostname(args, data_init)
    data_run = _cores2cores_w_shorthostname(args, data_run)
    return data_init, data_run


def _clean_hostname(hostname):
    '''
    *_clean_hostname* takes a *hostname*, copies it, and returns a mutated
    copy.  This function returns the left most label of a domain name as the
    hostname, and it ensures that the returned hostname is only contains
    characters present in string.ascii_lowercase, string.digits, and ','.

    >>> hostname = '$s%u&p*(e)/?rc0p;'
    >>> new_hostname = _clean_hostname(hostname)
    >>> new_hostname
    'superc0p'
    '''
    allowed = ascii_lowercase + digits + ','
    new_hostname = hostname.lower().split('.')[0]
    return ''.join(c for c in new_hostname if c in allowed)


def _cores2cores_w_shorthostname(args, cores):
    '''
    *_cores2cores_w_shorthostname* takes *args* and *cores* and returns an
    augmented copy of *cores* which has a new shorthostname key for each core.

    >>> _SIBLINGS_FMT = 'samples/topology/cpu{}'
    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> args = parse_args()
    >>> thread_ids = _machine2threadslist(info)
    >>> siblings = _threadslist2coreslist(thread_ids, _file=_SIBLINGS_FMT)
    >>> core_models = _machine2coremodels(_f=info)
    >>> cores = _cores2detailedcores(core_models, _fmt=_MAX_FREQ_FMT)
    >>> data_init, data_run = [], []
    >>> cpus = set()
    >>> for core in siblings:
    ...    variant = tuple(
    ...        (key, value)
    ...        for key, value in sorted(cores[core].items())
    ...        if key != 'processor'
    ...    )
    ...    if variant not in cpus:
    ...         cpus.add(variant)
    ...         data_init.append(cores[core])
    ...    data_run.append(cores[core])
    >>> data_init = _cores2cores_w_shorthostname(args, data_init)
    >>> data_run = _cores2cores_w_shorthostname(args, data_run)
    '''
    summary = set()
    new_cores = deepcopy(cores)
    for core in new_cores:
        summary.add(_make_hostname(args, core))
    if len(summary) == 1:
        shorthostname = summary.pop().split(',')[0]
        for core in new_cores:
            core['shorthostname'] = shorthostname
    else:
        for core in new_cores:
            core['shorthostname'] = _make_hostname(args, core)
    return new_cores


def data_do(cores, cmd, args):
    '''
    *data_do* runs a *cmd*:
    - fork/exec data-init
    - waitpid for data-init to finish

    >>> args = parse_args()
    >>> args.dry_run = True
    >>> args.verbose = False
    >>> args.hostname = 'host'
    >>> info = 'samples/proc_cpuinfo_x86_64'
    >>> _SIBLINGS_FMT = 'samples/topology/cpu{}'
    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> thread_ids = _machine2threadslist(info)
    >>> siblings = _threadslist2coreslist(thread_ids, _file=_SIBLINGS_FMT)
    >>> core_models = _machine2coremodels(_f=info)
    >>> cores = _cores2detailedcores(core_models, _fmt=_MAX_FREQ_FMT)
    >>> cores_data_init, cores_data_run = _siblings_cores2init_run(
    ...    args, siblings, cores
    ... )
    >>> data_do(cores_data_init, args.init_cmd, args)
    >>> data_do(cores_data_run, args.run_cmd, args)
    '''
    log = getLogger()
    log.debug('{} starting'.format(cmd))
    processes = []
    for core in cores:
        core_id = core['processor']
        shorthostname = core['shorthostname']
        task = [
            'env shorthostname='
            + shorthostname
            + ' taskset -c '
            + str(core_id)
            + ' '
            + cmd
        ]
        if args.dry_run:
            task = 'echo ' + task[0]
        log.debug('Running task: {}'.format(task))
        executed = subprocess.Popen(task, shell=True)
        log.debug('Started pid: {}'.format(executed.pid))
        processes.append(executed)
    log.debug('Waiting on {} task(s)'.format(len(processes)))
    for task in processes:
        task.wait()
        log.debug('Pid {} finished'.format(task.pid))
    log.debug("{} finished".format(cmd))


def parse_args():
    '''
    *parse_args* parses arguments and returns an argument structure.

    >>> parse_args() # doctest:+ELLIPSIS
    Namespace(cpu_info='/proc/cpuinfo'...
    '''
    parser = ArgumentParser(prog='data-do-biglittle')
    parser.add_argument(
        '-v', '--verbose', action='count', help='increase log level'
    )
    parser.add_argument('-d', '--dry-run', action='store_true', help='dry run')
    parser.add_argument(
        '-t', '--doc-tests', action='store_true', help='run doctests'
    )
    parser.add_argument(
        '-i',
        '--init-cmd',
        action='store',
        default='./data-init',
        help='data-init command',
    )
    parser.add_argument(
        '-r',
        '--run-cmd',
        action='store',
        default='./data-run',
        help='data-run command',
    )
    parser.add_argument(
        '-s',
        '--summarize-cmd',
        action='store',
        default='./data-summarize',
        help='data-summarize command',
    )
    parser.add_argument(
        '-H', '--hostname', action='store', default=node().split('.')[0], help='set hostname'
    )
    parser.add_argument(
        '-f',
        '--cpu-info',
        action='store',
        default=_CPUINFO,
        help='set file to read (default: ' + _CPUINFO + ')',
    )
    return parser.parse_args()


def main(
    args=None,
    _CPUINFO=_CPUINFO,
    _SIBLINGS_FMT=_SIBLINGS_FMT,
    _MAX_FREQ_FMT=_MAX_FREQ_FMT,
):
    '''
    *main* runs the show
    >>> _CPUINFO = 'samples/proc_cpuinfo_x86_64'
    >>> _SIBLINGS_FMT = 'samples/topology/cpu{}'
    >>> _MAX_FREQ_FMT = 'samples/cpuinfo_max_freq/cpu{}'
    >>> args = parse_args()
    >>> args.dry_run = True
    >>> args.doc_tests = False
    >>> args.verbose = False
    >>> args.cpu_info = _CPUINFO
    >>> args.hostname = 'host'
    >>> main(args=args, _CPUINFO=_CPUINFO, _SIBLINGS_FMT=_SIBLINGS_FMT, _MAX_FREQ_FMT=_MAX_FREQ_FMT)
    '''
    if not args:
        args = parse_args()
    if args.verbose:
        basicConfig(
            level=DEBUG, stream=stderr, datefmt=_DATE_FMT, format=_LOG_FMT
        )
    else:
        basicConfig(
            level=INFO, stream=stderr, datefmt=_DATE_FMT, format=_LOG_FMT
        )
    log = getLogger()

    if args.doc_tests:
        _exit(testmod())

    log.debug('Hostname: {}'.format(args.hostname))
    log.debug('OS: {}'.format(platform))
    log.debug('Arch: {}'.format(_get_arch()))
    thread_ids = _machine2threadslist(_file=args.cpu_info)
    log.debug('CPU threads: {}'.format(len(thread_ids)))
    siblings = _threadslist2coreslist(thread_ids, _file=_SIBLINGS_FMT)
    log.debug('CPU cores: {}'.format(len(siblings)))
    log.debug('CPU unique core ids: {}'.format(siblings))
    cores = _cores2detailedcores(
        _machine2coremodels(_f=_CPUINFO), _fmt=_MAX_FREQ_FMT
    )
    cores_data_init, cores_data_run = _siblings_cores2init_run(
        args, siblings, cores
    )
    cores_data_summarize = cores_data_init
    log.debug("data init tasks: {}".format(len(cores_data_init)))
    log.debug("data run tasks: {}".format(len(cores_data_run)))
    log.debug("data summarize tasks: {}".format(len(cores_data_summarize)))
    data_do(cores_data_init, args.init_cmd, args)
    data_do(cores_data_run, args.run_cmd, args)
    data_do(cores_data_summarize, args.summarize_cmd, args)


if __name__ == '__main__':
    main()