teems/App/launch

#!/usr/bin/env python3

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

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

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

# The TEEMS binary
TEEMS = "./teems"

def launch(node, manifest, config):
    manifestdata = manifest[node]
    cmdline = ''
    if 'launchprefix' in manifestdata:
        cmdline = manifestdata['launchprefix'] + ' '
    cmdline += TEEMS + " %s %s" % (manifestdata['sprvfile'], node)
    if 'args' in manifestdata:
        cmdline += ' ' + manifestdata['args']
    proc = subprocess.Popen(shlex.split(cmdline), stdin=subprocess.PIPE,
        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, bufsize=0)
    proc.stdin.write(config.encode('utf-8'))
    while True:
        line = proc.stdout.readline()
        if not line:
            break
        print(node + ": " + line.decode('utf-8'), end='')


if __name__ == "__main__":
    aparse = argparse.ArgumentParser(
        description='Launch TEEMS nodes'
    )
    aparse.add_argument('-m', default=MANIFEST,
        help='manifest.yaml file')
    aparse.add_argument('-p', default=PUBKEYS,
        help='pubkeys.yaml file')
    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)

    config = mkconfig.create_json(args.m, args.p, args.n)
    # 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"

    nodelist = args.n
    if nodelist is None or len(nodelist) == 0:
        nodelist = manifest.keys()

    threadlist = []
    for node in nodelist:
        if node == "params":
            continue
        thread = threading.Thread(target=launch,
            args=(node, manifest, config))
        thread.start()
        threadlist.append(thread)

    for thread in threadlist:
        thread.join()