Distributed wrk

wrkd.py

#!/usr/bin/env python3
import argparse
import subprocess
import sys

from collections import defaultdict
from statistics import mean


HELP = """Usage: wrkd <options> <url>
  Options:
    -c, --connections <N>  Connections to keep open
    -d, --duration    <T>  Duration of test
    -t, --threads     <N>  Number of threads to use

    -s, --script      <S>  Load Lua script file
    -H, --header      <H>  Add header to request
        --latency          Print latency statistics
        --timeout     <T>  Socket/request timeout
        --hosts            List of hosts
    -v, --version          Print version details

  Numeric arguments may include a SI unit (1k, 1M, 1G)
  Time arguments may include a time unit (2s, 2m, 2h)"""


def parse_args():
    parser = argparse.ArgumentParser(description='Distributed wrk wrapper', usage=HELP, add_help=False)
    parser.add_argument('-h', '--help', action='store_true', default=False)
    parser.add_argument('--hosts', action='store', type=str, nargs='*', required=True, help='host list')
    parser.add_argument('-c', '--connections', type=int, default=100)
    args, wrk_args = parser.parse_known_args()
    if args.help:
        print(HELP)
        return sys.exit(0)
    if not wrk_args:
        print(HELP)
        sys.exit(1)
    return args.hosts, args.connections, wrk_args


def aggregate_results(stdout):
    by_host = defaultdict(list)
    for l in stdout.splitlines():
        host, output = l.split(':', 1)
        by_host[host].append(output)

    results = {}
    for host, outputs in by_host.items():
        results[host] = {'Latency': {}, 'Req/Sec': {}}
        for line in outputs:
            if 'Latency' in line:
                results[host]['Latency']['Avg'], \
                    results[host]['Latency']['Stdev'], \
                    results[host]['Latency']['Max'], \
                    results[host]['Latency']['+/- Stdev'] = line.split()[1:]
            elif ' requests in ' in line:
                req, _, _, _, read, _ = line.split()
                results[host]['requests'], results[host]['read'] = req, read
            elif 'Requests/sec:' in line:
                results[host]['Requests/sec'] = float(line.split()[-1])
        results[host]['stdout'] = '\n'.join(outputs)
    return results


def print_results(results):
    for host, res in results.items():
        print(f'{host}{"-" * 20}')
        print(res['stdout'])

    rate = mean([v['Requests/sec'] for v in results.values()])
    ## TODO: Finish output format
    print(f'Aggregate {"-" * 20}')
    print(f'Requests/sec: {rate}')


def main(hosts, connections, wrk_args):
    wrk_cmd = ' '.join(['wrk', '-c', str(int(connections / len(hosts))), *wrk_args])
    print(f'Launching `{wrk_cmd}` on {", ".join(hosts)} with {int(connections/len(hosts))} connections per host')
    pdsh_cmd = ['pdsh', '-b', '-w', ','.join(hosts)]
    cmd = [*pdsh_cmd, wrk_cmd]
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stdin=subprocess.PIPE)
    results = aggregate_results(proc.stdout.decode())
    print_results(results)


if __name__ == '__main__':
    hosts, connections, wrk_args = parse_args()
    main(hosts, connections, wrk_args)