hpxrun-jetlag.py

#!/usr/bin/python3
#
# Copyright (c) 2014 Thomas Heller
#
# SPDX-License-Identifier: BSL-1.0
# Distributed under the Boost Software License, Version 1.0. (See accompanying
# file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

# This script is a simple startup script to start 1 or multiple HPX localities.
# It supports various startup wrappers for distributed runs.
#
# Usage:
#  hpxrun.py hpx-application [Options] [Additional options]
#
# Available options are:
#   -l Number of localities to start
#   -t Number of threads per locality
#   -p Parcelport to use
#   -r Which runwrapper to use
#   -e Expected return codes of all invoked processes
#   -v verbose output
#   -- delimiter for additional arguments

import sys, os, string, atexit
import functools, signal
import traceback
import time
import threading
import re

try:
    from paramiko import SSHClient
    has_paramiko = True
except:
    has_paramiko = False


def esc(v):
    nv = ''
    nvs = 0
    modified = False
    for g in re.finditer(r'[^\w/:\.-]', v):
        modified = True
        nv += v[nvs:g.start()]
        nv += '%' + str(ord(g.group(0)))
        nvs = g.end()
    nv += v[nvs:]
    return (nv, modified)

# This class allows Paramiko to behave
# like a subprocess.Popen() proc.
class ParamikoProc:

    def __init__(self, sout, serr):
        self.sout = sout
        self.serr = serr

    def wait(self):
        self.returncode = self.sout.channel.recv_exit_status()
        outs = self.sout.read().decode().strip()
        if outs != '':
            print(outs)
        errs = self.serr.read().decode().strip()
        if errs != '':
            print(errs)

from optparse import OptionParser

import subprocess
from subprocess import Popen

# Our global list of processes we started
procs = []
def subproc(cmd):
    kwargs = {}
    if sys.platform == 'win32':
        # For some reason or another ... this seems to be non existent sometimes
        try:
            kwargs['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP
        except:
            pass

        # If this is to be an ssh command and we are on windows
        # use the Paramiko (http://www.paramiko.org) instead.
        if cmd[0] == "ssh":
            ssh_subproc(cmd)
            return

    # Support for running from inside a singularity container
    # on Linux systems
    if cmd[0] == "ssh" and "SINGULARITY_CONTAINER" in os.environ:
        if "SING_OPTS" in os.environ:
            sing = ["singularity","exec"] + re.split(r'\s+',os.environ["SING_OPTS"].strip()) + [os.environ["SINGULARITY_CONTAINER"]]
        else:
            sing = ["singularity","exec", os.environ["SINGULARITY_CONTAINER"]]
        cmd = cmd[0:2] + sing + cmd[2:]

    # Pass environment variables along in Linux
    # Also save the current directory PWD. This
    # assumes a shared file system.
    if cmd[0] == "ssh":

        # Locate the env.py command and add it to the path
        xdir = os.path.abspath(os.path.join(__file__, os.pardir))
        ecmd = ["python3",os.path.join(xdir, "env.py")]

        # Encode all environment variables
        for e in options.environ.split(',')+['PWD']:
            if re.match(r'^\w+$', e):
                value = os.environ.get(e, "")
                value, modified = esc(value)
                if modified:
                    # The value contains special
                    # characters. Escaping them with
                    # quotes turns out to be kind of
                    # difficult.
                    ecmd += ["%s%%=%s" % (e,value)]
                else:
                    ecmd += ["%s=%s" % (e, value)]

        cmd = cmd[0:2] + ecmd + cmd[2:]


    proc = Popen(cmd, **kwargs)
    procs.append(proc)

    handler = functools.partial(cleanup, proc)
    signal.signal(signal.SIGINT, handler)
    if sys.platform == 'win32':
        signal.signal(signal.SIGBREAK, handler)
    else:
        signal.signal(signal.SIGTERM, handler)

def ssh_subproc(cmd):
    global SSHClient, procs
    if not has_paramiko:

        # Attempt to install Paramiko automatically...
        try:
            from pip import main as pip_main
            has_pip = True
        except:
            has_pip = False
        if not has_pip:
            try:
                from pip._internal import main as pip_main
                has_pip = True
            except:
                pass

        if has_pip:
            print("You need the Paramiko SSH client, installing...")
            pip_main(["install","--user","paramiko"])
            print("The Paramiko SSH client has been installed. Please re-run your command.")
        else:
            print("Please install 'pip install --user paramiko' in order to run this command.")
            raise Exception()

        exit(1)

    cmd_str = ''
    for c in cmd[2:]:
        if cmd_str != '':
            cmd_str += ' '
        if re.search(r'\s', c):
            if not re.search(r"'", c):
                cmd_str += "'" + c + "'"
            elif not re.search(r'"', c):
                cmd_str += '"' + c + '"'
            else:
                cmd_str += '"' + re.sub(r'"','\\"',c) + '"'
        else:
            cmd_str += c

    client = SSHClient()
    client.load_system_host_keys()
    client.connect(cmd[1])
    sin, sout, serr = client.exec_command(cmd_str)
    procs += [ParamikoProc(sout, serr)]

# Run with no run wrapper
# This is just starting "localities" processes on the local node
def run_none(cmd, localities, nodes, verbose):
    if nodes is None:
        if localities > 1:
            for locality in range(localities):
                exec_cmd = cmd + ['--hpx:node=' + str(locality)]
                if verbose:
                    print('Executing command: ' + ' '.join(exec_cmd))
                subproc(exec_cmd)
        else:
             if verbose:
                 print('Executing command: ' + ' '.join(cmd))
             subproc(cmd)
    else:
        assert len(nodes) == localities, "Number of localities must match number of nodes=%d." % len(nodes)
        if localities > 1:
            for locality in range(localities):
                node = nodes[locality]
                node0 = nodes[0]
                g = re.match(r'([\d.]+):(\d+)$', node)
                ip = g.group(1)
                exec_cmd = ['ssh', ip] + cmd + ['--hpx:hpx=' + node, '--hpx:agas=' + node0]
                if locality != 0:
                    exec_cmd += ["--hpx:worker"]
                if verbose:
                    print('Executing command: ' + ' '.join(exec_cmd))
                subproc(exec_cmd)
        else:
             if verbose:
                 print('Executing command: ' + ' '.join(cmd))
             subproc(cmd)

# Run with mpiexec
# This is executing mpiexec with the "-np" option set to the number of localities
def run_mpi(cmd, localities, verbose):
    mpiexec = '/usr/lib64/mpich/bin/mpiexec'
    if mpiexec == '':
        mpiexec = '/usr/lib64/mpich/bin/mpiexec'
    if mpiexec == '':
        msg = 'mpiexec not available on this platform. '
        msg += 'Please rerun CMake with HPX_PARCELPORT_MPI=True.'
        print(msg, sys.stderr)
        sys.exit(1)
    exec_cmd = [mpiexec, '-n', str(localities)] + cmd
    if verbose:
        print('Executing command: ' + ' '.join(exec_cmd))
    subproc(exec_cmd)

# Run with srun
# This is executing srun with the '-n' option set to the number of localities
def run_srun(cmd, localities, verbose):
    exec_cmd = ['srun', '-K', '-n', str(localities)] + cmd
    if verbose:
        print('Executing command: ' + ' '.join(exec_cmd))
    subproc(exec_cmd)

# Run with jsrun
# This is executing jsrun with the '-n' option set to the number of localities
def run_jsrun(cmd, localities, verbose):
    exec_cmd = ['jsrun', '-n', str(localities)] + cmd
    if verbose:
        print('Executing command: ' + ' '.join(exec_cmd))
    subproc(exec_cmd)

# Select the appropriate run function based on runwrapper
def run(cmd, runwrapper, localities, nodes, verbose):
    if runwrapper == 'none':
        run_none(cmd, localities, nodes, verbose)
    if runwrapper == 'mpi':
        assert nodes is None, "nodes option only valid with tcp parcelport."
        run_mpi(cmd, localities, verbose)
    if runwrapper == 'srun':
        assert nodes is None, "nodes option only valid with tcp parcelport."
        run_srun(cmd, localities, verbose)
    if runwrapper == 'jsrun':
        assert nodes is None, "nodes option only valid with tcp parcelport."
        run_jsrun(cmd, localities, verbose)

# Building the command line. This function concatenates the different options
def build_cmd(options, args):
    cmd = [args[0]]
    args.pop(0)

    # Append the remaining args
    for arg in args:
        cmd += [arg]

    if options.localities > 1:
        # Selecting the parcelport for hpx via hpx ini configuration
        select_parcelport = (lambda pp:
            ['--hpx:ini=hpx.parcel.verbs.priority=1000', '--hpx:ini=hpx.parcel.verbs.enable=1'] if pp == 'verbs'
            else ['--hpx:ini=hpx.parcel.ipc.priority=1000', '--hpx:ini=hpx.parcel.ipc.enable=1'] if pp == 'ipc'
            else ['--hpx:ini=hpx.parcel.mpi.priority=1000', '--hpx:ini=hpx.parcel.mpi.enable=1', '--hpx:ini=hpx.parcel.bootstrap=mpi'] if pp == 'mpi'
            else ['--hpx:ini=hpx.parcel.tcp.priority=1000', '--hpx:ini=hpx.parcel.tcp.enable=1'] if pp == 'tcp'
            else [])
        cmd += select_parcelport(options.parcelport)

    # set number of threads
    if options.threads == -1:
        cmd += ['--hpx:threads=all']
    if options.threads == -2:
        cmd += ['--hpx:threads=cores']
    if options.threads >= 1:
        cmd += ['--hpx:threads=' + str(options.threads)]

    # set number of localities
    if options.localities > 1 and (options.parcelport != "none"):
        cmd += ['--hpx:localities=' + str(options.localities)]

    return cmd

def check_options(parser, options, args):
    if 0 == len(args):
        print('Error: You need to specify at least the application to start\n', sys.stderr)
        parser.print_help()
        sys.exit(1)

    if not os.path.exists(args[0]):
        print('Executable ' + args[0] + ' does not exist', sys.stderr)
        sys.exit(1)

    if options.localities < 1:
        print('Can not start less than one locality', sys.stderr)
        sys.exit(1)

    if options.threads < 1 and options.threads != -1 and options.threads != -2:
        print('Can not start less than one thread per locality', sys.stderr)
        sys.exit(1)

    check_valid_parcelport = (lambda x:
            x == 'verbs' or x == 'ipc' or x == 'mpi' or x == 'tcp' or x == 'none');
    if not check_valid_parcelport(options.parcelport):
        print('Error: Parcelport option not valid\n', sys.stderr)
        parser.print_help()
        sys.exit(1)

    check_valid_runwrapper = (lambda x:
            x == 'none' or x == 'mpi' or x == 'srun' or x =='jsrun');
    if not check_valid_runwrapper(options.runwrapper):
        print('Error: Runwrapper option not valid\n', sys.stderr)
        parser.print_help()
        sys.exit(1)

# Send a SIGTERM/SIGBRAK to proc and wait for it to terminate.
def term(proc):
    if sys.platform == 'win32':
        try:
            proc.send_signal(signal.CTRL_BREAK_EVENT)
        except:
            proc.terminate()
    else:
        proc.terminate()
    proc.wait()

# Stop the sub-process child if signum is SIGTERM. Then terminate.
def cleanup(child, signum, frame):
    try:
        if child and ((sys.platform == 'win32') or signum != signal.SIGINT):
            # Forward SIGTERM on Linux or any signal on Windows
            term(child)
    except:
        traceback.print_exc()
    finally:
        sys.exit()

if __name__ == '__main__':

    help_message = 'Usage %proc hpx-application [Options] [-- Additional options]\n'
    help_message = help_message + '\n'
    help_message = help_message + 'This script is a simple startup script to start '
    help_message = help_message + 'one or multiple HPX localities. It supports '
    help_message = help_message + 'various startup wrappers for distributed runs.'

    parser = OptionParser(usage = help_message)

    default_env = (lambda env, default:
        os.environ[env] if env in os.environ else default)

    parser.add_option('-n', '--nodes'
      , action='store', type='str'
      , dest='nodes', default=default_env('HPXRUN_NODES', '')
      , help= \
'''Comma delimited list of ip addresses (possibly including port) to run on.
E.g.  167.96.129.220,167.96.131.223  or 167.96.129.220:7910,167.96.131.223:7910. 
If the port is not specified, port 7910 will be assumed.
Used by the tcp parcelport only.
      ''')

    parser.add_option('-l', '--localities'
      , action='store', type='int'
      , dest='localities', default=default_env('HPXRUN_LOCALITIES', '1')
      , help='Number of localities to run (environment variable '
              'HPXRUN_LOCALITIES')

    parser.add_option('-t', '--threads'
      , action='store', type='int'
      , dest='threads', default=default_env('HPXRUN_THREADS', '1')
      , help='Number of threads per locality (environment variable '
             'HPXRUN_THREADS)')

    # It is helpful to allow the users to
    # set a defaultport. Re-using 7910 as
    # the default may lead to hangs.
    parser.add_option('-d', '--defaultport'
      , action='store', type='int'
      , dest='defaultport', default=7910
      , help='The default port to use')

    parser.add_option('--environ'
      , action='store', type='str'
      , dest='environ', default=""
      , help='The environment to use')

    parser.add_option('-p', '--parcelport'
      , action='store', type='string'
      , dest='parcelport', default=default_env('HPXRUN_PARCELPORT', 'tcp')
      , help='Which parcelport to use (Options are: verbs, ipc, mpi, tcp) '
             '(environment variable HPXRUN_PARCELPORT')

    parser.add_option('-r', '--runwrapper'
      , action='store', type='string'
      , dest='runwrapper', default=default_env('HPXRUN_RUNWRAPPER', 'none')
      , help='Which runwrapper to use (Options are: none, mpi, srun, jsrun) '
             '(environment variable HPXRUN_RUNWRAPPER)')

    parser.add_option('-e', '--expected'
      , action='store', type='int'
      , dest='expected', default=default_env('HPXRUN_EXPECTED', '0')
      , help='Expected return codes of all invoked processes '
             '(environment variable HPXRUN_EXPECTED)')

    parser.add_option('-v', '--verbose'
      , action='store_true'
      , dest='verbose', default=False
            if default_env('HPXRUN_VERBOSE', '0') == '0' else True
      , help='Verbose output (environment variable HPXRUN_VERBOSE)')

    (options, args) = parser.parse_args()

    check_options(parser, options, args)
    if 'HPXRUN_ARGS' in os.environ:
        args += os.environ['HPXRUN_ARGS'].split()

    cmd = build_cmd(options, args)

    if options.verbose:
        print('Base command is "' + ' '.join(cmd) + '"')

    if options.nodes == '':
        nodes = None
    else:
        assert re.match(r'^\d+(\.\d+){3}(?::\d+|)(,\d+(\.\d+){3}(?::\d+|))*$', options.nodes), \
            "nodes must be a comma-delimited list of 'ip address:port' units."
        nodes = []
        for node in options.nodes.split(','):
            g = re.match(r'([\d.]+)(?::(\d+)|)$', node)
            ip = g.group(1)
            port = g.group(2)
            if port is None:
                port = str(options.defaultport)
            nodes += [ ip + ":" + port ]

    run(cmd, options.runwrapper, options.localities, nodes, options.verbose)

    if options.expected == 0:
        ret_expected = (lambda ret : True if ret == 0 else False)
    else:
        ret_expected = (lambda ret : False if ret == 0 else True)

    if len(procs) == 1:
        procs[0].wait()
        ret = procs[0].returncode
        if not ret_expected(ret):
            # Output which process failed
            msg = 'Process 0 failed with an unexpected error '
            msg += 'code of ' + str(ret) + ' (expected ' + str(options.expected)
            msg += ')'
            sys.exit(1)
        sys.exit(0)

    procs_lock = threading.Lock()
    returncode = 0

    def wait_on_proc(proc, which):
        global returncode
        proc.wait()
        ret = proc.returncode
        procs_lock.acquire()
        try:
            if not ret_expected(ret):
                returncode = 1
                # Output which process failed
                msg = 'Process ' + str(which) + ' failed with an unexpected error '
                msg += 'code of ' + str(ret) + ' (expected ' + str(options.expected)
                msg += ')'
                print(msg)
                while procs:
                    nextproc = procs.pop(0)
                    if nextproc != proc:
                        term(nextproc)
        except:
            pass
        finally: procs_lock.release()

    which = 0
    proc_watchdogs = []
    procs_lock.acquire()
    try:
        for proc in procs:
            proc_watchdog = threading.Thread(target=wait_on_proc, args=(proc, which))
            proc_watchdog.start()
            proc_watchdogs.append(proc_watchdog)
            which = which + 1
    except:
        pass
    finally: procs_lock.release()

    for proc_watchdog in proc_watchdogs:
        proc_watchdog.join()

    sys.exit(returncode)