teems/Client/clientlaunch

#!/usr/bin/env python3

import argparse
import os
import shlex
import subprocess
import sys
import threading
import yaml
sys.path.insert(0, os.getcwd())
sys.path.insert(1, '../App/')
import mkconfig

# The default manifest file
MANIFEST = "../App/manifest.yaml"

# The default pubkeys file
PUBKEYS = "../App/pubkeys.yaml"

# The client binary
CLIENTS = "./clients"

def launch(config, cmd, threads, lgfile, corelist):
    cmdline = ''
    if corelist is not None:
        cmdline = "numactl -C %s " % corelist
    cmdline += CLIENTS + " -t " + str(threads) + ""

    stdout_file = subprocess.PIPE
    if lgfile:
        stdout_file = open(lgfile, "a+")

    proc = subprocess.Popen(shlex.split(cmdline) + cmd,
        stdin=subprocess.PIPE, stdout=stdout_file,
        stderr=subprocess.STDOUT, bufsize=0)
    print(cmdline)
    proc.stdin.write(config.encode('utf-8'))

    if lgfile:
        proc.wait()
        stdout_file.close()
    else:
        while True:
            line = proc.stdout.readline()
            if not line:
                break
            print(line.decode('utf-8'), end='', flush=True)


if __name__ == "__main__":
    print("In clientlaunch")
    aparse = argparse.ArgumentParser(
        description='Launch CLIENTS'
    )
    aparse.add_argument('-m', default=MANIFEST,
        help='manifest.yaml file')
    aparse.add_argument('-p', default=PUBKEYS,
        help='pubkeys.yaml file')
    aparse.add_argument('-t', default=1,
        help='number of threads')
    aparse.add_argument('-T', default=None,
        help='CPU cores to use for clients')
    aparse.add_argument('-z', default=None,
        help='override message size')
    aparse.add_argument('-u', default=None,
        help='override max number of users')
    aparse.add_argument('-B', default=None,
        help='override max number of outgoing token channel messages per user per epoch')
    aparse.add_argument('-b', default=None,
        help='override max number of incoming token channel messages per user per epoch')
    aparse.add_argument('-C', default=None,
        help='override max number of outgoing ID channel messages per user per epoch')
    aparse.add_argument('-c', default=None,
        help='override max number of incoming ID channel messages per user per epoch')
    aparse.add_argument('-l', default=None,
        help='log file to store client simulator log for an experiment')
    aparse.add_argument('-n', nargs='*', help='nodes to include')
    aparse.add_argument('cmd', nargs='*', help='experiment to run')
    args = aparse.parse_args()

    with open(args.m) as mf:
        manifest = yaml.safe_load(mf)

    params_overrides = {
        'msg_size': args.z,
        'user_count': args.u,
        'token_out': args.B,
        'token_in': args.b,
        'id_out': args.C,
        'id_in': args.c,
    }

    config = mkconfig.create_json(args.m, args.p, args.n, params_overrides)

    # There must not be any newlines in the config json string
    if "\n" in config:
        print("Error: config.json must not contain embedded newlines")
        sys.exit(1)
    # Now add a trailing newline
    config += "\n"

    launch(config, args.cmd, args.t, args.l, args.T)