source: benchmark/rmit.py @ a77f25b

#!/usr/bin/python3
"""
Python Script to implement R.M.I.T. testing : Randomized Multiple Interleaved Trials

./rmit.py run COMMAND CANDIDATES
-t trials
-o option:values
"""


import argparse
import datetime
import itertools
import json
import os
import random
import re
import socket
import subprocess
import sys


def parse_range(x):
    result = []
    for part in x.split(','):
        if '-' in part:
            a, b = part.split('-')
            a, b = int(a), int(b)
            result.extend(range(a, b + 1))
        else:
            a = int(part)
            result.append(a)
    return result

class DependentOpt:
        def __init__(self, key, value):
                self.key = key
                self.value = value
                self.vars = re.findall("[a-zA-Z]", value)

def parse_option(key, values):
        try:
                num = int(values)
                return key, [num]
        except:
                pass

        if re.search("^[0-9-,]+$", values):
                values = parse_range(values)
                return key, [v for v in values]
        else:
                return key, DependentOpt(key, values)

def eval_one(fmt, vals):
        orig = fmt
        for k, v in vals:
                fmt = fmt.replace(k, str(v))

        if not re.search("^[0-9-/*+ ]+$", fmt):
                print('ERROR: pattern option {} (interpreted as {}) could not be evaluated'.format(orig, fmt), file=sys.stderr)
                sys.exit(1)

        return eval(fmt)

def eval_options(opts):
        dependents = [d for d in opts.values() if type(d) is DependentOpt]
        processed = []
        nopts = []
        for d in dependents:
                processed.append(d.key)
                lists = []
                for dvar in d.vars:
                        if not dvar in opts.keys():
                                print('ERROR: extra pattern option {}:{} uses unknown key {}'.format(d.key,d.value,dvar), file=sys.stderr)
                                sys.exit(1)

                        lists.append([(dvar, o) for o in opts[dvar]])
                        processed.append(dvar)

                kopt = []
                for vals in list(itertools.product(*lists)):
                        res = ['-{}'.format(d.key), "{}".format(eval_one(d.value, vals))]
                        for k, v in vals:
                                res.extend(['-{}'.format(k), "{}".format(v)])
                        kopt.append(res)
                nopts.append(kopt)


        for k, vals in opts.items():
                if k not in processed:
                        kopt = []
                        for v in vals:
                                kopt.append(['-{}'.format(k), "{}".format(v)])
                        nopts.append(kopt)

        return nopts

# returns the first option with key 'opt'
def search_option(action, opt):
        i = 0
        while i < len(action):
                if action[i] == opt:
                        i += 1
                        if i != len(action):
                                return action[i]
                i += 1

        return None

def actions_eta(actions):
        time = 0
        for a in actions:
                o = search_option(a, '-d')
                if o :
                        time += int(o)
        return time

taskset_maps = None

def init_taskset_maps():
        global taskset_maps
        known_hosts = {
                "jax": {
                        range(  1,  24) : "48-71",
                        range( 25,  48) : "48-71,144-167",
                        range( 49,  96) : "48-95,144-191",
                        range( 97, 144) : "24-95,120-191",
                        range(145, 192) : "0-95,96-191",
                },
        }

        if (host := socket.gethostname()) in known_hosts:
                taskset_maps = known_hosts[host]
                return True

        print("Warning unknown host '{}', disable taskset usage".format(host), file=sys.stderr)
        return False


def settaskset_one(action):
        o = search_option(action, '-p')
        if not o:
                return action
        try:
                oi = int(o)
        except ValueError:
                return action

        m = "Not found"
        for key in taskset_maps:
                if oi in key:
                        return ['taskset', '-c', taskset_maps[key], *action]

        print("Warning no mapping for {} cores".format(oi), file=sys.stderr)
        return action

def settaskset(actions):
        return [settaskset_one(a) for a in actions]

if __name__ == "__main__":
        # ================================================================================
        # parse command line arguments
        formats = ['raw', 'csv', 'json']
        parser = argparse.ArgumentParser(description='Python Script to implement R.M.I.T. testing : Randomized Multiple Interleaved Trials')
        parser.add_argument('--list', help='List all the commands that would be run', action='store_true')
        parser.add_argument('--file', nargs='?', type=argparse.FileType('w'), default=sys.stdout)
        parser.add_argument('--trials', help='Number of trials to run per combinaison', type=int, default=3)
        parser.add_argument('--notaskset', help='If specified, the trial will not use taskset to match the -p option', action='store_true')
        parser.add_argument('command', metavar='command', type=str, nargs=1, help='the command prefix to run')
        parser.add_argument('candidates', metavar='candidates', type=str, nargs='*', help='the candidate suffix to run')

        try:
                options, unknown =  parser.parse_known_args()

                options.option = []
                while unknown:
                        key = unknown.pop(0)
                        val = unknown.pop(0)

                        if key[0] != '-':
                                raise ValueError

                        options.option.append((key[1:], val))

        except:
                print('ERROR: invalid arguments', file=sys.stderr)
                parser.print_help(sys.stderr)
                sys.exit(1)

        # ================================================================================
        # Identify the commands to run
        command = './' + options.command[0]
        if options.candidates:
                commands = [command + "-" + c for c in options.candidates]
        else:
                commands = [command]
        for c in commands:
                if not os.path.isfile(c):
                        print('ERROR: invalid command {}, file does not exist'.format(c), file=sys.stderr)
                        sys.exit(1)

                if not os.access(c, os.X_OK):
                        print('ERROR: invalid command {}, file not executable'.format(c), file=sys.stderr)
                        sys.exit(1)


        # ================================================================================
        # Identify the options to run
        opts = dict([parse_option(k, v) for k, v in options.option])

        # Evaluate the options (options can depend on the value of other options)
        opts = eval_options(opts)

        # ================================================================================
        # Figure out all the combinations to run
        actions = []
        for p in itertools.product(range(options.trials), commands, *opts):
                act = [p[1]]
                for o in p[2:]:
                        act.extend(o)
                actions.append(act)

        # ================================================================================
        # Fixup the different commands

        # Add tasksets
        withtaskset = False
        if not options.notaskset and init_taskset_maps():
                withtaskset = True
                actions = settaskset(actions)

        # ================================================================================
        # Now that we know what to run, print it.
        # find expected time
        time = actions_eta(actions)
        print("Running {} trials{}".format(len(actions), "" if time == 0 else " (expecting to take {})".format(str(datetime.timedelta(seconds=int(time)))) ))

        # dry run if options ask for it
        if options.list:
                for a in actions:
                        print(" ".join(a))
                sys.exit(0)


        # ================================================================================
        # Prepare to run

        random.shuffle(actions)

        # ================================================================================
        # Run
        options.file.write("[")
        first = True
        for i, a in enumerate(actions):
                sa = " ".join(a[3:] if withtaskset else a)
                if first:
                        first = False
                else:
                        options.file.write(",")
                if options.file != sys.stdout:
                        print("{}/{} : {}          \r".format(i, len(actions), sa), end = '')
                fields = {}
                with subprocess.Popen( a, stdout  = subprocess.PIPE, stderr  = subprocess.PIPE) as proc:
                        out, err = proc.communicate()
                        if proc.returncode != 0:
                                print("ERROR: command '{}' encountered error, returned code {}".format(sa, proc.returncode), file=sys.stderr)
                                print(err.decode("utf-8"))
                                sys.exit(1)
                        for s in out.decode("utf-8").splitlines():
                                match = re.search("^(.*):(.*)$", s)
                                if match:
                                        try:
                                                fields[match.group(1).strip()] = float(match.group(2).strip().replace(',',''))
                                        except:
                                                pass

                options.file.write(json.dumps([a[3 if withtaskset else 0][2:], sa, fields]))
                options.file.flush()

        options.file.write("]\n")

        if options.file != sys.stdout:
                print("Done                                                                                ")