Dragon distributed runtime for HPC and AI applications and workflows
import os
import sys
from typing import Any

from sphinx.ext import autodoc

# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# http://www.sphinx-doc.org/en/master/config

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
# import os
# import sys
sys.path.insert(0, os.path.abspath("./../src/dragon"))
sys.path.insert(0, os.path.abspath("."))

# -- Project information -----------------------------------------------------

project = "Dragon"
DragonVersion = "0.61"
copyright = "2023, Hewlett Packard Enterprise"
author = "Michael Burke, Eric Cozzi, Zach Crisler, Julius Donnert, Veena Ghorakavi, Nick Hill, Maria Kalantzi, Ben Keen, Kent D. Lee, Pete Mendygral, Davin Potts and Nick Radcliffe"

# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    "breathe",
    "sphinx.ext.autodoc",
    "sphinx.ext.autosummary",
    "sphinx.ext.intersphinx",
    "sphinx.ext.autosectionlabel",
    "sphinxfortran.fortran_domain",
    "sphinxfortran.fortran_autodoc",
    "sphinx_copybutton"
]

# autodoc_typehints = 'description'
autodoc_typehints_format = "short"
autodoc_class_signature = "separated"
autodoc_member_order = "bysource"
autosummary_generate = True

# changed this from "any" because we just get too many duplicates that way.
default_role = "code"

intersphinx_mapping = {"python": ("https://docs.python.org/3", None)}

# Add any paths that contain templates here, relative to this directory.
templates_path = ["_templates"]

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]

# Sphinx Fortran options
fortran_src = ["./../src/lib/"]
fortran_ext = ["f90"]

# autosectionlabel options
autosectionlabel_prefix_document = True
autosectionlabel_maxdepth = None

# provides numbering of figures for free.
numfig = True

# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
html_theme = "sphinx_rtd_theme"

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ["_static"]

# Paths relative to html_static_path
html_css_files = [
    "css/custom.css",
]

# See https://sphinx-rtd-theme.readthedocs.io/en/stable/configuring.html
html_theme_options = {}

breathe_projects = {"dragon": "../src/doxygen/xml"}
breathe_default_project = "dragon"

add_module_names = False

variables_to_export = [
    "DragonVersion",
]
frozen_locals = dict(locals())
rst_epilog = '\n'.join(map(lambda x: f".. |{x}| replace:: {frozen_locals[x]}", variables_to_export))
del frozen_locals

# Add an autodoc class that only posts the docstring without function
# names. Useful for autodoc-ing the Dragon CLI commands
class AutoDocstringOnly(autodoc.MethodDocumenter):
    objtype = "docstringonly"

    #do not indent the content
    content_indent = ""

    @classmethod
    def can_document_member(cls, member: Any, membername: str, isattr: bool, parent: Any) -> bool:
        return False

    #do not add a header to the docstring
    def add_directive_header(self, sig):
        pass

def setup(app):
    app.add_autodocumenter(AutoDocstringOnly)

doc/conf.py
DragonHPC/dragon
import dragon
from os import environ, getcwd
import multiprocessing as mp
from functools import partial

from socket import gethostname
from subprocess import Popen, PIPE, TimeoutExpired
import logging

# Pull in a lot of non-user facing dragon infrastructure to make logging pretty
from dragon.managed_memory import MemoryPool
from dragon.infrastructure import messages as dmsg
from dragon.infrastructure.parameters import this_process
from dragon.dlogging.util import setup_dragon_logging, setup_FE_logging, setup_BE_logging
from dragon.dlogging.util import LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING, DragonLoggingServices as dls
from dragon.dlogging.logger import DragonLogger
from dragon.utils import B64


class Shim:
    def __init__(self, rank, retries, timeout):
        self._rank = rank
        self._retries = retries
        self._timeout = timeout
        self._shutdown = False

    def _run_cmdline(self, cmdline):
        """Execute the user binary via subproces.Popen and return results"""

        log = logging.getLogger(f'shim_{gethostname()}_{self._rank}').getChild('run_cmdline)')
        emptystring_utf8 = "".encode("utf-8")
        # currently dragon.process.Popen is under development but seems close?
        proc = Popen(cmdline,
                     stderr=PIPE,
                     stdout=PIPE,
                     stdin=PIPE,
                     cwd=None,
                     env=None)
        try:
            out, err = proc.communicate(timeout=self._timeout)
        except TimeoutExpired:
            proc.kill()
            proc.communicate()
            log.info("FATAL: Failed to execute command, timeout reached")
            return -1, emptystring_utf8, emptystring_utf8
        except OSError:
            log.info("FATAL: Exception while attempting to start a process")
            return 1, emptystring_utf8, emptystring_utf8

        return proc.returncode, out.decode("utf-8"), err.decode("utf-8")

    def run_image(self, cmdline, logger_sdesc, use_dragon):

        # Configure handlers to make sure our log messages go to the DragonLogger channel
        def_mpool = None
        if use_dragon:
            def_mpool = MemoryPool.attach(B64.from_str(this_process.default_pd).decode())
        level, fname = setup_BE_logging(dls.TEST, logger_sdesc=B64.from_str(logger_sdesc), mpool=def_mpool)

        # Execute the user-provided application until it is successful or we run out
        # of retry attempts. Log as we do it
        log = logging.getLogger(f'shim_{gethostname()}_{self._rank}').getChild('run_image')
        log.info(f"Beginning execution, cmdline={' '.join(cmdline)}")

        while True:
            code, sout, serr = self._run_cmdline(cmdline)
            if code < 0:
                log.info(f"Execution FAILED; exit code={code}")
                log.info(f"Error output: {serr}")
            elif code > 0:
                log.info(f"Execution unscuccesful; exit code={code}")
                log.info(f"Error output: {serr}")
                break
            else:
                log.info("Execution complete")
                log.info(f'stdout: {sout}')
                break

            # see if we have any retries left
            if self._retries == 0:
                log.info("No more retries")
                break
            self._retries -= 1


def launch_shim(cmdline, logging_sdesc, use_dragon, shim):
    shim.run_image(cmdline, logging_sdesc, use_dragon)


def run_logging(logger_sdesc, LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING, use_dragon):
    """Run a proc pulling messages out of logging channel and logging them away"""

    # Configure logging handlers for our logging environment
    def_mpool = None
    if use_dragon:
        basename = 'dragon-shim'

        # Attach to the default memory pool
        def_mpool = MemoryPool.attach(B64.from_str(this_process.default_pd).decode())
    else:
        basename = 'regular-shim'


    setup_FE_logging(log_device_level_map=LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING,
                     basename=basename, basedir=getcwd())

    # Log a hello for this process
    local_log = logging.getLogger(f'logger_{gethostname()}')
    local_log.info('hello from logging thread')

    dragon_logger = DragonLogger.attach(B64.from_str(logger_sdesc).decode(), mpool=def_mpool)

    # Pull messages out of the DragonLogger and log them until told to stop
    while True:
        msg = dmsg.parse(dragon_logger.get(logging.INFO, timeout=None))
        if isinstance(msg, dmsg.HaltLoggingInfra):
            break
        elif isinstance(msg, dmsg.LoggingMsg):
            log = logging.getLogger(msg.name)
            log.log(msg.level, msg.msg, extra=msg.get_logging_dict())


if __name__ == "__main__":
    import argparse

    # This looks awful, but it's how we do all our infrastructure logging
    LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING['DRAGON_LOG_DEVICE_DRAGON_FILE'] = 'DEBUG'
    environ['DRAGON_LOG_DEVICE_DRAGON_FILE'] = 'DEBUG'

    # parse args
    parser = argparse.ArgumentParser()

    # Args to this driver program
    parser.add_argument("--retries", type=int, default=-1, help="Number of attempts to run an executable before exiting")
    parser.add_argument("--timeout", type=int, default=15, help="Time to wait on executable to complete")
    parser.add_argument("--num_workers", type=int, default=4, help="Number of pool workers for deploying executables")
    parser.add_argument("--dragon", action="store_true", help="Run with dragon multiprocessing")

    # The binary executable and its args
    parser.add_argument('prog', type=str, nargs='?', help="Executable pool workers will launch")
    parser.add_argument('args', type=str, nargs=argparse.REMAINDER, default=[], help="Option to the executable")

    args = parser.parse_args()

    # Opt to use dragon or not
    if args.dragon:
        print("using Dragon-augmented multiprocessing")
        mp.set_start_method("dragon")
    else:
        print("using base multiprocessing")
        mp.set_start_method("spawn")

    cmdline = [args.prog] + args.args
    print(f'cmdline = {cmdline}')

    # Create my dragon logger and run it in its own process.
    print(f'creating logger on {gethostname()}', flush=True)

    if args.dragon:
        dragon_logger = DragonLogger(mpool=MemoryPool.attach(B64.from_str(this_process.default_pd).decode()))
    else:
        dragon_logger = setup_dragon_logging(0)
    logging_sdesc = B64(dragon_logger.serialize())
    logging_proc = mp.Process(target=run_logging, args=(str(logging_sdesc), LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING, args.dragon))
    print('starting dragon logger receiver', flush=True)
    logging_proc.start()

    # set up my experiments as a pool and run them
    worker_pool = mp.Pool(args.num_workers)
    print(f"created pool with {args.num_workers} workers")
    launch_partial = partial(launch_shim, cmdline, str(logging_sdesc), args.dragon)
    shims = [Shim(idx, args.retries, args.timeout) for idx in range(args.num_workers)]
    print("entering shim.run_image", flush=True)
    pool_proc = worker_pool.map(launch_partial, shims, 1)
    print("Pool work is complete.", flush=True)

    # Tell the dragon logger to halt
    dragon_logger.put(dmsg.HaltLoggingInfra(tag=0).serialize(), logging.INFO)
    print("signaled logging exit", flush=True)
    logging_proc.join()
    print('joined on logging thread', flush=True)

    if args.dragon:
        dragon_logger.destroy(destroy_pool=False)
    else:
        dragon_logger.destroy()


examples/client_logging/logging_shim.py
DragonHPC/dragon
import os
import random
import logging
import multiprocessing as mp

from dragon.managed_memory import MemoryPool
from dragon.dlogging.logger import DragonLogger

POOL_SIZE = 4194304
DEFAULT_UID = 123456


def worker(log_ser):
    logger = DragonLogger.attach(log_ser)  # Attach to an existing logger
    msg = logger.get()
    print(f"Worker process attached, retrieved message: {msg}")
    logger.put("Salutations!")


if __name__ == "__main__":
    mpool = MemoryPool(
        POOL_SIZE, f"log_test_{os.getuid()}_{os.getpid()}", DEFAULT_UID, None
    )  # Create a memory pool for the logger
    logger = DragonLogger(mpool)  # Initialize the logger
    logger.put("Hello world!")  # Insert a log, default priority is INFO
    msg = logger.get()  # Retrieve a log, default priority is INFO
    print(msg)  # Should print "Hello World!"

    logger.put("Hello slightly more important world!", logging.WARNING)
    msg = logger.get(logging.INFO)  # Get a message of at least INFO priority
    print(msg)

    logger.put("Hello not very important world!", logging.DEBUG)
    msg = logger.get(logging.WARNING)  # Try to get a message of at least WARNING priority
    print(msg)  # msg returned will be None since the next message was below WARNING priority

    logger.put("Hello world (from somewhere else!)")
    log_ser = logger.serialize()  # Get a serialized descriptor of the logger
    proc = mp.Process(target=worker, args=(log_ser,))
    proc.start()
    proc.join()
    msg = logger.get()  # Retrieve the log we inserted from the other process
    print(f"Got message from worker process: {msg}")

    logger.destroy()

examples/dragon_core/logging.py
DragonHPC/dragon
import subprocess
from pathlib import Path
import dragon.infrastructure.parameters as dparms

CH_P2P_LATENCY = 'ch_p2p_latency'
CH_P2P_MSG_RATE = 'ch_p2p_msg_rate'
CH_P2P_BANDWIDTH = 'ch_p2p_bandwidth'

CH_P2P_TESTS = [
    CH_P2P_LATENCY,
    CH_P2P_MSG_RATE,
    CH_P2P_BANDWIDTH
]


def main():
    try:
        nodes_to_channel_map = eval(input())
        num_nodes = len(nodes_to_channel_map)

        if num_nodes < 2:
            raise ValueError('There must be at least 2 nodes allocated to run this test.', flush=True)

        if num_nodes % 2 != 0:
            raise ValueError('The number of allocated nodes must be a multiple of 2.', flush=True)

        this_node_index = dparms.this_process.index
        next_node_index = (this_node_index + 1) % num_nodes

        channel_0_index = next_node_index if this_node_index % 2 == 0 else this_node_index
        channel_0_descr = nodes_to_channel_map[channel_0_index][0]

        channel_1_index = this_node_index if this_node_index % 2 == 0 else next_node_index
        channel_1_descr = nodes_to_channel_map[channel_1_index][0]

        for ch_p2p_test in CH_P2P_TESTS:
            print(f'{this_node_index}: Starting {ch_p2p_test}', flush=True)

            proc = subprocess.Popen(
                [f'{Path.cwd() / ch_p2p_test}',
                 str(this_node_index % 2),         # 0 / 1 node id
                 dparms.this_process.default_pd,   # descriptor for default memory pool
                 channel_0_descr,                  # descriptor for channel 0
                 channel_1_descr],                 # descriptor for channel 1
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE
            )

            stdout, stderr = proc.communicate()

            for stream_name, stream_handle in [('stdout', stdout), ('stderr', stderr)]:
                for line in stream_handle.splitlines():
                    if isinstance(line, bytes):
                        line = line.decode()
                    print(f'{this_node_index}-{stream_name}: {line}')

            print(f'{this_node_index}: app exited with {proc.returncode}')
            print(f'{this_node_index}: Done', flush=True)

    except Exception as ex:
        print("!!!!There was an EXCEPTION!!!!!")
        print(repr(ex))


if __name__ == "__main__":
    main()

examples/dragon_core/performance/ch_p2p_benchmark.py
DragonHPC/dragon
import dragon
import multiprocessing as mp
import subprocess
import sys
import dragon.channels as dch
import dragon.managed_memory as dm
import dragon.infrastructure.parameters as dp
import dragon.infrastructure.facts as df
import dragon.utils as du
import time

def start_ringproc(iterations, cuid, receive_from_channel_sdesc, ret_queue):
    proc = subprocess.Popen(['ringproc', str(iterations), str(cuid), receive_from_channel_sdesc], stdout=subprocess.PIPE)
    send_to_channel_sdesc = proc.stdout.readline()
    while len(send_to_channel_sdesc.strip()) == 0:
        send_to_channel_sdesc = proc.stdout.readline()
    ret_queue.put(send_to_channel_sdesc)
    proc.wait()
    if proc.returncode != 0:
        print('*******Proc exited with rc=', proc.returncode, flush=True)

def usage():
    print('usage: dragon ring.py <num_procs> <iterations>')
    print('    <num_procs> is the number of processes to start, one per node.')
    print('    <iterations> is the number of times each process forwards a message')
    print('                to the next node.')
    print('    The program creates a ring across the user specified number of')
    print('    nodes and sends a message around a ring of nodes. The num_procs')
    print('    and iterations must be greater than 0.')
    sys.exit(1)

def main():
    try:
        if len(sys.argv) != 3:
            raise ValueError()

        mp.set_start_method('dragon')
        ring_size = int(sys.argv[1])
        iterations = int(sys.argv[2])
        if iterations <= 0 or ring_size <= 0:
            raise ValueError()
    except:
        usage()

    pool = dm.MemoryPool.attach(du.B64.str_to_bytes(dp.this_process.default_pd))
    origin_channel = dch.Channel(pool, df.BASE_USER_MANAGED_CUID)
    receive_sdesc = du.B64.bytes_to_str(origin_channel.serialize())
    final_channel = dch.Channel(pool, df.BASE_USER_MANAGED_CUID+1)
    final_sdesc = du.B64.bytes_to_str(final_channel.serialize())
    origin_send_sdesc = receive_sdesc

    ret_queue = mp.Queue()
    mp_procs = []
    for i in range(1,ring_size):
        proc = mp.Process(target=start_ringproc, args=(str(iterations), str(i+df.BASE_USER_MANAGED_CUID+1), receive_sdesc, ret_queue))
        proc.start()
        mp_procs.append(proc)
        receive_sdesc = ret_queue.get().strip()

    # This final process starts on the current node and completes the ring. It
    # also provides the destination for the final message to be returned.
    proc = subprocess.Popen(['ringproc', str(iterations), str(df.BASE_USER_MANAGED_CUID), receive_sdesc, origin_send_sdesc, final_sdesc], stdout=subprocess.PIPE)

    reader = dch.ChannelRecvH(final_channel)
    writer = dch.ChannelSendH(origin_channel)
    reader.open()
    writer.open()
    start = time.perf_counter()
    writer.send_bytes(b'hello', timeout=None, blocking=True)
    last_msg = reader.recv_bytes(timeout=None, blocking=True)
    stop = time.perf_counter()

    avg_time = (stop - start) / (iterations*ring_size)
    proc.wait()
    print('Ring proc exited...', flush=True)
    for proc in mp_procs:
        proc.join()
        print('Ring proc exited...', flush=True)
    if last_msg == b'hello':
        print('Test Passed.', flush=True)
        print(f'The average time per message transfer was {avg_time*1e6} microseconds.')
    else:
        print('Test Failed.', flush=True)
    print('Main proc exiting...', flush=True)


if __name__ == '__main__':
    main()
examples/dragon_core/ring.py
DragonHPC/dragon
import enum
import time
import dragon
import multiprocessing as mp
import argparse
import string
import random
from dragon.data.distdictionary.dragon_dict import DragonDict

@enum.unique
class DictOp(enum.Enum):
    """Action to be performed on the dictionary item"""
    SET_ITEM = enum.auto()
    GET_ITEM = enum.auto()
    DEL_ITEM = enum.auto()


def do_dict_ops(_keys, _dict, client_id, iterations, msg_size, result_link, dict_op):
    """Function used to execute operations on the shared Dragon dictionary

    :param _keys: List of all the keys of the dictionary
    :type _keys: list
    :param _dict: A dragon dictionary object
    :type _dict: dragon dictionary
    :param client_id: Unique ID of the client
    :type client_id: int
    :param iterations: Number of iterations to perform a dictionary operation
    :type iterations: int
    :param msg_size: Number of characters used for the length of the value
    :type msg_size: int
    :param result_link: A pipe used for the communication of the results
    :type result_link: connection object
    :param dict_op: Enum that controls the operations of the dictionary
    :type dict_op: enum
    """

    if dict_op == DictOp.SET_ITEM or dict_op == DictOp.DEL_ITEM:
        letters = string.ascii_letters
        value = ''.join(random.choice(letters) for i in range(msg_size))

    if dict_op == DictOp.SET_ITEM:
        start = time.monotonic()
        for i in range(iterations):
            key = random.choice(_keys)
            _dict[key] = value
        end = time.monotonic()
    elif dict_op == DictOp.GET_ITEM:
        start = time.monotonic()
        for i in range(iterations):
            key = random.choice(_keys)
            val = _dict[key]
        end = time.monotonic()
    elif dict_op == DictOp.DEL_ITEM:
        start = time.monotonic()
        for i in range(iterations):
            key = random.choice(_keys)
            del _dict[key]
            _dict[key] = value
        end = time.monotonic()

    result_link.send((start, end))
    if client_id == 0:
        print(f"DictOp {dict_op.value}: I am client {client_id}. (start): {start} -- (end): {end}. (end - start): {end - start}. Elapsed time: {(end - start) / iterations} sec", flush=True)

    if mp.get_start_method() == "dragon":
        _dict.close()

def generate_keys(dict_size=100):
    """Generate a list including the keys that will be used for the dictionary.

    :param dict_size: Total number of keys to be populated in the dicitonary
    :type dict_size: int
    :return: List of keys to be stored along with values in the dictionary
    :rtype: list
    """
    my_keys = list()
    letters = string.ascii_letters

    for _ in range(dict_size):
        # each key is 30 characters long
        key = ''.join(random.choice(letters) for i in range(30)) # characters can be repeated
        my_keys.append(key)

    assert len(my_keys) == dict_size
    return my_keys


def assign_keys(_dict, keys, value_size):
    """Initiate the dictionary. Assign the values to each key in the provided list.
    Each value is a string of msg_size characters long.

    :param _dict: A dragon dictionary object
    :type _dict: dragon dictionary
    :param keys: List of keys to assign values in the dictionary
    :type keys: list
    :param value_size: Number of characters used for the length of the value
    :type value_size: int
    """
    for key in keys:
        letters = string.ascii_letters
        value = ''.join(random.choice(letters) for i in range(value_size))
        _dict[key] = value


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Distributed dictionary benchmark')
    parser.add_argument('--dragon', action='store_true', help='run using dragon')
    parser.add_argument('--dict_size', type=int, default=500,
                        help='Number of (key,value) pairs inside the dict')
    parser.add_argument('--value_size', type=int, default=1024,
                        help='size of the value (bytes) that are stored in the dict')
    parser.add_argument('--num_nodes', type=int, default=1,
                        help='number of nodes the dictionary distributed across')
    parser.add_argument('--clients', type=int, default=10,
                        help='number of client processes performing operations on the dict')
    parser.add_argument('--managers_per_node', type=int, default=1,
                        help='number of managers per node for the dragon dict')
    parser.add_argument('--total_mem_size', type=int, default=1,
                        help='total managed memory size for dictionary in GB')
    parser.add_argument('--batch_size', type=int, default=100,
                        help='number of kv pairs added by each process before dict operations')
    parser.add_argument('--iterations', type=int, default=1000,
                        help='number of iterations')
    parser.add_argument('--dict_ops', type=int, default=0,
                        help='choose the operations to be performed on the dict -- '
                        '0 to set the values, '
                        '1 to get the values, '
                        '2 for both, '
                        '3 for deletes (includes setting the values for now to continue further deletes)')

    my_args = parser.parse_args()

    # The following doesn't seem to work for setting the seed when multiple managers and Dragon.
    # We need to set 'export PYTHONHASHSEED=1234' for now.

    # hashseed = os.getenv('PYTHONHASHSEED')
    # if not hashseed:
    #     os.environ['PYTHONHASHSEED'] = '0'
    #     os.execv(sys.executable, [sys.executable] + sys.argv)

    if my_args.dragon:
        print('using Dragon')
        mp.set_start_method("dragon")
        # Convert the total size of the dictionary (in GB) to bytes
        total_mem_size = my_args.total_mem_size * (1024*1024*1024)
        dd = DragonDict(my_args.managers_per_node, my_args.num_nodes, total_mem_size)
    else:
        print('using multiprocessing')
        mp.set_start_method("spawn")
        manager = mp.Manager()
        dd = manager.dict()

    num_clients = my_args.clients
    value_size = my_args.value_size

    _keys = generate_keys(dict_size=my_args.dict_size)

    # Parallelize the initialization with batch size of key/value entries
    num_keys = len(_keys)
    batch_size = my_args.batch_size
    num_batches = num_keys // batch_size + int(num_keys % batch_size != 0)
    jobs = []
    for i in range(num_batches):
        if (i == num_batches - 1) and (num_keys % batch_size != 0):
            batch_keys = _keys[i * batch_size:]
        else:
            batch_keys = _keys[i * batch_size:(i+1) * batch_size]
        jobs.append(mp.Process(target=assign_keys, args=(dd, batch_keys, value_size)))

    # Complete the initialization of the dictionary
    _ = [p.start() for p in jobs]
    _ = [p.join() for p in jobs]
    _ = [p.terminate() for p in jobs] # make sure we clean everything

    length = len(dd)
    print(f'Length of the dictionary is {length}', flush=True)

    dict_ops = []
    if my_args.dict_ops == 0:
        dict_ops.append(DictOp.SET_ITEM)
    elif my_args.dict_ops == 1:
        dict_ops.append(DictOp.GET_ITEM)
    elif my_args.dict_ops == 2:
        dict_ops.append(DictOp.SET_ITEM)
        dict_ops.append(DictOp.GET_ITEM)
    elif my_args.dict_ops == 3:
        dict_ops.append(DictOp.DEL_ITEM)

    for ii in range(len(dict_ops)):
        result_links = [mp.Pipe(duplex=False) for _ in range(num_clients)]
        try:
            procs = []
            for i in range(num_clients):
                client_proc = mp.Process(target=do_dict_ops,
                                         args=(_keys, dd, i, my_args.iterations,
                                               value_size, result_links[i][1], dict_ops[ii],))
                client_proc.start()
                procs.append(client_proc)

            for i in range(len(procs)):
                procs[i].join()

            # min_start = 1.0e9
            # since we have joined the processes, we know this value will be greater than the processes' corresponding values
            min_start = time.monotonic()
            max_end = 0.0
            for i in range(num_clients):
                start, end = result_links[i][0].recv()
                min_start = min(min_start, start)
                max_end = max(max_end, end)

            result = (max_end - min_start) / my_args.iterations
            rate = (my_args.iterations * num_clients) / (max_end - min_start) # aggregated rate

            print(f"\n{dict_ops[ii]}:", flush=True)
            print(f"Msglen [B]   Lat [sec]")
            print(f"{value_size}  {result}")
            print(f"Msglen [B]   Rate ")
            print(f"{value_size}  {rate}\n")

            for i in range(len(procs)):
                procs[i].kill()

        except Exception as e:
            print(e)

    if my_args.dragon:
        dd.stop()

examples/dragon_data/dist_dict_bench.py
DragonHPC/dragon
import dragon
import argparse
import multiprocessing as mp

from dragon.data.distdictionary.dragon_dict import DragonDict

def _retrieve_value(_dict, key, client_id):
    value = _dict[key]
    print(f'Retrieving value:{value} for key:{key} for client id:{client_id} from the dictionary', flush=True)
    print(f'Closing the connection with client id:{client_id}', flush=True)
    _dict.close()
    return value

def _store_key_value(_dict, key, value, client_id):
    print(f'Storing key:{key} and value:{value} from client id:{client_id} into the dictionary', flush=True)
    _dict[key] = value
    print(f'Closing the connection with client id:{client_id}', flush=True)
    _dict.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Distributed dictionary example')
    parser.add_argument('--num_nodes', type=int, default=1,
                        help='number of nodes the dictionary distributed across')
    parser.add_argument('--managers_per_node', type=int, default=1,
                        help='number of managers per node for the dragon dict')
    parser.add_argument('--total_mem_size', type=int, default=1,
                        help='total managed memory size for dictionary in GB')

    my_args = parser.parse_args()
    mp.set_start_method("dragon")

    # Instantiate the dictionary and start the processes
    total_mem_size = my_args.total_mem_size * (1024*1024*1024)
    dd = DragonDict(my_args.managers_per_node, my_args.num_nodes, total_mem_size)

    client_proc_1 = mp.Process(target=_store_key_value, args=(dd, "Hello", "Dragon", 1))
    client_proc_1.start()
    client_proc_1.join()

    client_proc_2 = mp.Process(target=_retrieve_value, args=(dd, "Hello", 2))
    client_proc_2.start()
    client_proc_2.join()

    print(f'Done here. Stopping the Dragon Dictionary', flush=True)
    dd.stop()
examples/dragon_data/dist_dict_client.py
DragonHPC/dragon
"""Single producer and single consumer communication scheme with Dragon Channels.

This example demonstrates the optimal way to communicate data that will be serialized
with Pickle or other serializers which interact with a file-like interface. It shows
how large payloads (larger than the Dragon Managed Memory Pool) can be communicated.

* Create a single Dragon Channel.
* Create one process that acts as the producer and one as the consumer. The producer
    sends a large size message, in this example 8GB is used, and the consumer waits
    until receiving the message.
"""

import sys
import os
import random
import pickle

import dragon.globalservices.channel as dgchan
import dragon.globalservices.process as dgprocess

import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts
from dragon.infrastructure.process_desc import ProcessOptions

from dragon.channels import Channel, Peer2PeerWritingChannelFile, Peer2PeerReadingChannelFile


def producer(chan_name, msg):
    print(f"I am the producer.", end=" ")

    channel_desc = dgchan.query(chan_name)
    chan = Channel.attach(channel_desc.sdesc)

    my_size = sys.getsizeof(msg)
    print(f"Message size to be sent: {my_size} bytes.", flush=True)

    try:
        adapter = Peer2PeerWritingChannelFile(chan)
        adapter.open()
        pickle.dump(msg, file=adapter, protocol=5)
    finally:
        adapter.close()


def consumer(chan_name, msg):
    channel_desc = dgchan.query(chan_name)
    chan = Channel.attach(channel_desc.sdesc)

    try:
        adapter = Peer2PeerReadingChannelFile(chan)
        adapter.open()
        rcv_msg = pickle.load(adapter)
    finally:
        adapter.close()

    assert rcv_msg == msg, "The msg received is not correct. Something went wrong."
    print(f"I am the consumer. Successful communication.", flush=True)


def main():
    default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)

    # ask GS to create a dragon channel
    channel = dgchan.create(default_muid, user_name="InputChannel")

    # create processes
    cmd = sys.executable
    wdir = "."
    options = ProcessOptions(make_inf_channels=True)
    proc1 = dgprocess.create(
        cmd, wdir, [os.path.basename(__file__), "sub", "producer", channel.name], None, options=options
    )
    proc2 = dgprocess.create(
        cmd, wdir, [os.path.basename(__file__), "sub", "consumer", channel.name], None, options=options
    )

    dgprocess.join(proc1.p_uid)
    dgprocess.join(proc2.p_uid)

    dgchan.destroy(channel.c_uid)


if __name__ == "__main__":

    # get size of default pool
    pool_size = int(os.getenv("DRAGON_DEFAULT_SEG_SZ"))

    print(f"{pool_size=}", flush=True)

    msg_size = 2 * pool_size  # message does not fit into pool
    msg = bytearray(msg_size)

    if len(sys.argv) > 1 and sys.argv[1] == "sub":
        if sys.argv[2] == "producer":
            producer(sys.argv[3], msg)
        else:
            consumer(sys.argv[3], msg)
    else:
        print(f"Memory Pool Size is {pool_size} bytes", flush=True)

        main()

examples/dragon_gs_client/connection_demo.py
DragonHPC/dragon
""" Pass a message between two Dragon processes using one Dragon channel. Uses Global Services and Channels
from Dragon API.
    * Create a single channel and process on default Dragon pool.  In the new process, attach to the Channel
    * passed as an argument via the serialized descriptor.  In the new process, use subprocess.Popen to execute
    * the shell command and get the output. Write that output into the attached Channel.  In the original process
    * in main, output the message.  Join the channel and the process with Global Services.
"""


import dragon.globalservices.channel as channel
import dragon.globalservices.process as process

from dragon.infrastructure.process_desc import ProcessOptions
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts

from dragon.channels import Channel, Message

import subprocess
import sys


def main():
    """The channel dragon is created in main. The data from dragon is printed in main."""

    try:

        if len(sys.argv) != 2:
            print("You must pass a shell command as the first argument")
            return

        shell_cmd = sys.argv[1]
        default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
        c_desc = channel.create(default_muid, user_name="dragon")

        p = process.create(
            "python3",
            ".",
            ["dragon_popen_api.py", "sub", c_desc.name, shell_cmd],
            None,
            options=ProcessOptions(make_inf_channels=True),
        )

        ch = Channel.attach(c_desc.sdesc)
        recv_h = ch.recvh()
        recv_h.open()

        reading_output = True
        while reading_output:
            recv_msg = recv_h.recv()
            recv_msgview = recv_msg.bytes_memview()
            data_flag = recv_msgview[:4].tobytes().decode("utf-8")

            if data_flag == "EOT ":
                reading_output = False

            else:
                data = recv_msgview[4:].tobytes().decode("utf-8")
                print(data)

            recv_msg.destroy()

        process.join(p.p_uid)
        ch.detach()
        channel.destroy(c_desc.c_uid)
        print("Exiting main")

    except Exception as e:
        print(e)


def submain():
    """The output from the shell command passed from main to submain is written in the Dragon channel by the submain process."""

    try:

        shell_cmd = sys.argv[3]
        channel_name = sys.argv[2]

        c_desc = channel.join(channel_name)
        serialized_channel = c_desc.sdesc
        ch = Channel.attach(serialized_channel)
        mpool = ch.get_pool()

        send_handle = ch.sendh()
        send_handle.open()

        # Subprocess.Popen captures the output of the shell command while the program runs.
        proc = subprocess.Popen([shell_cmd], stdout=subprocess.PIPE)
        data = proc.stdout.read()

        while data != None and len(data) != 0:
            msg = Message.create_alloc(mpool, len(data) + 4)
            msg_view = msg.bytes_memview()
            msg_view[:4] = bytes("data", "utf-8")
            msg_view[4:] = data
            send_handle.send(msg)
            data = proc.stdout.read()

        msg = Message.create_alloc(mpool, 4)
        msg_view = msg.bytes_memview()
        msg_view[:] = bytes("EOT ", "utf-8")
        send_handle.send(msg)
        send_handle.close()

        msg.destroy()
        ch.detach()
        mpool.detach()

    except Exception as e:
        print(e)


if __name__ == "__main__":
    if len(sys.argv) >= 2 and sys.argv[1] == "sub":
        submain()
    else:
        main()

examples/dragon_gs_client/dragon_popen_api.py
DragonHPC/dragon
""" Pass a message between two Dragon processes using one Dragon channel. Uses Global Services and Channels
from Dragon API.
    * Create a single channel and process on default Dragon pool.  In the new process, attach to the Channel
    * passed as an argument via the serialized descriptor.  In the new process, use subprocess.run to execute the
    * shell command and get the output. Write that output into the attached Channel.  In the original process in
    * main, output the message.  Join the channel and the process with Global Services.
"""

import dragon.globalservices.channel as channel
import dragon.globalservices.process as process
from dragon.infrastructure.process_desc import ProcessOptions
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts

from dragon.channels import Channel, Message

import sys
import subprocess


def main():
    """The channel dragon is created in main. The data from dragon is printed in main."""

    try:

        if len(sys.argv) != 2:
            print("You must pass a shell command as the first argument")
            return

        shell_cmd = sys.argv[1]
        default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
        c_desc = channel.create(default_muid, user_name="dragon")

        p = process.create(
            "python3",
            ".",
            ["dragon_run_api.py", "sub", c_desc.name, shell_cmd],
            None,
            options=ProcessOptions(make_inf_channels=True),
        )

        ch = Channel.attach(c_desc.sdesc)

        recv_h = ch.recvh()
        recv_h.open()
        recv_msg = recv_h.recv()
        recv_msgview = recv_msg.bytes_memview()
        data = recv_msgview[:].tobytes().decode("utf-8")
        print(data)
        recv_msg.destroy()

        process.join(p.p_uid)
        ch.detach()
        channel.destroy(c_desc.c_uid)
        print("Exiting main")

    except Exception as e:
        print(e)


def submain():
    """The output from the shell command passed from main to submain is written in the Dragon channel by the
    submain process.
    """

    try:

        shell_cmd = sys.argv[3]
        channel_name = sys.argv[2]

        c_desc = channel.join(channel_name)
        serialized_channel = c_desc.sdesc
        ch = Channel.attach(serialized_channel)
        mpool = ch.get_pool()

        send_handle = ch.sendh()
        send_handle.open()

        # Subprocess.run captures the output of the shell command after the program exits.
        proc = subprocess.run([shell_cmd], stdout=subprocess.PIPE)
        data = proc.stdout
        msg = Message.create_alloc(mpool, len(data))
        msg_view = msg.bytes_memview()
        msg_view[:] = data
        send_handle.send(msg)
        send_handle.close()

        msg.destroy()
        ch.detach()
        mpool.detach()

    except Exception as e:
        print(e)


if __name__ == "__main__":
    if len(sys.argv) >= 2 and sys.argv[1] == "sub":
        submain()
    else:
        main()

examples/dragon_gs_client/dragon_run_api.py
DragonHPC/dragon
""" Pass a message between two Dragon processes using one Dragon channel. Uses Global Services and Channels
from Dragon API.
    * Create a server process that communicates with the client.
    * The Dragon channel is created on the default memory pool.
    * The client handles user input from the requests.txt file.
    * The client passes the request information on the Dragon channel to the server process.
    * The server creates a server:client channel to pass the messages to the client.
    * After all the requests are processed or the stop request is processed, the server stops listening for requests.
"""


import dragon.globalservices.channel as channel
import dragon.globalservices.process as process

import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts
from dragon.infrastructure.process_desc import ProcessOptions

from dragon.channels import Channel, Message

import sys
import os


def response_message_array(request):
    """This function creates the response for the request to be placed in the response channel."""

    messages = []
    request_types = ["bananas", "apples", "rodents", "stop"]

    if request not in request_types:
        print("Request value ", request, " is not supported.")
    if request == "apples":
        messages = ["honeycrisp", "red delicious", "braeburn"]
    if request == "bananas":
        messages = ["plantain", "cavendish", "java"]
    if request == "rodents":
        messages = ["squirrels", "mice", "capybaras"]
    if request == "stop":
        messages = "stop"

    output = "The request: " + request + " has the messages " + str(messages)
    return output


def create_requests(filename):
    """This function will be replaced once input() works with Dragon. The request textfile is read into an
    array of request."""
    request_array = []

    with open(filename, "r") as f:
        for line in f:
            request_array.append(line.strip())
    f.close()

    return request_array


def server(channel_name):
    """The server creates the request channel and the response channel. It provides the response messages for
    the requests from the user."""
    receiving = True

    # Join the server channel that receives the request from user
    server_descriptor = channel.join(channel_name)
    server_serialized_channel = server_descriptor.sdesc
    server_channel = Channel.attach(server_serialized_channel)
    mpool = server_channel.get_pool()
    server_receive_handle = server_channel.recvh()
    server_receive_handle.open()

    while receiving:
        request_msg = server_receive_handle.recv()
        request_msgview = request_msg.bytes_memview()
        request_puid_channel_name = request_msgview[:].tobytes().decode("utf-8")
        request = request_puid_channel_name.split("-")[0]
        request_msg.destroy()

        if "stop" == request:
            receiving = False

        # Create the response channel to send the messages of the responses to the request
        response_channel_name = request_puid_channel_name
        default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
        response_c_desc = channel.create(default_muid, user_name=response_channel_name)
        response_serialized_channel = response_c_desc.sdesc
        response_channel = Channel.attach(response_serialized_channel)

        response_send_handle = response_channel.sendh()
        response_send_handle.open()
        response = response_message_array(request)
        response_msg = Message.create_alloc(mpool, len(response))
        response_msg_view = response_msg.bytes_memview()
        response_msg_view[:] = bytes(response, "utf-8")
        response_send_handle.send(response_msg)
        response_send_handle.close()

        response_msg.destroy()
        response_channel.detach()

    server_receive_handle.close()


def main():
    """The main provides an interface for the user to communicate to the server. This is where the requests
    are sent and the responses are printed.
    """

    default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
    server_channel_create = channel.create(default_muid, user_name="dragon")
    channel_name = server_channel_create.name

    requests = create_requests("requests.txt")

    # Create the server process
    cmd = sys.executable
    wdir = "."
    options = ProcessOptions(make_inf_channels=True)
    server_process = process.create(
        cmd, wdir, [os.path.basename(__file__), "server", channel_name], None, options=options
    )

    # Create the server channel where the user passes the requests to the server
    server_cdesc = channel.join(channel_name)
    server_serialized_channel = server_cdesc.sdesc
    server_channel = Channel.attach(server_serialized_channel)
    mpool = server_channel.get_pool()
    server_send_handle = server_channel.sendh()
    server_send_handle.open()

    for request in requests:

        # Send the request to the server. The server will create the request channel named after the request.
        puid = dparm.this_process.my_puid
        response_channel = str(request) + "-" + str(puid)
        msg = Message.create_alloc(mpool, len(response_channel))
        msg_view = msg.bytes_memview()
        msg_view[:] = bytes(response_channel, "utf-8")
        server_send_handle.send(msg)
        msg.destroy()

        # Output the messages from the response channel.
        response_desc = channel.join(response_channel)
        response_serialized_channel = response_desc.sdesc
        response_channel = Channel.attach(response_serialized_channel)
        response_recv_h = response_channel.recvh()
        response_recv_h.open()
        response_msg = response_recv_h.recv()
        response_msgview = response_msg.bytes_memview()
        response_data = response_msgview[:].tobytes().decode("utf-8")

        if response_data != "stop":
            print(response_data, flush=True)

        response_msg.destroy()
        response_recv_h.close()
        channel.destroy(response_desc.c_uid)

    process.join(server_process.p_uid)
    server_send_handle.close()
    server_channel.destroy()
    mpool.detach()


if __name__ == "__main__":
    try:
        if len(sys.argv) > 1:
            server(sys.argv[2])
        else:
            main()
    except Exception as e:
        print(e)

examples/dragon_gs_client/dragon_server_client_api.py
DragonHPC/dragon
""" The serialized descriptor for the managed memory allocation is passed from the input channel to the output channel to access the allocation where the pids
for the processes are written and stored.

    * Coordinating process creates two channels (input, output) and gets serialized descriptors of each
    * Coordinating process creates an allocation from the default Pool of 32 B and writes its PID into the first 8 byes and 0 everywhere else.
    * For the three processes started by the coordinating process, the pid and any payload is written to the first 0s found in the 32 bytes.
    * Once all the descriptors and payload are written to the message, the message is printed and the GS objects are cleaned up.

"""
import sys
import os

import dragon.infrastructure.parameters as dparm
import dragon.utils as du
import dragon.infrastructure.facts as dfacts
from dragon.infrastructure.process_desc import ProcessOptions
import dragon.globalservices.channel as dgchan
import dragon.globalservices.process as dgprocess
from dragon.managed_memory import MemoryPool, MemoryAlloc
from dragon.channels import Channel


def worker(input_channel_sdesc, output_channel_sdesc):
    """The workers write their pids to the managed memory allocation."""

    input_channel_desc = du.B64.str_to_bytes(input_channel_sdesc)
    output_channel_desc = du.B64.str_to_bytes(output_channel_sdesc)

    ch_in = Channel.attach(input_channel_desc)
    ch_out = Channel.attach(output_channel_desc)

    recv_handle = ch_in.recvh()
    send_handle = ch_out.sendh()

    recv_handle.open()
    send_handle.open()

    serialized_mem_descr = recv_handle.recv_bytes()

    # Worker process attaches to the memory allocation created by coordinating process

    mem = MemoryAlloc.attach(serialized_mem_descr)
    memview = mem.get_memview()

    i = 0
    found = False
    while i < 32 and not found:

        # write pid of process to first set of 8 bytes equal to 0

        if int.from_bytes(memview[i : i + 8], "little") == 0:
            memview[i : i + 8] = os.getpid().to_bytes(8, "little")
            found = True
        i += 8

    send_handle.send_bytes(serialized_mem_descr)

    mem.detach()

    recv_handle.close()
    send_handle.close()

    ch_in.detach()
    ch_out.detach()


def main():
    """The coordinating process creates the managed memory allocation, the two channels, and the processes.
    The first 8 bytes of the managed memory are dedicated to the coordinating process's pid."""

    default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)

    input_channel = dgchan.create(default_muid, user_name="input-channel")
    output_channel = dgchan.create(default_muid, user_name="output-channel")

    # The coord_in channel is where the coordinating process waits for the message on output. The coord_out channel is the input channel where the coordinating process sends the serialized memory descriptor for the memory allocation to the worker processes.
    coord_in = Channel.attach(output_channel.sdesc)
    coord_out = Channel.attach(input_channel.sdesc)

    recv_handle = coord_in.recvh()
    send_handle = coord_out.sendh()

    recv_handle.open()
    send_handle.open()

    # Coordinating process attaches to the default memory pool

    mpool = MemoryPool.attach(coord_in.get_pool().serialize())

    # Memory allocation of 32 bytes is created in the default memory pool

    mem = mpool.alloc(32)
    memview = mem.get_memview()

    # Write coordinating process PID to first 8 bytes

    memview[:8] = os.getpid().to_bytes(8, "little")

    cmd = sys.executable
    wdir = "."
    cmd_params = [
        os.path.basename(__file__),
        "worker",
        du.B64.bytes_to_str(input_channel.sdesc),
        du.B64.bytes_to_str(output_channel.sdesc),
    ]

    options = ProcessOptions(make_inf_channels=True)
    processes = []

    for _ in range(3):
        p = dgprocess.create(cmd, wdir, cmd_params, None, options=options)
        processes.append(p.p_uid)

    for _ in range(len(processes)):
        send_handle.send_bytes(mem.serialize())
        recv_handle.recv_bytes()

    memview = mem.get_memview()

    # Convert memory allocation data from bytes to integers/pids

    pid_list = [int.from_bytes(memview[i : i + 8], "little") for i in range(0, len(memview), 8)]
    print("Completed message with all pids:", pid_list)

    recv_handle.close()
    send_handle.close()

    coord_in.detach()
    coord_out.detach()

    dgprocess.multi_join(processes, join_all=True)

    dgchan.destroy(input_channel.c_uid)
    dgchan.destroy(output_channel.c_uid)

    mpool.destroy()


if __name__ == "__main__":
    try:
        if len(sys.argv) > 1:
            worker(sys.argv[2], sys.argv[3])
        else:
            sys.stdout.flush()
            main()
    except Exception as e:
        print(e)

examples/dragon_gs_client/managed_mem_demo.py
DragonHPC/dragon
""" Calculate pi in parallel using the Dragon GS Client interface.

    * Create two Dragon Channels to be used as an in-queue and out-queue for work
    * Start up a collection of processes (number is an argument to the script) each of
       which will execute some function based on input data. Let's have that function do the simple thing of
       helping to estimate Pi by taking two input random number and checking if the sum of the squares is less
       than 1. If so, have the function return True. Pass the descriptors to the 2 Channels to each process as
       startup information.
    * Have all processes attach to the Channels and use one of them to pass the two numbers and the other to
       return True or False.
    * The coordinating process then generates random number pairs in the range of [0 .. 1] and passes pairs of
       them as messages into one of the Channels.
    * The worker then should wait for messages on that Channel as work and return the result in the other
       Channel.
    * The coordinator will need to generate random pairs and accumulate results at the same time and use the ratio
       of True results to the total number of pairs to compute Pi
    * The coordinator then should send some terminate message into the work queue Channel and clean everything up.
"""

import sys
import os
import random
import pickle
import math

# Dragon GS Client API
import dragon.globalservices.pool as dgpool
import dragon.globalservices.channel as dgchan
import dragon.globalservices.process as dgprocess

# Dragon infrastructure API
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts
from dragon.infrastructure.process_desc import ProcessOptions

# Dragon Core API
from dragon.managed_memory import MemoryPool
from dragon.channels import Channel


def pi_helper(output_channel_name, input_channel_name):

    # get serialized channel descriptors from channel names
    input_channel_desc = dgchan.query(input_channel_name)
    output_channel_desc = dgchan.query(output_channel_name)

    ch_in = Channel.attach(input_channel_desc.sdesc)
    ch_out = Channel.attach(output_channel_desc.sdesc)

    # we are using the message interface to channels here, not
    # the file-like interface.
    recv_handle = ch_in.recvh()
    send_handle = ch_out.sendh()

    recv_handle.open()
    send_handle.open()

    while True:  # keep working

        msg = recv_handle.recv_bytes()
        x, y = pickle.loads(msg)

        if x == -1 and y == -1:  # we're done here
            break

        r = x * x + y * y

        if r > 1.0:
            msg = pickle.dumps(False)
        else:
            msg = pickle.dumps(True)

        send_handle.send_bytes(msg)

    recv_handle.close()
    send_handle.close()

    ch_in.detach()
    ch_out.detach()


def main(num_workers):

    # use default memory pool of this process
    default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)

    # ask GS to create 2 dragon channels
    input_channel = dgchan.create(default_muid, user_name="InputChannel")
    output_channel = dgchan.create(default_muid, user_name="OutputChannel")

    # create processes
    cmd = sys.executable
    wdir = "."
    cmd_params = [
        os.path.basename(__file__),
        "sub",
        input_channel.name,
        output_channel.name,
    ]
    options = ProcessOptions(make_inf_channels=True)
    processes = []
    for _ in range(num_workers):
        p = dgprocess.create(cmd, wdir, cmd_params, None, options=options)
        processes.append(p.p_uid)

    # attach to channels
    ch_in = Channel.attach(input_channel.sdesc)
    ch_out = Channel.attach(output_channel.sdesc)

    recv_handle = ch_in.recvh()
    send_handle = ch_out.sendh()

    recv_handle.open()
    send_handle.open()

    # start calculation
    random.seed(a=42)

    num_true = 0
    num_send = 0
    delta = 1
    count = 0

    while abs(delta) > 1e-6:

        # send new values to everyone
        for _ in range(num_workers):
            x = random.uniform(0.0, 1.0)
            y = random.uniform(0.0, 1.0)

            msg = pickle.dumps((x, y))

            send_handle.send_bytes(msg)  # easy but slow way to send 2 x 8 bytes
            num_send += 1

        # receive results from everyone
        for _ in range(num_workers):
            msg = recv_handle.recv_bytes()

            is_inside = pickle.loads(msg)

            if is_inside:
                num_true += 1

        # calculate result of this iteration
        value = 4.0 * float(num_true) / float(num_send)
        delta = (value - math.pi) / math.pi

        if count % 512 == 0:
            print(f"{count}: pi={value}, error={delta}")
            sys.stdout.flush()

        count += 1

    print(f"Final value after {count} iterations: pi={value}, error={delta}")

    # shut down all managed processes
    for _ in range(num_workers):
        msg = pickle.dumps((-1.0, -1.0))  # termination message
        send_handle.send_bytes(msg)

    recv_handle.close()
    send_handle.close()

    ch_in.detach()
    ch_out.detach()

    # join a set of processes at once; wait until all the processes are finished
    dgprocess.multi_join(processes, join_all=True)

    dgchan.destroy(input_channel.c_uid)
    dgchan.destroy(output_channel.c_uid)


def usage():
    print(f"USAGE: dragon pi_demo.py $N")
    print(f"  N : number of worker processes to start")


if __name__ == "__main__":

    if len(sys.argv) == 1:
        usage()
        sys.exit(f"")

    if sys.argv[1] == "sub":
        pi_helper(sys.argv[2], sys.argv[3])
    else:

        print(f"\npi-demo: Calculate π = 3.1415 ... in parallel using the Dragon native API.\n")

        try:
            num_workers = int(sys.argv[1])
        except:
            usage()
            sys.exit(f"Wrong argument '{sys.argv[1]}'")

        print(f"Got num_workers = {num_workers}")
        sys.stdout.flush()
        main(num_workers)

examples/dragon_gs_client/pi_demo.py
DragonHPC/dragon
"""Multiple producers multiple consumers communication scheme with Dragon
Channels.

* Create one Dragon Channel for the communication.
* Start up two processes that act as producers and two processes as consumers.
* Each producer generates a python object consisting of random data (a set of
    random strings) and a set of random functions to operate on the data (string
    operations such as upper()). Then, it uses the Channel to send the object as
    a message.
* Each consumer blocks until they receive a message on the channel. Then they apply
    the functions on the data they received and print the result in the output.
"""

import sys
import os
import random
import cloudpickle
import string

import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts
import dragon.globalservices.channel as dgchan
import dragon.globalservices.process as dgprocess
from dragon.infrastructure.process_desc import ProcessOptions

from dragon.channels import Channel, Many2ManyWritingChannelFile, Many2ManyReadingChannelFile


def producer(chan_name, funcs):
    data = ""
    for i in range(5):
        data = data + " " + "".join(random.choices(string.ascii_lowercase, k=i + 3))

    print(f'I am producer {os.getpid()} and I\'m sending data: "{data}" and string ops:', end=" ")

    channel_desc = dgchan.query(chan_name)
    _chan = Channel.attach(channel_desc.sdesc)

    n = random.randint(1, len(funcs))

    chosen = random.sample(list(funcs.items()), n)  # random selection without replacement

    for item in chosen:
        print(item[0], end=" ")
    print(flush=True)

    obj = (chosen, data)
    try:
        adapter = Many2ManyWritingChannelFile(_chan)
        adapter.open()
        cloudpickle.dump(obj, file=adapter, protocol=5)
    finally:
        adapter.close()


def consumer(chan_name, funcs):
    print(f"I am consumer {os.getpid()} --", end=" ")

    channel_desc = dgchan.query(chan_name)
    _chan = Channel.attach(channel_desc.sdesc)

    try:
        adapter = Many2ManyReadingChannelFile(_chan)
        adapter.open()
        funcs, data = cloudpickle.load(adapter)
    finally:
        adapter.close()

    for identifier, f in funcs:
        print(f'{identifier}: "{f(data)}"', end=" ")
        data = f(data)
    print(flush=True)


def main():
    # use default memory pool of this process
    default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)

    # ask GS to create a dragon channel
    _channel = dgchan.create(default_muid, user_name="InputChannel")

    # create processes
    cmd = sys.executable
    wdir = "."
    options = ProcessOptions(make_inf_channels=True)
    processes = []
    for _ in range(2):
        p = dgprocess.create(
            cmd, wdir, [os.path.basename(__file__), "sub", "producer", _channel.name], None, options=options
        )
        processes.append(p.p_uid)

    for _ in range(2):
        p = dgprocess.create(
            cmd, wdir, [os.path.basename(__file__), "sub", "consumer", _channel.name], None, options=options
        )
        processes.append(p.p_uid)

    # join a set of processes at once; wait until all the processes are finished
    dgprocess.multi_join(processes, join_all=True)

    dgchan.destroy(_channel.c_uid)


if __name__ == "__main__":
    # predefined functions for the producers to choose from
    # the dict is also used to verify the correctness of send/receive ops
    funcs = {
        "upper": lambda a: a.upper(),
        "lower": lambda a: a.lower(),
        "strip": lambda a: a.strip(),
        "capitalize": lambda a: a.capitalize(),
        'replace(" ", "")': lambda a: a.replace(" ", ""),
    }

    if len(sys.argv) > 1 and sys.argv[1] == "sub":
        if sys.argv[2] == "producer":
            producer(sys.argv[3], funcs)
        else:
            consumer(sys.argv[3], funcs)
    else:
        main()

examples/dragon_gs_client/queue_demo.py
DragonHPC/dragon
import os
import re
import logging
import queue

import dragon
import multiprocessing as mp

from dragon.channels import Channel
from dragon.infrastructure.connection import Connection
from dragon.globalservices import node
from dragon.globalservices import group
from dragon.globalservices import process
from dragon.globalservices import policy_eval
from dragon.infrastructure import process_desc
from dragon.utils import B64

logging.basicConfig(level=logging.INFO)


def producer_proc(producer_id: int, num_ranks: int, result_queue: mp.Queue) -> None:
    """
    The producer process will start `num_ranks` copies of the osu_alltoall
    MPI application. The stdout output from the head process will be piped via
    Dragon to a separate process responsible for parsing the osu_alltoall
    results and putting the results onto the results queue.

    :param producer_id: Numerical id identifying this producer
    :type producer_id: int
    :param num_ranks: The number of MPI ranks to start
    :type num_ranks: int
    :param result_queue: Handle to a queue where the results should be placed
    :type result_queue: mp.Queue
    """

    log = logging.getLogger(f"producer {producer_id}")
    log.info("Starting producer (num_ranks=%d)", num_ranks)

    exe = os.path.join(os.getcwd(), "osu_alltoall")
    args = ["--warmup", "10", "-m", "4096"]
    run_dir = os.getcwd()

    # Pipe the stdout output from the head process to a Dragon connection
    rank_that_prints_template = process.get_create_message(
        exe=exe,
        run_dir=run_dir,
        args=args,
        pmi_required=True,
        env=None,
        stdout=process.PIPE,
    )

    # All other ranks should have their output go to DEVNULL
    ranks_that_dont_print_template = process.get_create_message(
        exe=exe,
        run_dir=run_dir,
        args=args,
        pmi_required=True,
        env=None,
        stdout=process.DEVNULL,
    )

    # Establish the list and number of process ranks that should be started
    group_items = [
        (1, rank_that_prints_template.serialize()),
        (num_ranks - 1, ranks_that_dont_print_template.serialize()),
    ]

    # Ask Dragon to create the process group
    grp = group.create(items=group_items, policy=policy_eval.Policy(), soft=False)

    # Inpect the returned group descriptor
    group_puids = []
    parser_proc = None
    ranks_per_node = {n: 0 for n in range(len(node.get_list()))}
    for resources in grp.sets:
        group_puids.extend(
            [
                resource.desc.p_uid
                for resource in resources
                if resource.desc.state == process_desc.ProcessDescriptor.State.ACTIVE
            ]
        )

        for resource in resources:
            ranks_per_node[resource.placement] = ranks_per_node[resource.placement] + 1
            if resource.desc.stdout_sdesc:
                log.info(
                    "Starting parse_results process for p_uid=%d", resource.desc.p_uid
                )
                parser_proc = mp.Process(
                    target=parse_results_proc,
                    args=(producer_id, resource.desc.stdout_sdesc, result_queue),
                )
                parser_proc.start()

    log.info(", ".join([f"node {n} has {r} ranks" for n, r in ranks_per_node.items()]))
    log.info("Waiting for group to finish")
    if len(group_puids) > 0:
        process.multi_join(group_puids, join_all=True)

    if parser_proc:
        parser_proc.join()

    grp = group.destroy(grp.g_uid)

    log.info("Done")


def parse_results_proc(producer_id: int, stdout_sdesc: str, result_queue: mp.Queue) -> None:
    """
    Read stdout from the Dragon connection. Parse statistical data
    and put onto result queue.

    :param producer_id: Numerical id identifying this producer
    :type producer_id: int
    :param stdout_conn: Dragon serialized Channel descriptor used to read stdout data
    :type stdout_conn: str
    :param result_queue: Handle to a queue where the results should be placed
    :type result_queue: mp.Queue
    """

    log = logging.getLogger(f"parse_results {producer_id}")
    log.info("Parsing stdout from stdout connection")

    ch = Channel.attach(B64.from_str(stdout_sdesc).decode())
    conn = Connection(inbound_initializer=ch)

    try:
        result_matcher = re.compile(r"^(\d+)\s+([\d.]+)")
        while True:
            line = conn.recv()
            result = result_matcher.search(line)
            if result:
                result_queue.put(
                    {
                        producer_id: (
                            result[1],
                            result[2],
                        )
                    }
                )
    except EOFError:
        pass

    ch.detach()
    log.info("Done")


def consumer_proc(result_queue: mp.Queue, shutdown_event: mp.Event) -> None:
    """
    Read the values out of the result queue and
    just print them to the log

    :param result_queue: Handle to a queue where the results should be placed
    :type result_queue: mp.Queue
    :param shutdown_event: Event used to signal that the consumer process should exit
    :type shutdown_event: mp.Event
    """

    log = logging.getLogger("consumer")
    log.info("reading from result_queue")

    while not shutdown_event.is_set():
        try:
            values = result_queue.get(timeout=0.1)
            log.info(values)
        except queue.Empty:
            pass

    log.info("Done")


def main() -> None:
    mp.set_start_method("dragon")

    log = logging.getLogger("main")

    result_queue = mp.Queue()

    total_runs = 5
    current_runs = 0
    simultaneous_producers = 2
    producer_num = 0

    num_nodes = len(node.get_list())
    reserved_cores = (
        num_nodes * 2
    )  # Reserve a couple of cores for Dragon infrastructure
    num_real_cores = mp.cpu_count() // 2
    ranks_per_job = (num_real_cores - reserved_cores) // simultaneous_producers

    shutdown_event = mp.Event()
    log.info("Starting consumer process")
    consumer = mp.Process(
        target=consumer_proc,
        args=(
            result_queue,
            shutdown_event,
        ),
    )
    consumer.start()

    producers = set()
    active_producers = 0
    while current_runs < total_runs:
        while active_producers < min(simultaneous_producers, total_runs - current_runs):
            log.info("Starting a new producer")
            producer = mp.Process(
                target=producer_proc, args=(producer_num, ranks_per_job, result_queue)
            )
            producer.start()
            producers.add(producer.pid)
            active_producers += 1
            producer_num += 1

        exited_list, _ = process.multi_join(producers, join_all=False)
        log.info("at least one producer has exited")
        exited_puids = [] if exited_list is None else [puid for puid, _ in exited_list]
        current_runs = current_runs + len(exited_puids)
        active_producers = active_producers - len(exited_puids)
        producers -= set(exited_puids)

    log.info("Shutting down")
    shutdown_event.set()
    consumer.join()
    log.info("Done")


if __name__ == "__main__":
    main()

examples/dragon_native/mpi/hpc_workflow_demo.py
DragonHPC/dragon
import os
import re
import logging
import queue

import dragon
import multiprocessing as mp

from dragon.globalservices import node
from dragon.globalservices.process import multi_join
from dragon.infrastructure.connection import Connection
from dragon.native.process import MSG_PIPE, MSG_DEVNULL, Process, TemplateProcess
from dragon.native.process_group import ProcessGroup

logging.basicConfig(level=logging.INFO)


def producer_proc(producer_id: int, num_ranks: int, result_queue: mp.Queue) -> None:
    """
    The producer process will start `num_ranks` copies of the osu_alltoall
    MPI application. The stdout output from the head process will be piped via
    Dragon to a separate process responsible for parsing the osu_alltoall
    results and putting the results onto the results queue.

    :param producer_id: Numerical id identifying this producer
    :type producer_id: int
    :param num_ranks: The number of MPI ranks to start
    :type num_ranks: int
    :param result_queue: Handle to a queue where the results should be placed
    :type result_queue: mp.Queue
    """

    logging.basicConfig(level=logging.INFO)
    log = logging.getLogger(f"producer {producer_id}")
    log.info("Starting producer (num_ranks=%d)", num_ranks)

    exe = os.path.join(os.getcwd(), "osu_alltoall")
    args = ["--warmup", "10", "-m", "4096"]
    run_dir = os.getcwd()

    grp = ProcessGroup(restart=False, pmi_enabled=True)

    # Pipe the stdout output from the head process to a Dragon connection
    grp.add_process(
        nproc=1,
        template=TemplateProcess(target=exe, args=args, cwd=run_dir, stdout=MSG_PIPE)
    )

    # All other ranks should have their output go to DEVNULL
    grp.add_process(
        nproc=num_ranks-1,
        template=TemplateProcess(target=exe, args=args, cwd=run_dir, stdout=MSG_DEVNULL)
    )

    grp.init()
    grp.start()
    child_resources = [Process(None, ident=puid) for puid in grp.puids]
    parser_proc = None
    ranks_per_node = {n: 0 for n in range(len(node.get_list()))}
    for child_resource in child_resources:
        ranks_per_node[child_resource.node] = ranks_per_node[child_resource.node] + 1
        if child_resource.stdout_conn:
            log.info("Starting parse_results process for puid=%d", child_resource.puid)
            parser_proc = Process(
                target=parse_results_proc,
                args=(producer_id, child_resource.stdout_conn, result_queue),
            )
            parser_proc.start()

    log.info(", ".join([f"node {n} has {r} ranks" for n, r in ranks_per_node.items()]))
    log.info("Waiting for group to finish")
    if len(child_resources) > 0:
        grp.join()

    if parser_proc:
        parser_proc.join()

    grp.stop()

    log.info("Done")


def parse_results_proc(producer_id: int, stdout_conn: Connection, result_queue: mp.Queue) -> None:
    """
    Read stdout from the Dragon connection. Parse statistical data
    and put onto result queue.

    :param producer_id: Numerical id identifying this producer
    :type producer_id: int
    :param stdout_conn: Dragon Connection object to read stdout data from
    :type stdout_conn: Connection
    :param result_queue: Handle to a queue where the results should be placed
    :type result_queue: mp.Queue
    """

    logging.basicConfig(level=logging.INFO)
    log = logging.getLogger(f"parse_results {producer_id}")
    log.info("Parsing stdout from stdout connection")

    try:
        result_matcher = re.compile(r"^(\d+)\s+([\d.]+)")
        while True:
            line = stdout_conn.recv()
            result = result_matcher.search(line)
            if result:
                result_queue.put(
                    {
                        producer_id: (
                            result[1],
                            result[2],
                        )
                    }
                )
    except EOFError:
        pass

    log.info("Done")


def consumer_proc(result_queue: mp.Queue, shutdown_event: mp.Event) -> None:
    """
    Read the values out of the result queue and
    just print them to the log

    :param result_queue: Handle to a queue where the results should be placed
    :type result_queue: mp.Queue
    :param shutdown_event: Event used to signal that the consumer process should exit
    :type shutdown_event: mp.Event
    """

    logging.basicConfig(level=logging.INFO)
    log = logging.getLogger("consumer")
    log.info("reading from result_queue")

    while not shutdown_event.is_set():
        try:
            values = result_queue.get(timeout=0.1)
            log.info(values)
        except queue.Empty:
            pass

    log.info("Done")


def main() -> None:
    mp.set_start_method("dragon")

    log = logging.getLogger("main")

    result_queue = mp.Queue()

    total_runs = 5
    current_runs = 0
    simultaneous_producers = 2
    producer_num = 0

    num_nodes = len(node.get_list())
    reserved_cores = (
        num_nodes * 2
    )  # Reserve a couple of cores for Dragon infrastructure
    num_real_cores = mp.cpu_count() // 2
    ranks_per_job = (num_real_cores - reserved_cores) // simultaneous_producers

    shutdown_event = mp.Event()
    log.info("Starting consumer process")
    consumer = Process(
        target=consumer_proc,
        args=(
            result_queue,
            shutdown_event,
        ),
    )
    consumer.start()

    producers = set()
    active_producers = 0
    while current_runs < total_runs:
        while active_producers < min(simultaneous_producers, total_runs - current_runs):
            log.info("Starting a new producer")
            producer = Process(
                target=producer_proc, args=(producer_num, ranks_per_job, result_queue)
            )
            producer.start()
            producers.add(producer.puid)
            active_producers += 1
            producer_num += 1

        exited_list, _ = multi_join(producers, join_all=False)
        log.info("at least one producer has exited")
        exited_puids = [] if exited_list is None else [puid for puid, _ in exited_list]
        current_runs = current_runs + len(exited_puids)
        active_producers = active_producers - len(exited_puids)
        producers -= set(exited_puids)

    log.info("Shutting down")
    shutdown_event.set()
    consumer.join()
    log.info("Done")


if __name__ == "__main__":
    main()

examples/dragon_native/mpi/hpc_workflow_demo_highlevel.py
DragonHPC/dragon
""" 
Run a sample MPI Hello World application with 1 rank per allocated node.
"""

import sys
import os

from dragon.native.pool_workers import PoolWorkers, Worker, DragonPoolWorkersError
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.globalservices.node import get_list


def main():
    nnodes = len(get_list())
    mpi_hello_cmd = os.path.join(os.getcwd(), "mpi_hello")
    args = []
    cwd = os.getcwd()
    options = ProcessOptions(make_inf_channels=True)

    if not nnodes:
        print("No slurm allocation detected.")
        os.exit(-1)

    mpi_workers = [
        Worker(
            cmd=mpi_hello_cmd,
            args=args,
            cwd=cwd,
            processOptions=options,
            env=None,
            restartPolicy=Worker.RestartPolicy.NONE,
        )
        for _ in range(nnodes)
    ]

    pool = PoolWorkers(mpi_workers)
    pool.start()
    pool.join()
    pool.stop()

    return 0


if __name__ == "__main__":
    main()

examples/dragon_native/mpi/mpi_pool_workers_demo.py
DragonHPC/dragon
""" Calculate pi in parallel using the standard Monte Carlo technique 
    with the Dragon native interface. 

    NOTE: due to the missing Process object in Dragon native for version v0.3
    we are use the Multiprocessing process API here. This will change with
    v1.0
"""

import sys
import os
import random
import pickle
import math

import dragon
from dragon.native.queue import Queue

from multiprocessing import Process as Process  # only for Dragon v0.3


def pi_helper(q_in: Queue, q_out: Queue) -> None:

    while True:  # keep working

        x, y = q_in.get(timeout=None)

        if x == -1 and y == -1:  # we're done here
            break

        r = x * x + y * y

        if r > 1.0:
            q_out.put(False)
        else:
            q_out.put(True)


def main(num_workers):

    # create Dragon queues
    q_send = Queue()
    q_recv = Queue()

    # create & start processes
    processes = []
    for _ in range(num_workers):
        p = Process(target=pi_helper, args=(q_send, q_recv))
        p.start()
        processes.append(p)

    # start calculation
    random.seed(a=42)

    num_true = 0
    num_send = 0
    delta = 1
    count = 0

    while abs(delta) > 1e-6:

        # send new values to everyone
        for _ in range(num_workers):
            x = random.uniform(0.0, 1.0)
            y = random.uniform(0.0, 1.0)

            q_send.put((x, y), timeout=0)
            num_send += 1

        # receive results from everyone
        for _ in range(num_workers):

            is_inside = q_recv.get(timeout=None)

            if is_inside:
                num_true += 1

        # calculate result of this iteration
        value = 4.0 * float(num_true) / float(num_send)
        delta = (value - math.pi) / math.pi

        if count % 512 == 0:
            print(f"{count:04}: pi={value:10.7f}, error={delta:8.1e}", flush=True)

        count += 1

    print(f"Final value after {count} iterations: pi={value}, error={delta}")

    # shut down all managed processes
    for _ in range(num_workers):
        q_send.put((-1.0, -1.0), timeout=0)  # termination message

    # wait for all processes to be finished
    for p in processes:
        p.join(timeout=None)


def usage():
    print(f"USAGE: dragon pi_demo.py $N")
    print(f"  N : number of worker puids to start")


if __name__ == "__main__":

    print(f"\npi-demo: Calculate π = 3.1415 ... in parallel using the Dragon native API.\n")

    try:
        num_workers = int(sys.argv[1])
    except:
        usage()
        sys.exit(f"Wrong argument '{sys.argv[1]}'")

    print(f"Got num_workers = {num_workers}", flush=True)
    main(num_workers)

examples/dragon_native/pi_demo.py
DragonHPC/dragon
import socket

def hello(arg=None):
    print(f'Hello DRAGON from {socket.gethostname()}!!', flush=True)


if __name__ == "__main__":
    hello()

examples/jupyter/example.py
DragonHPC/dragon
import dragon
import multiprocessing as mp
import dragon.globalservices.process as gproc
import dragon.infrastructure.util as dutil
import example as ex

def main():
    mp.set_start_method('dragon')
    dutil.enable_logging()
    gproc.start_capturing_child_mp_output()

    proc = mp.Process(target=ex.hello, args=())
    proc.start()
    proc.join()
    gproc.stop_capturing_child_mp_output()

if __name__ == "__main__":
    main()


examples/jupyter/getout.py
DragonHPC/dragon
import scipy.signal
import dragon
import os
import socket
import multiprocessing as mp

def hello(arg=None):
    print(f'Hello DRAGON from {socket.gethostname()} with {mp.cpu_count()} cores!!', flush=True)

def f(args):
    image, random_filter = args
    # Do some image processing.
    return scipy.signal.convolve2d(image, random_filter)[::5, ::5]

if __name__ == "__main__":
    hello()
examples/jupyter/imageproc.py
DragonHPC/dragon
import dragon
import multiprocessing as mp
from dragon.globalservices.node import get_list, query
from transformers import BlenderbotTokenizer, BlenderbotForConditionalGeneration
import torch
from py3nvml.py3nvml import *
import time
import os
import queue
import socket


def get_device_count():
    """Uses py3nvml to get the number of gpus on a node.

    :return: the number of devices on a single node
    :rtype: int
    """
    if torch.cuda.is_available(): 
        nvmlInit()
        device_count = nvmlDeviceGetCount()
        nvmlShutdown()
    else: 
        device_count=0
    return device_count


def build_device_queues():
    """Builds a dictionary of device queues.

    :return: a dict of dragon queues
    :rtype: dict
    """
    node_dict = {}
    num_devices = get_device_count()
    print(f"num_devices = {num_devices}", flush=True)
    node_list = get_list()
    for node in node_list:
        node_dict[node] = mp.Queue()
    for node in node_list:
        for device in range(num_devices):
            node_dict[node].put(device)
    return node_dict


def get_huid():
    """Gets huid for a worker's node.

    :return: returns h_uid
    :rtype: int
    """
    name = os.uname().nodename
    desc = query(str(name))
    return desc.h_uid


def get_device(device_queue):
    # grabs an unoccupied device from the nodes unique queue
    huid = get_huid()
    if torch.cuda.is_available():
        try:
            available_cuda_device = device_queue[huid].get(timeout=10)
            gpu_available = True
        except queue.Empty:
            gpu_available = False
        if gpu_available:
            device = torch.device("cuda", available_cuda_device)
        else:
            # if we don't have a device that is free, we use the cpu
            device = torch.device("cpu")
    else:
        # if we don't have a device that is free, we use the cpu
        device = torch.device("cpu")
    return device


def start_model():
    """Initializes the model and tokenizer.

    :return: returns the model and tokenizer
    :rtype: tuple
    """
    mname = "facebook/blenderbot-400M-distill"
    model = BlenderbotForConditionalGeneration.from_pretrained(mname)
    tokenizer = BlenderbotTokenizer.from_pretrained(mname)
    return model, tokenizer


class ReadWorker:
    """Handles reading from the response queue and returning the response to the front end."""

    def __init__(self, q, end_ev):
        self.q = q
        self.end_ev = end_ev

    def read(self, idx):
        """The main function for returning responses to the prompter.
        Worker idx sits inside the while loop and continuously reads from a response queue. The worker
        exits when the queue is empty and the end_ev is set.

        :param idx: worker id
        :type idx: int
        """
        print(f"read_worker {idx} started", flush=True)
        # sit in this while loop and wait for work to read
        while True:
            try:
                # gets the prompt from the queue
                response_info = self.q.get(timeout=1)
                # parses response and prompt id pairs
                prompt = response_info[0]
                prompt_id = response_info[1]
                response = response_info[2]
                responder_id = response_info[3]
                # this print could be replaced by passing the response back to the prompt_id if there are many prompters
                print(
                    f" {prompt_id} asked: {prompt} \n worker {responder_id} responded: {response} \n read worker {idx} handled this response",
                    flush=True,
                )
            except queue.Empty:
                if self.end_ev.is_set():
                    # if the queue is empty and the end event is set then we shut down
                    print(f"Shutting down read worker {idx} ", flush=True)
                    break
                else:
                    time.sleep(1)
            except Exception as e:
                print(f"Exception caught: {e}", flush=True)


class InfWorker:
    """The main inference worker."""

    def __init__(self, q_in, q_out, device_queue, end_ev):
        self.q_in = q_in
        self.q_out = q_out
        self.end_ev = end_ev
        self.device_queue = device_queue

    def infer(self, idx):
        """Inference worker function. Worker idx gets a device, initializes the model, and places the model on
        that device. It then enters a while loop that continually checks a shared dragon queue that contains the prompt
        with the prompter's ID. It tokenizes the prompt, places the prompt on the device, and then generates a response.
        It places this response in the shared response queue. The worker exits once both the prompt queue is empty and
        the end event is set.

        :param idx: inference worker id
        :type idx: int
        """
        print(f"inf_worker {idx} started", flush=True)
        device = get_device(self.device_queue)
        # starting the model and tokenizer in the worker avoids pickling and keeps start up costs equal
        model, tokenizer = start_model()
        model.to(device)
        print(f" worker {idx} has device {device} on {socket.gethostname()}", flush=True)
        # sit in this while loop and wait for work
        while True:
            try:
                # gets the prompt from the queue
                prompt_id_pair = self.q_in.get(timeout=1)
                # parses id and prompt pair
                prompt = prompt_id_pair[0]
                id = prompt_id_pair[1]
                # generates reply from the model
                reply = self._respond(prompt, model, tokenizer, device)
                # removes some special characters.
                reply = str(reply[0]).strip("<s>")
                reply = reply.strip("</s>")
                self.q_out.put((prompt, id, reply, idx))
            except queue.Empty:
                if self.end_ev.is_set():
                    # if the queue is empty and the end event is set then we shut down
                    print(f"Shutting down inference worker {idx} ", flush=True)
                    break
                else:
                    time.sleep(1)
            except Exception as e:
                print(f"Exception caught: {e}", flush=True)

    def _respond(self, prompt, model, tokenizer, device):
        """Generates the response.

        :param prompt: input prompt
        :type prompt: str
        :param model: language model
        :type model: transformers.modeling_utils.PreTrainedModel
        :param tokenizer: tokenizer
        :type tokenizer: transformers.tokenization_utils_base.PreTrainedTokenizerBase
        :param device: device where response is generated
        :type device: torch.device
        :return: response
        :rtype: str
        """
        print("reading prompt", flush=True)
        input_ids = tokenizer([prompt], return_tensors="pt")
        input_ids = input_ids.to(device)
        output = model.generate(**input_ids, min_new_tokens=100, max_new_tokens=300)

        # Decode the generated text and return it
        reply_ids = tokenizer.batch_decode(output)
        return reply_ids


class InfWorkers:
    """Orchestrates the inference and response workers."""

    def __init__(self, num_workers, num_readers):
        self.num_workers = num_workers
        self.num_readers = num_readers
        # handles device allocation on node
        self.device_queue = build_device_queues()
        # event to signal we're done with the chat
        self.end_ev = mp.Event()
        # shared queue where we place prompts
        self.q_in = mp.Queue()
        self.q_out = mp.Queue()
        # pool and base inference worker that does inference work on prompt
        self.inf_pool = mp.Pool(self.num_workers)
        self.worker = InfWorker(self.q_in, self.q_out, self.device_queue, self.end_ev)
        # pool and base read worker that handles reading and eventually routing of replys
        self.read_pool = mp.Pool(self.num_readers)
        self.reader = ReadWorker(self.q_out, self.end_ev)

    def start(self):
        """Asychronously launches the inference and read workers."""
        print(f"starting inf pool with {self.num_workers} and read pool with {self.num_readers}", flush=True)
        # list to get reasonable placement of workers and keep track of who is responding
        worker_list = [i for i in range(self.num_workers)]
        # actually launch job. Use async so we return and can read prompts
        print(f"worker list = {worker_list}", flush=True)
        workers = self.inf_pool.map_async(self.worker.infer, worker_list, 1)
        reader_list = [i for i in range(self.num_readers)]
        # actually launch job. Use async so we return and can read prompts
        print(f"reader list = {reader_list}", flush=True)
        readers = self.read_pool.map_async(self.reader.read, reader_list, 1)

    def close(self):
        """Sets the end event and then closes and joines the inference and response worker pool."""
        # sets the end event and closes everything
        print(f"closing everything", flush=True)
        self.end_ev.set()
        self.inf_pool.close()
        self.inf_pool.join()
        self.read_pool.close()
        self.read_pool.join()
        self.q_in.close()
        self.q_out.close()

    def prompt(self, prompt):
        """Accepts a prompt and then couples that prompt with an id before placing in the shared input queue.

        :param prompt: user provided response
        :type prompt: str
        """
        # place prompt with h_uid in queue
        huid = get_huid()
        print(f"putting a prompt in the queue from {huid}", flush=True)
        self.q_in.put((prompt, huid))

examples/jupyter/llm_backend.py
DragonHPC/dragon
import dragon
import sys
import dragon.globalservices.process as gproc
from jupyter_core.command import main

def startup():
    gproc.start_capturing_child_mp_output()
    sys.argv = ["jupyter", "notebook",  "--ip", "0.0.0.0", "--no-browser", "--allow-root"]
    rc = main()
    sys.exit(rc)

if __name__ == '__main__':
    startup()

examples/jupyter/start_jupyter.py
DragonHPC/dragon
import dragon
import multiprocessing as mp
from dragon.globalservices.node import get_list
from py3nvml.py3nvml import *
import py3nvml.nvidia_smi as smi
from prometheus_client import Gauge, Histogram, start_http_server
import time
import os
import socket


class TelemWorkers:
    def __init__(self):
        self.procs = []
        self.end_ev = mp.Event()

    def start(self):
        """Launches a telemetry worker on each node"""

        node_list = get_list()
        nnodes = len(node_list)
        # Create a process on each node for monitoring
        for _ in range(nnodes):
            proc = mp.Process(target=self.telem_work, args=(self.end_ev,))
            proc.start()
            self.procs.append(proc)

    def close(self):
        """Stops the collection of data and closes the procs"""
        # wait on the monitoring processes
        self.end_ev.set()
        for proc in self.procs:
            proc.join()

    def telem_work(self, end_ev):
        """Updates a prometheus server with telemetry data from cpus and gpus on each node

        :param end_ev: the event used to signal the end of the telemetry data collection
        :type end_ev: mp.Event
        """
        print(f"This is a telemetry process on node {os.uname().nodename}.", flush=True)
        # Create Prometheus metrics
        gpu_utilization = Gauge(
            "gpu_utilization", "GPU utilization percentage", ["hostname", "gpu_index", "uuid"]
        )
        gpu_memory_utilization = Gauge(
            "gpu_memory_utilization", "GPU memory utilization percentage", ["hostname", "gpu_index", "uuid"]
        )
        gpu_memory_used = Gauge("gpu_memory_used", "GPU memory used ", ["hostname", "gpu_index", "uuid"])
        gpu_memory_free = Gauge("gpu_memory_free", "GPU memory free ", ["hostname", "gpu_index", "uuid"])
        gpu_memory_total = Gauge("gpu_memory_total", "GPU memory total ", ["hostname", "gpu_index", "uuid"])
        system_load_average = Gauge("system_load_average", "System load average over 1 minute")
        request_latency = Histogram("request_latency_seconds", "Request latency in seconds")

        # Start the Prometheus metrics server
        start_http_server(8000)

        while True:
            # TELEMETRY WITH PROMETHEUS
            # Initialize NVML
            nvmlInit()

            # Process requests and update metrics
            # Record the start time of the request
            start_time = time.time()

            # Get the system load averages
            load1, _, _ = os.getloadavg()

            # Update the system_load_average gauge with the new value
            system_load_average.set(load1)

            # Get the GPU utilization and memory utilization for each device
            device_count = nvmlDeviceGetCount()
            for i in range(device_count):
                handle = nvmlDeviceGetHandleByIndex(i)
                uuid = nvmlDeviceGetUUID(handle)
                utilization = nvmlDeviceGetUtilizationRates(handle)
                memory = nvmlDeviceGetMemoryInfo(handle)
                gpu_utilization.labels(socket.gethostname(), i, uuid).set(utilization.gpu)
                gpu_memory_utilization.labels(socket.gethostname(), i, uuid).set(utilization.memory)
                gpu_memory_used.labels(socket.gethostname(), i, uuid).set(memory.used >> 20)
                gpu_memory_free.labels(socket.gethostname(), i, uuid).set(memory.free >> 20)
                gpu_memory_total.labels(socket.gethostname(), i, uuid).set(memory.total >> 20)

            # Record the end time of the request and update the request_latency histogram
            end_time = time.time()
            request_latency.observe(end_time - start_time)

            # Shut down NVML
            nvmlShutdown()
            # END

            time.sleep(1)

            # check if the end event is set. If yes, exit.
            if end_ev.is_set():
                print(f"Telemetry process on node {os.uname().nodename} exiting ...", flush=True)
                break

examples/jupyter/telemetry.py
DragonHPC/dragon
import dragon

import argparse
import multiprocessing as mp
import time
import threading


def get_args():
    parser = argparse.ArgumentParser(description="All-all connection/queue test")

    parser.add_argument("--num_workers", type=int, default=4, help="number of workers in all-all")

    parser.add_argument("--iterations", type=int, default=100, help="number of iterations to do")

    parser.add_argument(
        "--burn_iterations", type=int, default=5, help="number of iterations to burn first time"
    )

    parser.add_argument(
        "--lg_message_size", type=int, default=4, help="log base 2 of msg size between each pair"
    )

    parser.add_argument("--dragon", action="store_true", help="run with dragon objs in dragon-style pattern")
    parser.add_argument("--dragon_compare", action="store_true", default=False, help="run with dragon objs and mp-native compatible patterns")
    parser.add_argument("--queues", action="store_true", help="measure aa with per-worker queues")
    my_args = parser.parse_args()

    assert my_args.lg_message_size <= 24, "arbitrary limit: 16MB message size"

    return my_args


def test_worker_conn(links, result_link, msg_size, total_iterations):
    """
    Worker function for both native multiprocessing and dragon
    Resources are generated by the initial process, which for dragon means poor multi-node scaling
    """
    # Output
    result_timings = []
    # Empty msg of size N bytes
    my_msg = bytearray(msg_size)

    def send_stuff():
        for link in links:
            link.send(my_msg)

    def recv_stuff():
        for link in links:
            _ = link.recv()

    for _ in range(total_iterations):
        start = time.perf_counter()

        send_th = threading.Thread(target=send_stuff)
        recv_th = threading.Thread(target=recv_stuff)

        send_th.start()
        recv_th.start()
        send_th.join()
        recv_th.join()

        result_timings.append(time.perf_counter() - start)

    result_link.send(result_timings)


# messages can mix between workers between iterations - there really
# ought to be a barrier - but the books balance at the end and this
# is meant to be more of a stress/comparison test anyway than
# a serious all-all.
def test_worker_queue(queues, my_idx, result_link, msg_size, total_iterations):
    """
    Worker function for both native multiprocessing and dragon
    Resources are generated by the initial process, which for dragon means poor multi-node scaling
    """
    result_timings = []
    my_msg = bytearray(msg_size)

    my_q = queues[my_idx]
    their_q = []
    for idx, q in enumerate(queues):
        if idx != my_idx:
            their_q.append(q)

    def send_stuff():
        for q in their_q:
            q.put(my_msg)

    def recv_stuff():
        for _ in range(len(their_q)):
            msg = my_q.get()
            assert msg_size == len(msg)

    for it_cnt in range(total_iterations):
        start = time.perf_counter()

        send_th = threading.Thread(target=send_stuff)
        recv_th = threading.Thread(target=recv_stuff)

        send_th.start()
        recv_th.start()
        send_th.join()
        recv_th.join()

        result_timings.append(time.perf_counter() - start)

    result_link.send(result_timings)



def test_worker_dragon_queue(queues, my_idx, result_link, msg_size, total_its):
    """
    Worker function for Dragon queues
    Native multiprocessing cannot share Queue objects between spawned processes, but Dragon can
    This offers better performance, as well as scalability to multi-node
    """
    # Create a per-process queue that Dragon can pass around
    # This provides significant on-node performance as well as multi-node performance
    my_q = mp.Queue(maxsize=len(queues)-1)
    my_msg = bytearray(msg_size)
    result_timings = []

    # Send process-made queue to all other workers
    for idx, q in enumerate(queues):
        if my_idx != idx:
            q.put(my_q)

    # Collect all the queues that have been sent to us
    their_q = []
    for i in range(len(queues)-1):
        q = queues[my_idx].get()
        their_q.append(q)

    del queues # Attempt to free up resource

    def send_stuff():
        for q in their_q:
            q.put(my_msg)

    def recv_stuff():
        for _ in range(len(their_q)):
            msg = my_q.get()
            assert msg_size == len(msg)

    for _ in range(total_its):
        start = time.perf_counter()

        send_th = threading.Thread(target=send_stuff)
        recv_th = threading.Thread(target=recv_stuff)

        send_th.start()
        recv_th.start()
        send_th.join()
        recv_th.join()

        result_timings.append(time.perf_counter() - start)

    result_link.send(result_timings)


def test_worker_dragon_conn(queues, my_idx, result_link, msg_size, total_iterations):
    """
    Worker function for both native multiprocessing and dragon, but written similar to dragon_queue
    This generates resources per process.
    Native and Dragon can both pass these resources between processes
    For Dragon this also enables multi-node scaling
    """
    my_msg = bytearray(msg_size)
    result_timings = []

    # Generate connection objects and pass them around
    # Queue is used here for this purpose, may be a less resource-hungry method to do so
    my_links = []
    for idx, q in enumerate(queues):
        if my_idx != idx:
            a, b = mp.Pipe()
            my_links.append(a)
            q.put(b)

    their_links = []
    for i in range(len(queues)-1):
        conn = queues[my_idx].get()
        their_links.append(conn)

    del queues # Attempt to free up resource

    def send_stuff():
        for conn in their_links:
            conn.send(my_msg)

    def recv_stuff():
        for conn in my_links:
            msg = conn.recv()
            assert msg_size == len(msg)


    for _ in range(total_iterations):
        start = time.perf_counter()

        send_th = threading.Thread(target=send_stuff)
        recv_th = threading.Thread(target=recv_stuff)

        send_th.start()
        recv_th.start()
        send_th.join()
        recv_th.join()

        result_timings.append(time.perf_counter() - start)

    result_link.send(result_timings)

def run_test(with_queues, nworkers, its, burns, msg_sz, is_dragon, comparison):
    # burns = warmup iterations
    total_iterations = burns + its

    # Return links for results of each process
    result_links = [mp.Pipe(duplex=False) for _ in range(nworkers)]

    queues = [mp.Queue(maxsize=nworkers) for _ in range(nworkers)]
    all_workers = []
    if with_queues:
        # Build list of queues for processes
        print(f"using {nworkers} Queues")

        # Build list of proccesses for all workers
        # Underlying implementation details mean Dragon can pass queue objects after spawn, but native cannot
        # Hence two slightly different functions
        if is_dragon and not comparison:
            all_workers = [
                mp.Process(
                    target=test_worker_dragon_queue,
                    args=(queues, i, result_links[i][1], msg_sz, total_iterations)
                )
                for i in range(nworkers)
            ]
        else:
            all_workers = [
                mp.Process(
                    target=test_worker_queue,
                    args=(queues, i, result_links[i][1], msg_sz, total_iterations)
                )
                for i in range(nworkers)
            ]
    else:
        # Dragon generates connection objections in spawned processes
        if is_dragon and not comparison:
            all_workers = [
                mp.Process(
                    target=test_worker_dragon_conn,
                    args=(queues, i, result_links[i][1], msg_sz, total_iterations)
                )
                for i in range(nworkers)
            ]

        else:
            # Generate all pipes in main process
            print(f"using {nworkers * (nworkers - 1) // 2} Pipes")
            all_links = {}

            # full duplex links; nworkers * (nworkers-1)/2 of them
            for i in range(nworkers):
                for j in range(i, nworkers):
                    all_links[i, j] = mp.Pipe()

            for i in range(nworkers):
                # first end if other's index less than yours,
                # second end if other's index greater than yours
                links = []
                for k in range(nworkers):
                    if k != i:
                        links.append(all_links[min(i, k), max(i, k)][int(k < i)])

                # Build a list of processes with an approrpriate start-end link and output link for results
                proc = mp.Process(
                    target=test_worker_conn, args=(links, result_links[i][1], msg_sz, total_iterations)
                )

                all_workers.append(proc)

    # Start all workers
    for worker in all_workers:
        worker.start()

    # Collect averages and totals
    avg_times = []
    total_times = []
    for i in range(nworkers):
        # Exclude "spin up" burn times from data
        result = result_links[i][0].recv()[burns:]
        total_times.append(sum(result))
        avg_times.append(total_times[i] / its)

    # Math
    avg_aa_completion_time = sum(avg_times) / nworkers  # seconds
    bandwidth = msg_sz * nworkers / avg_aa_completion_time  # bytes/second

    msg_size_in_k = msg_sz / (2**10)
    bw_in_mb_sec = bandwidth / (2**20)
    completion_time_in_ms = avg_aa_completion_time * 1000

    # Close all processes
    for worker in all_workers:
        worker.join()

    print(f"{nworkers=} {its=} {msg_size_in_k=} {completion_time_in_ms=:.3f} {bw_in_mb_sec=:.2f}")


def main():
    args = get_args()

    if args.dragon or args.dragon_compare:
        print("using dragon runtime")
        mp.set_start_method("dragon")
    else:
        # NOTE: Do not call with dragon launcher and using queues
        #       Results in an odd bug in v0.4
        print("using regular mp libs/calls")
        mp.set_start_method("spawn")

    run_test(
        with_queues=args.queues,
        nworkers=args.num_workers,
        msg_sz=2**args.lg_message_size,
        its=args.iterations,
        burns=args.burn_iterations,
        is_dragon=args.dragon,
        comparison=args.dragon_compare
    )


if __name__ == "__main__":
    exit(main())

examples/multiprocessing/aa_bench.py
DragonHPC/dragon
import dragon
import multiprocessing as mp
import heapq
import random

# In a real use of the MergePool class defined in this module, the sorter
# function below would be replaced by possibly many different functions that all
# produce ordered data that needs to be merged together into one sequence in an
# application. Think of many simulations that have a timecode associated with
# steps in the simulation. If you want to correlate the data by timecode, this
# MergePool class could be used and each of the functions producing ordered data
# would be different invocations of a simulation. Of course, there would be many
# more possible uses of merging sorted data into one stream.

def sorter(queue):
    # For this example, we'll generate some data, randomize the order, then sort it.
    # Presumably this would be the data that your function produces that you want to
    # combine, in sorted order, with data from other sources.
    data = list(range(10))
    random.shuffle(data)
    data.sort()

    # Then the process writes its sorted data to the provided queue to be merged with
    # other data from other sources
    for item in data:
        queue.put(item)

    # When this process's data is exhausted, we close the queue.
    queue.close()

# A SentinelQueue is a queue that raise EOFError when end of
# file is reached. It knows to do this by writing a sentinel
# into the queue when the queue is closed. A SentinelQueue
# contains a multiprocessing Queue. Other methods could
# be implemented for SentinelQueue, but are not needed in
# this example.
class SentinelQueue:
    _SENTINEL = 'sentinel_queue_sentinel'

    def __init__(self, maxsize=10):
        self.queue = mp.Queue(maxsize)

    def get(self):
        item = self.queue.get()
        if item == SentinelQueue._SENTINEL:
            raise EOFError('At EOF')

        return item

    def put(self, item):
        self.queue.put(item)

    def close(self):
        self.queue.put(SentinelQueue._SENTINEL)
        self.queue.close()

# The PQEntry class is needed by the priority
# queue which is used to always know which queue
# to get the next value from. The __lt__ orders
# the priority queue elements by their original
# values. But the queue index of where the value
# came from is carried along with the value in the
# priority queue so the merging algorithm knows where
# to get the next value. In this way, the total number
# of entries in the priority queue is never more than the
# fanin value of the MergePool.
class PQEntry:
    def __init__(self, value, queue_index):
        self._value = value
        self._queue_index = queue_index

    def __lt__(self, other):
        return self.value < other.value

    @property
    def queue_index(self):
        return self._queue_index

    @property
    def value(self):
        return self._value

    def __repr__(self):
        return 'PQEntry('+str(self.value)+','+str(self.queue_index)+')'

    def __str__(self):
        return repr(self)

# The MergePool is written so that start should be called to start
# a new MergePool object. Calling start returns a new process and a
# SentinelQueue from which to get the sorted values. The MergePool is
# given a list of functions which are the sources of sorted data that
# is to be merged together. This merging is done in a scalable fashion so
# this algorithm could be used to produce a sorted sequence of literally
# billions of items. The fanin_size can be set to a relatively high value
# since this will result in one priority queue with the number of elements
# equal to the fanin_size. The algorithm creates a tree of processes, each
# with a priority_queue of no more than fanin_size. The data is then merged
# up the tree until the original caller gets a sorted sequence of all the
# data from all the source functions.
class MergePool:
    def __init__(self):
        raise ValueError('Call MergePool.start to create a MergePool and run it.')

    # This is the hidden constructor that should only be called internally.
    def _init(sink_queue, source_functions, source_args=None, fanin_size=10, queue_depth=10):
        self = MergePool.__new__(MergePool)
        num_sources = len(source_functions)
        self._sources = []
        self._procs = []
        if num_sources > fanin_size:
            chunk_size = num_sources // fanin_size
            for i in range(0, num_sources, chunk_size):
                fun_chunk = source_functions[i:i+chunk_size]
                args_chunk = None if source_args is None else source_args[i:i+chunk_size]
                proc, sorted_queue = MergePool.start(fun_chunk, args_chunk, fanin_size, queue_depth)
                self._procs.append(proc)
                self._sources.append(sorted_queue)
        else:
            for i in range(len(source_functions)):
                sorted_queue = SentinelQueue(queue_depth)
                args = (sorted_queue,) if source_args is None else (sorted_queue,) + source_args[i]
                proc = mp.Process(target=source_functions[i], args=args)
                self._procs.append(proc)
                proc.start()
                self._sources.append(sorted_queue)

        self.sink_queue = sink_queue

        self.priority_queue = []

        self._prime_heap()

        return self

    def _prime_heap(self):
        for i in range(len(self._sources)):
            try:
                item = self._sources[i].get()
                heapq.heappush(self.priority_queue, PQEntry(item, i))
            except EOFError:
                pass

    def _start_merge_pool(sink_queue, source_functions, source_args, fanin_size, queue_depth):
        pool = MergePool._init(sink_queue, source_functions, source_args, fanin_size, queue_depth)
        pool._merge()

    def _merge(self):
        while len(self.priority_queue) > 0:
            item = heapq.heappop(self.priority_queue)
            self.sink_queue.put(item.value)

            try:
                next = self._sources[item.queue_index].get()
                heapq.heappush(self.priority_queue, PQEntry(next, item.queue_index))
            except EOFError:
                pass

        self.sink_queue.close()
        for proc in self._procs:
            proc.join()

    def start(source_functions, source_args=None, fanin_size=10, queue_depth=10):
        sink_queue = SentinelQueue(queue_depth)
        proc = mp.Process(target=MergePool._start_merge_pool,
                          args=(sink_queue, source_functions, source_args, fanin_size, queue_depth))
        proc.start()
        return (proc, sink_queue)

def main():
    mp.set_start_method('dragon')

    functions = [sorter] * 20
    proc, queue = MergePool.start(functions)

    try:
        while True:
            print(queue.get(), flush=True)
    except EOFError:
        pass

    proc.join()

if __name__ == "__main__":
    main()
examples/multiprocessing/distmerge.py
DragonHPC/dragon
#!/usr/bin/env python3
import dragon
import multiprocessing as mp
import time
import argparse


def numAcq(nlock, acquisitions):
    for _ in range(acquisitions):
        is_locked = nlock.acquire(block=True)
        if is_locked:
            nlock.release()


class LockPerformance:
    def workerAcquisition(self, min_workers, max_workers, worker_inc, acquisitions):

        avg_acq_time_num_workers = {}
        lock = mp.Lock()

        print(type(lock))
        print("Number of processes   Average acquire time(in usecs)")
        for num_workers in range(min_workers, max_workers + 1, worker_inc):

            processes = []

            start = time.monotonic()

            for _ in range(0, num_workers):
                p = mp.Process(target=numAcq, args=(lock, acquisitions))
                processes.append(p)
                p.start()

            for index in range(len(processes)):
                processes[index].join()

            end = time.monotonic()
            total_time = end - start
            avg_time_per_acquisition = (total_time / (num_workers * acquisitions)) * 1000000
            avg_acq_time_num_workers[num_workers] = avg_time_per_acquisition
            print(f"%19d%23f" % (num_workers, avg_time_per_acquisition))


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Dragon Lock Performance Tests")

    parser.add_argument("--dragon", action="store_true", help="run using dragon lock")
    parser.add_argument("--min_workers", type=int, default=300, help="workers for processes to do")
    parser.add_argument("--max_workers", type=int, default=400, help="workers for processes to do")
    parser.add_argument("--worker_inc", type=int, default=10, help="workers for processes to do")
    parser.add_argument("--acquisitions", type=int, default=1000, help="acquisitions per worker")

    args = parser.parse_args()

    if args.dragon:
        mp.set_start_method("dragon")

    else:
        mp.set_start_method("spawn")

    test = LockPerformance()
    test.workerAcquisition(
        min_workers=args.min_workers,
        max_workers=args.max_workers,
        worker_inc=args.worker_inc,
        acquisitions=args.acquisitions,
    )

examples/multiprocessing/lock_performance.py
DragonHPC/dragon
"""A simple recursive merge sort implementation using Python Multiprocessing.
It starts a process for every sublist until the sublist contains less than 
CUTOFF elements.
`__main__` actually executes (MAX_SIZE-MIN_SIZE)/INCREMENT merge sorts of
increasingly large sublists, measures the time and prints the timing results.

See e.g. Knuth, The Art of Computer Programming, 1998, Vol. 3, section 5.2.4.
"""

import random
import sys
import time
import math

import dragon
import multiprocessing as mp

# values are chosen so memory usage fits into the default Dragon memory pool of 4GB
CUTOFF = 20000
MIN_SIZE = 100000
MAX_SIZE = 1000000
INCREMENT = MIN_SIZE


def merge(left: list, right: list) -> list:
    """This function merges two lists.

    :param left: First list containing data
    :type left: list
    :param right: Second list containing data
    :type right: list
    :return: Merged data
    :rtype: list
    """

    merged_list = [None] * (len(left) + len(right))

    i = 0
    j = 0
    k = 0

    while i < len(left) and j < len(right):
        if left[i] < right[j]:
            merged_list[k] = left[i]
            i = i + 1
        else:
            merged_list[k] = right[j]
            j = j + 1
        k = k + 1

    # When we are done with the while loop above
    # it is either the case that i > midpoint or
    # that j > end but not both.

    # finish up copying over the 1st list if needed
    while i < len(left):
        merged_list[k] = left[i]
        i = i + 1
        k = k + 1

    # finish up copying over the 2nd list if needed
    while j < len(right):
        merged_list[k] = right[j]
        j = j + 1
        k = k + 1

    return merged_list


def parallel_merge_sort(chunk: list, cutoff: int, sorted_chunk_queue: object) -> None:
    """Recursive merge sort function. Below `cutoff` items, sort immediately.
    Otherwise, divide the input list `chunk` into two, start two processes
    executing the same function on one of the sublist. Wait for their result in
    the `result_queue`. Merge the two results and put the resulting list into
    `sorted_chunk_queue`.

    :param chunk: sub-list to recursively sort
    :type chunk: list
    :param cutoff: number of items below which the list will be sorted immediately
    :type cutoff: int
    :param sorted_chunk_queue: Queue to put the merged list into
    :type sorted_chunk_queue: mp.Queue object
    """

    if len(chunk) <= cutoff:
        chunk.sort()
        sorted_chunk_queue.put(chunk)

    else:
        midpoint = len(chunk) // 2

        left_chunk = chunk[:midpoint]
        right_chunk = chunk[midpoint:]

        result_queue = mp.Queue()

        left_proc = mp.Process(target=parallel_merge_sort, args=(left_chunk, cutoff, result_queue))
        right_proc = mp.Process(target=parallel_merge_sort, args=(right_chunk, cutoff, result_queue))

        left_proc.start()
        right_proc.start()

        result_a = result_queue.get(timeout=None)  # blocking
        right_b = result_queue.get(timeout=None)

        result = merge(result_a, right_b)

        sorted_chunk_queue.put(result)


def merge_sort(data: list, size: int, cutoff: int) -> int:
    """Kick off merge sort on a slice of `data` of size `size`,
    measure the runtime and return it.

    :param data: The whole input data
    :type data: list
    :param size: the size of the slice to sort
    :type size: int
    :param cutoff: when to stop recursing
    :type cutoff: int
    :return: runtime in seconds
    :rtype: int
    """

    the_lst = data[:size]

    start = time.perf_counter()

    result_queue = mp.Queue()
    parallel_merge_sort(the_lst, cutoff, result_queue)
    result = result_queue.get()

    the_lst.clear()
    the_lst.extend(result)

    end = time.perf_counter()
    delta = end - start

    return delta


def find_number_of_processes(n: int, cutoff: int) -> int:
    """Return the number of processes started by effectively
    replaying the recursion.

    :param n: number of elements
    :type n: int
    :param cutoff: umber of items below which no additional process is started
    :type cutoff: int
    :return: number of started processes
    :rtype: int
    """

    procs = 1

    while True:
        left = n // 2
        right = n - left
        procs = 2 * procs + 1

        if left <= cutoff:
            if right > cutoff:
                procs = procs + 2 * left

            return procs

        n = n // 2


if __name__ == "__main__":

    if "dragon" in sys.argv:
        mp.set_start_method("dragon")

    data: list = [random.randrange(MAX_SIZE) for i in range(MAX_SIZE)]

    print(
        f"    List Size    Time (seconds)    Processes    Channels (or Queues) with cutoff={CUTOFF}",
        flush=True,
    )

    for size in range(MIN_SIZE, MAX_SIZE + 1, INCREMENT):

        delta = merge_sort(data, size, CUTOFF)
        proc_count = find_number_of_processes(size, CUTOFF)
        channel_count = proc_count // 2
        print(f"{size:13d}    {delta:14.6f}{proc_count:12}{channel_count:12}")

examples/multiprocessing/merge_sort.py
DragonHPC/dragon
#!/usr/bin/env python3

from mpi4py import MPI
import time
import argparse
import sys


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="All-to-all communication test")

    parser.add_argument("--iterations", type=int, default=100, help="number of iterations to do")

    parser.add_argument(
        "--burn_iterations", type=int, default=5, help="number of iterations to burn first time"
    )

    parser.add_argument(
        "--lg_message_size", type=int, default=21, help="log base 2 of msg size between each pair"
    )

    parser.add_argument("--with_bytes", action="store_true", help="creates a payload using bytearray")

    my_args = parser.parse_args()

    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()

    burns = my_args.burn_iterations
    its = my_args.iterations
    workers = comm.Get_size()

    msg_size = 2**my_args.lg_message_size

    if my_args.with_bytes:
        my_msg = bytearray(msg_size)
    else:
        my_msg = []
        my_size = sys.getsizeof(my_msg)
        while my_size < msg_size:
            my_msg.append("a")
            my_size = sys.getsizeof(my_msg)

    senddata = workers * [my_msg]

    timer = 0.0
    for i in range(its + burns):
        start = time.perf_counter()
        recvdata = comm.alltoall(senddata)
        stop = time.perf_counter()

        if i >= burns:
            timer += stop - start

    timing = timer / its
    avg_time = comm.reduce(timing, op=MPI.SUM, root=0)

    if rank == 0:
        completion_time = avg_time / workers
        completion_time_in_ms = 1000 * completion_time
        msg_size_in_k = msg_size / (2**10)
        bandwidth = msg_size * workers / completion_time  # bytes/second
        bw_in_mb_sec = bandwidth / (2**20)
        print(f"{workers=} {its=} {msg_size_in_k=} {completion_time_in_ms=:.3f} {bw_in_mb_sec=:.2f}")

examples/multiprocessing/numpy-mpi4py-examples/mpi4py_aa_bench.py
DragonHPC/dragon
#!/usr/bin/env python3

from mpi4py import MPI
import time
import argparse
import sys

BURN_ITERS = 2

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="P2P latency test")

    parser.add_argument("--iterations", type=int, default=100, help="number of iterations to do")

    parser.add_argument(
        "--lg_max_message_size", type=int, default=4, help="log base 2 of size of message to pass in"
    )

    parser.add_argument("--with_bytes", action="store_true", help="creates a payload using bytearray")

    my_args = parser.parse_args()

    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()

    msg_size = 2
    if rank == 0:
        print(f"Msglen [B]   Lat [usec]")

    while msg_size <= 2**my_args.lg_max_message_size:
        start = 0

        if my_args.with_bytes:
            my_msg = bytearray(msg_size)
        else:
            my_msg = []
            my_size = sys.getsizeof(my_msg)
            while my_size < msg_size:
                my_msg.append("a")
                my_size = sys.getsizeof(my_msg)

        if rank == 0:
            for i in range(my_args.iterations + BURN_ITERS):
                if i == (BURN_ITERS + 1):
                    start = time.perf_counter()

                comm.send(my_msg, dest=1, tag=11)
                data = comm.recv(source=1, tag=12)

            avg_time = 1e6 * (time.perf_counter() - start) / my_args.iterations
            print(f"{msg_size}  {avg_time}")
        elif rank == 1:
            for i in range(my_args.iterations + BURN_ITERS):
                data = comm.recv(source=0, tag=11)
                comm.send(data, dest=0, tag=12)

        msg_size *= 2

examples/multiprocessing/numpy-mpi4py-examples/mpi4py_p2p_lat.py
DragonHPC/dragon
import dragon
import argparse
import multiprocessing
import numpy as np
import time


def get_args():
    parser = argparse.ArgumentParser(description="Basic NumPy test")

    parser.add_argument("--num_workers", type=int, default=4, help="number of workers")

    parser.add_argument("--iterations", type=int, default=100, help="number of iterations to do")

    parser.add_argument(
        "--burns", type=int, default=5, help="number of iterations to burn/ignore in order to warm up"
    )

    parser.add_argument("--dragon", action="store_true", help="run with dragon objs")

    parser.add_argument("--work_size", type=int, default=1, help="size of the array in thousands")

    parser.add_argument(
        "--work_items",
        type=int,
        default=100,
        help="number of work items, i.e., number of numpy arrays that the workers will work on",
    )

    parser.add_argument("--debug", action="store_true", default=False)

    my_args = parser.parse_args()

    return my_args


def g(size):
    y = np.ones(size)
    return np.exp(y)


if __name__ == "__main__":
    import sys

    args = get_args()

    if args.dragon:
        multiprocessing.set_start_method("dragon")
    else:
        multiprocessing.set_start_method("spawn")

    if args.debug:
        import logging

        multiprocessing.util.log_to_stderr(logging.DEBUG)

    print(f"Multiprocessing start method: {multiprocessing.get_start_method()}")

    num_cpus = args.num_workers
    ar_size = int(1e3 * args.work_size)

    with multiprocessing.Pool(num_cpus) as pool:
        print(f"Created pool: {type(pool)}", file=sys.stderr)

        all_iter = args.iterations + args.burns
        results = np.zeros(all_iter)
        for i in range(all_iter):
            print(f"trial {i}: ", end="", file=sys.stderr)
            start = time.perf_counter()
            pool.map(g, [ar_size] * args.work_items)
            finish = time.perf_counter()
            results[i] = finish - start
            print(results[i], file=sys.stderr)

        print(f"Average time: {round(np.mean(results[args.burns:]), 4)} second(s)", file=sys.stderr)
        print(f"Standard deviation: {round(np.std(results[args.burns:]), 4)} second(s)", file=sys.stderr)

examples/multiprocessing/numpy-mpi4py-examples/numpy_scale_work.py
DragonHPC/dragon
import dragon
import argparse
import multiprocessing
import numpy as np
import scipy.signal
import time
import itertools


def get_args():
    parser = argparse.ArgumentParser(description="Basic SciPy test")

    parser.add_argument("--num_workers", type=int, default=4,
                        help="number of workers")

    parser.add_argument("--iterations", type=int, default=10,
                        help="number of iterations to do")

    parser.add_argument("--burns", type=int, default=2,
                        help="number of iterations to burn/ignore in order to warm up")

    parser.add_argument("--dragon", action="store_true",
                        help="run with dragon objs")

    parser.add_argument("--size", type=int, default=1000,
                        help="size of the array")

    parser.add_argument("--mem", type=int, default=(512 * 1024 * 1024),
                        help="overall footprint of image dataset to process")

    parser.add_argument('--work_time', type=float, default=0.0,
                        help='how many seconds of compute per image')

    my_args = parser.parse_args()

    return my_args


def f(args):
    # Do some image processing
    image, random_filter, work_time = args
    elapsed = 0.
    start = time.perf_counter()

    # Explicitly control compute time per image
    while elapsed < work_time:
        scipy.signal.convolve2d(image, random_filter)[::5, ::5]
        elapsed = time.perf_counter() - start
    return scipy.signal.convolve2d(image, random_filter)[::5, ::5]


if __name__ == "__main__":
    args = get_args()

    # Start Dragon or base Multiprocessing
    if args.dragon:
        print("using dragon runtime")
        multiprocessing.set_start_method("dragon")
    else:
        print("using regular mp libs/calls")
        multiprocessing.set_start_method("spawn")

    # Image and filter generation
    image = np.zeros((args.size, args.size))
    nimages = int(float(args.mem) / float(image.size))
    print(f"Number of images: {nimages}", flush=True)
    images = []
    for j in range(nimages):
        images.append(np.zeros((args.size, args.size)))
    filters = [np.random.normal(size=(4, 4)) for _ in range(nimages)]

    num_cpus = args.num_workers
    print(f"Number of workers: {num_cpus}", flush=True)

    # Initialize the pool of workers
    pool = multiprocessing.Pool(num_cpus)

    # Main body of the computation
    times = []
    for i in range(args.iterations + args.burns):
        start = time.perf_counter()
        res = pool.map(f, zip(images, filters, itertools.repeat(float(args.work_time))))
        del res
        times.append(time.perf_counter() - start)
        print(f"Time for iteration {i} is {times[i]} seconds", flush=True)

    pool.close()
    pool.join()
    # Print out the mean and standard deviation, excluding the burns iterations
    print(f"Average time: {round(np.mean(times[args.burns:]), 2)} second(s)")
    print(f"Standard deviation: {round(np.std(times[args.burns:]), 2)} second(s)")

examples/multiprocessing/numpy-mpi4py-examples/scipy_scale_work.py
DragonHPC/dragon
#!/usr/bin/env python3

import dragon
import multiprocessing as mp
import time
import argparse

BURN_ITERS = 2
WINDOW = 32


def worker_conn(id, send_link, recv_link, result_link, msg_size, total_iterations, use_bytes):
    start = 0
    if id == 0:
        my_msg = bytearray(msg_size)
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            for j in range(WINDOW):
                if use_bytes:
                    send_link.send_bytes(my_msg)
                else:
                    send_link.send(my_msg)
            msg = recv_link.recv()

    else:
        my_msg = bytearray(8)
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            for j in range(WINDOW):
                if use_bytes:
                    msg = recv_link.recv_bytes()
                else:
                    msg = recv_link.recv()
            send_link.send(my_msg)

    avg_time = (time.perf_counter() - start) / total_iterations
    bw = 1e-6 * msg_size * WINDOW / avg_time
    result_link.send(bw)


def worker_queue(id, q, q_ack, result_link, msg_size, total_iterations):
    start = 0
    if id == 0:
        my_msg = bytearray(msg_size)
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            for j in range(WINDOW):
                q.put(my_msg)
            msg = q_ack.get()

    else:
        my_msg = bytearray(8)
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            for j in range(WINDOW):
                msg = q.get()
            q_ack.put(my_msg)

    avg_time = (time.perf_counter() - start) / total_iterations
    bw = 1e-6 * msg_size * WINDOW / avg_time
    result_link.send(bw)


def run_p2p_bw(iterations=100, max_msg_sz=1024, use_bytes=False, with_queues=False):

    result_links = [mp.Pipe(duplex=False), mp.Pipe(duplex=False)]

    if with_queues:
        q = mp.Queue(maxsize=WINDOW)
        q_ack = mp.Queue(maxsize=WINDOW)
    else:
        left_right_link = mp.Pipe(duplex=False)
        right_left_link = mp.Pipe(duplex=False)

    msg_sz = 2
    print(f"Msglen [B]   BW [MiB/s]")
    while msg_sz <= max_msg_sz:
        if with_queues:
            proc0 = mp.Process(
                target=worker_queue, args=(0, q, q_ack, result_links[0][1], msg_sz, iterations)
            )
            proc1 = mp.Process(
                target=worker_queue, args=(1, q, q_ack, result_links[1][1], msg_sz, iterations)
            )
        else:
            proc0 = mp.Process(
                target=worker_conn,
                args=(
                    0,
                    left_right_link[1],
                    right_left_link[0],
                    result_links[0][1],
                    msg_sz,
                    iterations,
                    use_bytes,
                ),
            )
            proc1 = mp.Process(
                target=worker_conn,
                args=(
                    1,
                    right_left_link[1],
                    left_right_link[0],
                    result_links[1][1],
                    msg_sz,
                    iterations,
                    use_bytes,
                ),
            )

        proc0.start()
        proc1.start()

        bw = result_links[0][0].recv()

        proc0.join()
        proc1.join()

        print(f"{msg_sz}  {bw}")

        msg_sz *= 2


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="P2P bandwidth test")

    parser.add_argument("--iterations", type=int, default=100, help="number of iterations to do")

    parser.add_argument(
        "--lg_max_message_size", type=int, default=4, help="log base 2 of size of message to pass in"
    )

    parser.add_argument("--dragon", action="store_true", help="run using dragon Connection")

    parser.add_argument(
        "--with_bytes", action="store_true", help="use send_bytes/recv_bytes instead of send/recv"
    )

    parser.add_argument("--queues", action="store_true", help="use queues for the communication")

    my_args = parser.parse_args()

    if my_args.dragon:
        print("using Dragon")
        mp.set_start_method("dragon")

    else:
        print("using multiprocessing")
        mp.set_start_method("spawn")

    run_p2p_bw(
        iterations=my_args.iterations,
        max_msg_sz=2**my_args.lg_max_message_size,
        use_bytes=my_args.with_bytes,
        with_queues=my_args.queues,
    )

examples/multiprocessing/p2p_bw.py
DragonHPC/dragon
#!/usr/bin/env python3

import dragon
import multiprocessing as mp
import time
import argparse

BURN_ITERS = 2


def worker_conn(id, send_link, recv_link, result_link, msg_size, total_iterations, use_bytes):

    my_msg = bytearray(msg_size)

    start = 0
    if id == 0:
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            if use_bytes:
                send_link.send_bytes(my_msg)
                msg = recv_link.recv_bytes()
            else:
                send_link.send(my_msg)
                msg = recv_link.recv()

    else:
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            if use_bytes:
                msg = recv_link.recv_bytes()
                send_link.send_bytes(my_msg)
            else:
                msg = recv_link.recv()
                send_link.send(my_msg)

    avg_time = (time.perf_counter() - start) / total_iterations
    result_link.send(avg_time)


def worker_queue(id, send_link, recv_link, result_link, msg_size, total_iterations):

    my_msg = bytearray(msg_size)

    start = 0
    if id == 0:
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            send_link.put(my_msg)
            msg = recv_link.get()

    else:
        for i in range(total_iterations + BURN_ITERS):
            if i == (BURN_ITERS + 1):
                start = time.perf_counter()

            msg = recv_link.get()
            send_link.put(my_msg)

    avg_time = (time.perf_counter() - start) / total_iterations
    result_link.send(avg_time)


def run_p2p_lat(iterations=100, max_msg_sz=1024, use_bytes=False, with_queues=False):

    result_links = [mp.Pipe(duplex=False), mp.Pipe(duplex=False)]

    if with_queues:
        q0 = mp.Queue(maxsize=2)
        q1 = mp.Queue(maxsize=2)
    else:
        left_right_link = mp.Pipe(duplex=False)
        right_left_link = mp.Pipe(duplex=False)

    msg_sz = 2
    print(f"Msglen [B]   Lat [usec]")
    while msg_sz <= max_msg_sz:
        if with_queues:
            proc0 = mp.Process(target=worker_queue, args=(0, q1, q0, result_links[0][1], msg_sz, iterations))
            proc1 = mp.Process(target=worker_queue, args=(1, q0, q1, result_links[1][1], msg_sz, iterations))
        else:
            proc0 = mp.Process(
                target=worker_conn,
                args=(
                    0,
                    left_right_link[1],
                    right_left_link[0],
                    result_links[0][1],
                    msg_sz,
                    iterations,
                    use_bytes,
                ),
            )
            proc1 = mp.Process(
                target=worker_conn,
                args=(
                    1,
                    right_left_link[1],
                    left_right_link[0],
                    result_links[1][1],
                    msg_sz,
                    iterations,
                    use_bytes,
                ),
            )

        proc0.start()
        proc1.start()

        time_avg = 0
        time_avg += result_links[0][0].recv()
        time_avg += result_links[1][0].recv()
        time_avg = 1e6 * time_avg / 2

        proc0.join()
        proc1.join()

        print(f"{msg_sz}  {time_avg}")

        msg_sz *= 2


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="P2P latency test")

    parser.add_argument("--iterations", type=int, default=1000, help="number of iterations to do")

    parser.add_argument(
        "--lg_max_message_size", type=int, default=4, help="log base 2 of size of message to pass in"
    )

    parser.add_argument("--dragon", action="store_true", help="run using dragon")

    parser.add_argument(
        "--with_bytes", action="store_true", help="use send_bytes/recv_bytes instead of send/recv"
    )

    parser.add_argument("--queues", action="store_true", help="use per-worker queues for the communication")

    my_args = parser.parse_args()

    if my_args.dragon:
        print("using Dragon")
        mp.set_start_method("dragon")

    else:
        print("using multiprocessing")
        mp.set_start_method("spawn")

    run_p2p_lat(
        iterations=my_args.iterations,
        max_msg_sz=2**my_args.lg_max_message_size,
        use_bytes=my_args.with_bytes,
        with_queues=my_args.queues,
    )

examples/multiprocessing/p2p_lat.py
DragonHPC/dragon
#!/usr/bin/env python3

import dragon
import multiprocessing as mp
import dragon.perf as dperf

if __name__ == "__main__":
    num_procs = 8
    with dperf.Session(num_procs) as session:
        kernel = session.new_kernel('small msg all-to-all')

        small_msg_size = 64
        large_msg_size = 1024 * 1024
        timeout_in_sec = 999

        # create an all-to-all kernel

        for src_ch_idx in range(num_procs):
            for dst_ch_idx in range(num_procs):
                kernel.append(dperf.Opcode.SEND_MSG, src_ch_idx, dst_ch_idx, small_msg_size, timeout_in_sec)

        for src_ch_idx in range(num_procs):
            for _ in range(num_procs):
                kernel.append(dperf.Opcode.GET_MSG, src_ch_idx, src_ch_idx, timeout_in_sec)

        # run the kernel

        kernel.run()


examples/multiprocessing/perf.py
DragonHPC/dragon
import dragon
import multiprocessing as mp

class PrimeNumberPipeline:
    SENTINEL = 0
    MAX_QUEUE_DEPTH = 10

    # The stage_size specifies how many prime numbers to keep in each stage of the pipeline.
    # The max_number is the largest number to count to.
    def __init__(self, stage_size=100, max_number=99999999):
        if stage_size <= 0:
            raise ValueError('Cannot create PrimeNumberPipeline with stage_size of 0 or less.')

        if max_number < 2:
            raise ValueError('Must provide a maximum number greater than 1 for a PrimeNumberPipeline')

        self._stage_size = stage_size
        self._max_number = max_number

    @staticmethod
    def source(limit, pipeline_queue):
        # This serves as the source for all the primes.
        # The values are written to the pipeline_queue to
        # serve as a source for the pipeline. We start at 3
        # because we prime the first pipeline stage with 2.
        for i in range(3, limit+1):
            pipeline_queue.put(i)

        # Used to signal the end of the pipeline.
        pipeline_queue.put(PrimeNumberPipeline.SENTINEL)

    @staticmethod
    def is_relatively_prime(number, primes):
        # Check that number is relatively prime to the primes.
        for prime in primes:
            # This is an optimization
            if number < prime*2:
                return True

            if number % prime == 0:
                return False

        return True

    @staticmethod
    def pipeline_stage(stage_size, prime, in_queue, out_queue):
        # A stage is given a prime from which to start. All other
        # primes it finds (up to the stage size) are added to this
        # list of primes to check.
        primes = [prime]
        stage_proc = None

        # This is the number of primes (relative to primes) that this
        # stage has found.
        cur_count = 1

        while True:
            number = in_queue.get()

            # We terminate via a sentinel value.
            if number == PrimeNumberPipeline.SENTINEL:
                out_queue.put(PrimeNumberPipeline.SENTINEL)
                if stage_proc is not None:
                    stage_proc.join()
                print(f'\nprimes in pipeline stage are {primes}', flush=True)
                return

            if PrimeNumberPipeline.is_relatively_prime(number, primes):
                # It is relatively prime, so send it to the output queue
                out_queue.put(number)

                # If it is found to be relatively prime, we add it to the
                # list of primes or we create a new stage with it (which
                # in turn builds its list of primes).
                if cur_count < stage_size:
                    primes.append(number)

                elif cur_count == stage_size:
                    # create a new pipeline stage
                    new_stage_queue = mp.Queue(maxsize=PrimeNumberPipeline.MAX_QUEUE_DEPTH)
                    stage_proc = mp.Process(target=PrimeNumberPipeline.pipeline_stage, args=(stage_size, number, new_stage_queue, out_queue))
                    out_queue = new_stage_queue
                    stage_proc.start()
                else:
                    # It was checked/will be checked by other stages in the pipeline
                    pass

                # Number of relatively prime primes found for this stage.
                cur_count+=1



    def start(self):
        self._source_queue = mp.Queue(maxsize=PrimeNumberPipeline.MAX_QUEUE_DEPTH)
        self._sink_queue = mp.Queue(maxsize=PrimeNumberPipeline.MAX_QUEUE_DEPTH)
        self._source_proc = mp.Process(target=PrimeNumberPipeline.source, args=(self._max_number, self._source_queue))
        self._source_proc.start()
        self._sink_queue.put(2) # "Prime" the sink queue ;)
        self._stage_proc = mp.Process(target=PrimeNumberPipeline.pipeline_stage, args=(self._stage_size, 2, self._source_queue, self._sink_queue))
        self._stage_proc.start()

    def stop(self):
        self._source_proc.join()
        self._stage_proc.join()

    def get(self):
        return self._sink_queue.get()


def main():
    # We set the start method to dragon to use dragon multiprocessing.
    mp.set_start_method('dragon')

    # We can control the pipeline by max_number and stage_size. The
    # stage_size can control the amount of parallelism.
    prime_pipeline = PrimeNumberPipeline(stage_size=3, max_number=100)
    prime_pipeline.start()
    print('Prime numbers:')
    count = 0
    for prime in iter(prime_pipeline.get, PrimeNumberPipeline.SENTINEL):
        text = f'{prime} '
        print(text, end="", flush=True)
        count+= len(text)
        if count > 80:
            count = 0
            print(flush=True)
    print(flush=True)
    prime_pipeline.stop()


if __name__ == "__main__":
    main()
examples/multiprocessing/prime_numbers.py
DragonHPC/dragon
"""Multiple producers multiple consumers communication scheme with
Multiprocessing and Dragon.

The code demonstrates the following key concepts working with Dragon:
* Serialization of data to pass it to another process. We use cloudpickle
  here, which is Python only. Multiprocessing with Dragon uses standard pickle,
  as default.
* Shared communication objects between processes, here a queue.
* Creating, starting and joining worker processes.
* Observe out-of-order execution and printing
"""

import random
import string
import cloudpickle

import dragon
import multiprocessing as mp


def producer(serialized_args: bytes) -> None:
    """Generate some string data, bundle it up with some random functions, add
    it to a queue.

    :param pickled_args: arguments to the function
    :type funcs: bytes
    """

    q, funcs = cloudpickle.loads(serialized_args)

    data = random.choices(string.ascii_lowercase, k=1)[0]
    for i in range(5):
        data = data + " " + "".join(random.choices(string.ascii_lowercase, k=i + 3))

    print(
        f'I am producer {mp.current_process().pid} and I\'m sending data: "{data}" and string ops:', end=" "
    )

    n = random.randint(1, len(funcs))
    chosen = random.sample(list(funcs.items()), n)  # random selection without replacement

    for item in chosen:
        print(item[0], end=" ")
    print(flush=True)

    work_pkg = cloudpickle.dumps((chosen, data))

    q.put(work_pkg)


def consumer(q: mp.queues.Queue) -> None:
    """Retrieve data from a queue, do some work on it and print the result.

    :param q: Queue to retrieve from
    :type q: mp.queues.Queue
    """

    # gives multi-node compatible Dragon puid, not OS pid.
    print(f"I am consumer {mp.current_process().pid} --", end=" ")

    serialized_data = q.get()  # implicit timeout=None here, blocking
    funcs, data = cloudpickle.loads(serialized_data)

    for identifier, f in funcs:
        print(f'{identifier}: "{f(data)}"', end=" ")
        data = f(data)
    print(flush=True)


if __name__ == "__main__":

    mp.set_start_method("dragon")

    # define some string transformations
    funcs = {
        "upper": lambda a: a.upper(),
        "lower": lambda a: a.lower(),
        "strip": lambda a: a.strip(),
        "capitalize": lambda a: a.capitalize(),
        'replace(" ", "")': lambda a: a.replace(" ", ""),
    }

    # use a queue for communication
    q = mp.Queue()

    # serialize producer arguments: Dragon uses pickle as default that doesn't
    # work with lambdas
    serialized_args = cloudpickle.dumps((q, funcs))

    # create & start processes
    processes = []
    for _ in range(8):
        p = mp.Process(target=producer, args=(serialized_args,))
        processes.append(p)
        p = mp.Process(target=consumer, args=(q,))
        processes.append(p)

    for p in processes:
        p.start()

    # wait for processes to finish
    for p in processes:
        p.join()

examples/multiprocessing/queue_demo.py
DragonHPC/dragon
"""SciPy Image Convolution Benchmark

This little benchmark can be run with `dragon` or standard `python3`
to compare performance on your machine. 

* Check which context to use - dragon or standard Multiprocessing
* Create pool of workers
* Create some random data
* Apply function `convolve_image_filter` to each of the image in parallel
* Measure time of every run, output mean and standard deviation
"""

import os
import time
import numpy as np

import dragon
import multiprocessing as mp
from scipy import signal

NUM_WORKERS = 4
IMAGE_SIZE = 1024
MEMSIZE = 32 * 1024**2  # 32 MB
NUM_ITER = 10
NUM_BURN = 2


def convolve_image_filter(args: tuple) -> np.ndarray:
    """Use scipy to convolve an image with a filter

    :param args: tuple containing image and filter
    :type args: tuple
    :return: convolved image
    :rtype: np.array
    """
    image, random_filter = args
    return signal.convolve2d(image, random_filter)[::5, ::5]


if __name__ == "__main__":
    
    # check if we should use the Dragon
    if "DRAGON_PATCH_MP" in os.environ:
        mp.set_start_method("dragon")
    else:
        mp.set_start_method("spawn")

    # create the worker pool
    pool = mp.Pool(NUM_WORKERS)

    # create images, stay within MEMSIZE
    image = np.zeros((IMAGE_SIZE, IMAGE_SIZE))
    num_images = int(float(MEMSIZE) / float(image.size))

    rng = np.random.default_rng(42)
    filters = [rng.standard_normal(size=(4, 4)) for _ in range(num_images)]
    images = [np.zeros((IMAGE_SIZE, IMAGE_SIZE)) for _ in range(num_images)]

    print(f"# of workers = {NUM_WORKERS}", flush=True)
    print(f"memory footprint = {MEMSIZE}", flush=True)
    print(f"# of images = {num_images}", flush=True)

    # run a small benchmark

    results = []
    for _ in range(NUM_BURN + NUM_ITER):
        beg = time.perf_counter()
        pool.map(convolve_image_filter, zip(images, filters))
        end = time.perf_counter()
        results.append(end - beg)

    results = results[NUM_BURN:]  # throw away burn in results
    
    pool.close()
    pool.join()
    print(f"Average execution time {round(np.mean(results), 2)} second(s)", flush=True)
    print(f"Standard deviation: {round(np.std(results), 2)} second(s)")

examples/multiprocessing/scipy_image_demo.py
DragonHPC/dragon
"""This contains all of the functions used for the SciPy convolution jobs.
"""
import scipy.signal
import numpy as np


# Generate the scipy data
def init_data(args):
    """Generates SciPy Image Data

    :param args: input args from command line
    :type args: ArgumentParser object
    :return: returns image and filter pairing
    :rtype: tuple
    """
    image = np.zeros((args.size, args.size))
    nimages = int(float(args.mem) / float(image.size))
    print(f"Number of images: {nimages}", flush=True)
    images = []
    for j in range(nimages):
        images.append(np.zeros((args.size, args.size)))
    filters = [np.random.normal(size=(4, 4)) for _ in range(nimages)]
    return zip(images, filters)


def f(args):
    """Use scipy to convolve an image with a filter

    :param args: tuple containing image and filter
    :type args: tuple
    :return: convolved image
    :rtype: np.array
    """
    image, random_filter = args
    # Do some image processing.
    return scipy.signal.convolve2d(image, random_filter)[::5, ::5]

examples/multiprocessing/torch-scipy-telemetry/conv.py
DragonHPC/dragon
"""This contains the functions used to run the mnist jobs and utilizes dragon queues to orchestrate GPU placement.
"""

import dragon
from dragon.globalservices.node import get_list, query_total_cpus, query
import argparse
import torch
import torch.multiprocessing as mp
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.optim.lr_scheduler import StepLR
import random
import os
import socket
import time
import telem as tm


class Net(nn.Module):
    """Convolutional neural network (two convolutional layers and two fully connected layers)
    based on the PyTorch neural network module. The RelU activation function adds nonlinearity
    and the max pool reduces the noise. The dropouts help reduce overfitting.
    """

    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, 3, 1)
        self.conv2 = nn.Conv2d(32, 64, 3, 1)
        self.dropout1 = nn.Dropout(0.25)
        self.dropout2 = nn.Dropout(0.5)
        self.fc1 = nn.Linear(9216, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        """Defines the computation done for a forward pass

        :param x: Input grayscaled image passed to the network
        :type x: torch.Tensor
        :return: Prediction
        :rtype: torch.Tensor
        """
        x = self.conv1(x)
        x = F.relu(x)
        x = self.conv2(x)
        x = F.relu(x)
        x = F.max_pool2d(x, 2)
        x = self.dropout1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        x = F.relu(x)
        x = self.dropout2(x)
        x = self.fc2(x)
        output = F.log_softmax(x, dim=1)
        return output


def train(model, device, train_loader, optimizer):
    """Trains the model on the specified device. It utilizes a
    PyTorch dataloader to iterate over the data

    :param model: Neural network model that defines layers and data flow
    :type model: mnist.Net
    :param device: PyTorch device to use for training
    :type device: torch.device
    :param train_loader: PyTorch data loader for training dataset
    :type train_loader: torch.utils.data.dataloader.DataLoader
    :param optimizer: PyTorch optimizer used to update model parameters
    :type optimizer: torch.optim
    """
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()


def test(model, device, test_loader):
    """Computes the accuracy and loss on the test data set

    :param model: Neural network model that defines layers and data flow
    :type model: mnist.Net
    :param device: PyTorch device to use for training
    :type device: torch.Device
    :param test_loader: PyTorch data loader for test dataset
    :type test_loader: torch.utils.data.dataloader.DataLoader
    :return: Test loass and test accuracy for current model parameters
    :rtype: tuple
    """
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.nll_loss(output, target, reduction="sum").item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    test_accuracy = 100.0 * correct / len(test_loader.dataset)
    return test_loss, test_accuracy


def buildDeviceQueue():
    """This constructs the dictinonary of device queues. It relies
    upon the ability to get the number of devices on a node and that
    the nodelist has unique node IDs. It assumes that the nodes
    all have the same number of GPUs per node.

    :return deviceQueue: a dictionary of multiprocessing queues that hold device numbers
    :rtype deviceQueue: dict[mp.queues.Queue]
    """
    nodeDict = {}
    numDevices = tm.getDeviceCount()
    nodelist = get_list()
    for node in nodelist:
        nodeDict[node] = mp.Queue()
    for node in nodelist:
        for device in range(numDevices):
            nodeDict[node].put(device)
    return nodeDict


def mnist_lr_sweep(args, deviceQueue, lr_p):
    """This trains on the MNIST dataset with variable learning rate
    and returns the final loss and accuracy. Each process gets a gpu
    that is determined by first getting a unique node identifier and
    then using that to access the queue of available devices.

    :param args: input args from command line
    :type args: ArgumentParser object
    :param deviceQueue: a dictionary of multiprocessing queues that hold device numbers
    :type deviceQueue: dict[mp.queues.Queue]
    :param lr_p: learning rate
    :type lr_p: float
    :return: learning rate, final test lost, and final test accuracy
    :rtype: list of floats
    """
    use_cuda = not args.no_cuda and torch.cuda.is_available()
    name = os.uname().nodename
    desc = query(str(name))
    # grabs an unoccupied device from the nodes unique queue
    available_cuda_device = deviceQueue[desc.h_uid].get()

    torch.manual_seed(args.seed)

    if use_cuda:
        device = torch.device("cuda", available_cuda_device)
    else:
        device = torch.device("cpu")

    train_kwargs = {"batch_size": args.batch_size}
    test_kwargs = {"batch_size": args.test_batch_size}
    if use_cuda:
        cuda_kwargs = {"num_workers": 0, "pin_memory": True, "shuffle": True}
        train_kwargs.update(cuda_kwargs)
        test_kwargs.update(cuda_kwargs)

    transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
    dataset1 = datasets.MNIST("../data", train=True, download=True, transform=transform)
    dataset2 = datasets.MNIST("../data", train=False, transform=transform)
    train_loader = torch.utils.data.DataLoader(dataset1, **train_kwargs)
    test_loader = torch.utils.data.DataLoader(dataset2, **test_kwargs)

    model = Net().to(device)
    optimizer = optim.Adadelta(model.parameters(), lr=lr_p)

    scheduler = StepLR(optimizer, step_size=1, gamma=args.gamma)

    for epoch in range(1, args.epochs + 1):
        train(model, device, train_loader, optimizer)
        test_loss, test_accuracy = test(model, device, test_loader)
        scheduler.step()

    # places device back in queue
    deviceQueue[desc.h_uid].put(available_cuda_device)
    return [lr_p, test_loss, test_accuracy]

examples/multiprocessing/torch-scipy-telemetry/mnist.py
DragonHPC/dragon
"""This file has all the functions used to gather telemetry data on each node. It relies heavily on py3nvml to gather this data. 
"""

from py3nvml.py3nvml import *
import py3nvml.nvidia_smi as smi
from tqdm import tqdm
import os
import sys
import socket
import time


def call_nvml():
    """Gathers telemetry data for Nvidia GPUs and the CPU.

    :return: telemetry information for all devices on a node
    :rtype: list of lists
    """
    nvmlInit()
    load1, load5, load15 = os.getloadavg()
    deviceCount = nvmlDeviceGetCount()
    gpu_info_list = []
    name = socket.gethostname()
    for i in range(deviceCount):
        handle = nvmlDeviceGetHandleByIndex(i)
        utilization = nvmlDeviceGetUtilizationRates(handle)
        uuid = nvmlDeviceGetUUID(handle)
        gpu_info_list.append([name, i, uuid, utilization.gpu, utilization.memory, load1])
    nvmlShutdown()
    return gpu_info_list


def getDeviceCount():
    """Uses py3nvml to get the number of Nvidia GPUs on a node

    :return: number of devices on a node
    :rtype: int
    """
    nvmlInit()
    deviceCount = nvmlDeviceGetCount()
    nvmlShutdown()
    return deviceCount


def buildTelemDict(hostname, tqdm_dict, deviceID=0):
    """Initializes tqdm bars so that they print in a somewhat nice fashion

    :param hostname: hostname of node
    :type hostname: str
    :param tqdm_dict: dictionary that is initially empty and populated with tqdm bars
    :type tqdm_dict: dict
    :param deviceID: device number that we want to print. If set to None then it will print all, defaults to 0
    :type deviceID: int, optional
    """
    offset = 0
    num_devices = 2
    if deviceID is None:
        num_gpus = getDeviceCount()
        num_devices = num_gpus
        num_devices += 1
    counter = len(tqdm_dict)
    tqdm_dict[hostname] = {}
    location = offset + counter * num_devices
    bar = tqdm(total=100, position=location, file=sys.stdout)
    bar.set_description_str(f"{hostname} cpu load avg.")
    tqdm_dict[hostname]["cpu"] = bar
    if deviceID is not None:
        location = offset + counter * num_devices + 1
        bar = tqdm(total=100, position=location, file=sys.stdout)
        bar.set_description_str(f"{hostname} device {deviceID} util")
        tqdm_dict[hostname]["device"] = bar
    else:
        num_gpus = getDeviceCount()
        num_devices = num_gpus
        num_devices += 1
        for gpu_device in range(num_gpus):
            location = offset + counter * num_devices + 1 + gpu_device
            bar = tqdm(total=100, position=location, file=sys.stdout)
            bar.set_description_str(f"{hostname} device {gpu_device} util")
            tqdm_dict[hostname]["device" + str(gpu_device)] = bar
    sys.stdout.flush()


def updateTelemDict(results_telem, tqdm_dict, deviceID=0):
    """Updates the information stored in the tqdm bars

    :param results_telem: telemetry information for all devices on a node
    :type results_telem: list of lists
    :param tqdm_dict: dictionary that is populated with tqdm bars
    :type tqdm_dict: dict
    :param deviceID: device number that we want to print. If set to None then it will print all, defaults to 0
    :type deviceID: int, optional
    """
    hostname = results_telem[0][0]
    try:
        tqdm_dict[hostname]["cpu"].n = results_telem[0][5]
    except:
        buildTelemDict(hostname, tqdm_dict, deviceID=deviceID)
        tqdm_dict[hostname]["cpu"].n = results_telem[0][5]
    tqdm_dict[hostname]["cpu"].refresh()
    for gpu_info in iter(results_telem):
        device_number = gpu_info[1]
        if deviceID is not None:
            if device_number != deviceID:
                continue
            else:
                tqdm_dict[hostname]["device"].n = gpu_info[3]
                tqdm_dict[hostname]["device"].refresh()
                sys.stdout.flush()
                break
        else:
            tqdm_dict[hostname]["device" + str(device_number)].n = gpu_info[3]
            tqdm_dict[hostname]["device" + str(device_number)].refresh()
            sys.stdout.flush()


def printTelem(results_telem, deviceID=0):
    """Prints telemetry data without tqdm bars

    :param results_telem: telemetry information for all devices on a node
    :type results_telem: list of lists
    :param deviceID: device number that we want to print. If set to None then it will print all, defaults to 0
    :type deviceID: int, optional
    """
    for gpu_info in iter(results_telem):
        hostname = gpu_info[0]
        device_number = gpu_info[1]
        gpu_uuid = gpu_info[2]
        gpu_util = gpu_info[3]
        gpu_memutil = gpu_info[4]
        load1 = gpu_info[5]
        if deviceID is not None:
            if device_number != deviceID:
                continue
            else:
                print(
                    "nodename: {} cpu load average 1 minute: {} device # {} utilization: {:.2f}%".format(
                        hostname, load1, device_number, gpu_util
                    ),
                    flush=True,
                )

examples/multiprocessing/torch-scipy-telemetry/telem.py
DragonHPC/dragon
"""This is the main file. It imports the other files and orchestrates the telemetry work. It contains telem_work, which
is the function launched on every node that gathers telemetry data and pushes it to a shared queue, and post_process, which is launched only
on one node and reads the telemetry data from the queue and then prints that information."""

import dragon
from dragon.globalservices.node import get_list, query_total_cpus
import multiprocessing as mp
import argparse
import time
import os
import queue

import functools
import mnist
import telem as tm
import conv


def get_args():
    """Get the user provided arguments
    :return args: input args from command line
    :rtype args: ArgumentParser object
    """
    parser = argparse.ArgumentParser(description="SciPy and MNIST test with telemetry")
    parser.add_argument("--scipy_workers", type=int, default=2, help="number of scipy workers (default: 2)")
    parser.add_argument("--mnist_workers", type=int, default=2, help="number of mnist workers (default: 2)")
    parser.add_argument(
        "--bars", action="store_true", default=False, help="uses tqdm bars to print telemetry info"
    )
    parser.add_argument("--no-cuda", action="store_true", default=False, help="disables CUDA training")
    parser.add_argument("--size", type=int, default=1024, help="size of the array (default: 1024)")
    parser.add_argument(
        "--mem",
        type=int,
        default=(1024 * 1024 * 1024),
        help="overall footprint of image dataset to process (default: 1024^3)",
    )
    parser.add_argument(
        "--batch-size", type=int, default=64, metavar="N", help="input batch size for training (default: 64)"
    )
    parser.add_argument(
        "--test-batch-size",
        type=int,
        default=1000,
        metavar="N",
        help="input batch size for testing (default: 1000)",
    )
    parser.add_argument(
        "--epochs", type=int, default=14, metavar="N", help="number of epochs to train (default: 14)"
    )
    parser.add_argument(
        "--gamma", type=float, default=0.7, metavar="M", help="Learning rate step gamma (default: 0.7)"
    )
    parser.add_argument("--seed", type=int, default=1, metavar="S", help="random seed (default: 1)")
    args = parser.parse_args()

    my_args = parser.parse_args()
    return my_args


def telem_work(q, start_ev, end_ev):
    """This is used by every monitoring process. It gathers telemetry data
    for CPU and GPU utilization and pushed it into the shared queue.

    :param q: shared queue that stores the telemetry data for each node
    :type q: Dragon Multiprocessing Queue
    :param start_ev: event that signals the beginning of monitoring
    :type start_ev: Event
    :param end_ev: event that signals the end of monitoring
    :type end_ev: Event
    """
    print(f"This is a telemetry process on node {os.uname().nodename}.", flush=True)
    start_ev.wait()  # wait until the starting event is set
    while True:
        gpu_info_list = tm.call_nvml()
        # one process on each node adds to a shared queue
        q.put(gpu_info_list)
        time.sleep(1)

        # check if the end event is set. If yes, exit.
        if end_ev.is_set():
            print(f"Telemetry process on node {os.uname().nodename} exiting ...", flush=True)
            break


def post_process(q, start_ev, end_ev, tqdm_bars):
    """This is used by the single post-processing process
    that gets the telemetry data from the shared queue and
    prints it.

    :param q: shared queue that stores the telemetry data for each node
    :type q: Dragon Multiprocessing Queue
    :param start_ev: event that signals the beginning of monitoring
    :type start_ev: Event
    :param end_ev: event that signals the end of monitoring
    :type end_ev: Event
    :param tqdm_bars: flag that signals whether to use bars or not for the presentation of the telemetry data
    :type tqdm_bars: Boolean
    """
    print(f"This is the postprocessing process, {os.uname().nodename}.", flush=True)
    start_ev.wait()  # wait until the starting event is set
    tqdm_dict = {}  # used when bars are used for the presentation of the telemetry data
    while True:
        # single process reads from the shared queue and prints results
        try:
            results_telem = q.get(timeout=10)
            if tqdm_bars:
                tm.updateTelemDict(results_telem, tqdm_dict, deviceID=None)
            else:
                tm.printTelem(results_telem)
        # when the queue is empty, exit
        except queue.Empty:
            print("Post process is exiting", flush=True)
            break
        except Exception as e:
            print(f"Exception caught: {e}", flush=True)


if __name__ == "__main__":
    args = get_args()
    print(f"Hello from main process {os.uname().nodename}.", flush=True)
    print("using dragon runtime", flush=True)
    mp.set_start_method("dragon")

    # get the list of nodes from Global Services
    nodeslist = get_list()
    nnodes = len(nodeslist)

    num_mnist_workers = args.mnist_workers
    assert num_mnist_workers > 1
    num_cpus = args.scipy_workers
    print(f"Number of nodes: {nnodes}", flush=True)
    print(f"Number of scipy workers: {num_cpus}", flush=True)
    print(f"Number of MNIST workers: {num_mnist_workers}", flush=True)

    # variable used to signal whether to use bars for the presentation of data or not
    use_bars = args.bars

    # Initialize the shared queue among the nodes
    # that is used for the communication of the telemetry data
    q = mp.Queue()

    # event used to signal the beginning of monitoring processes
    start_ev = mp.Event()
    # event used to signal the end of monitoring processes
    end_ev = mp.Event()

    # Create a process that gets and processes the telemetry data
    post_proc = mp.Process(target=post_process, args=(q, start_ev, end_ev, use_bars))
    post_proc.start()

    # Create a process on each node for monitoring
    procs = []
    for _ in range(nnodes):
        proc = mp.Process(target=telem_work, args=(q, start_ev, end_ev))
        proc.start()
        procs.append(proc)

    # Create a pool of workers for the scipy work
    scipy_data = conv.init_data(args)
    scipy_pool = mp.Pool(num_cpus)

    # Create a pool of workers for the mnist work
    deviceQueue = mnist.buildDeviceQueue()
    lr_list = [1 / (num_mnist_workers - 1) * i + 0.5 for i in range(num_mnist_workers)]
    mnist_lr_sweep_partial = functools.partial(mnist.mnist_lr_sweep, args, deviceQueue)
    mnist_pool = mp.Pool(num_mnist_workers)

    # start telemetry
    start_ev.set()

    # launch scipy and mnist jobs
    print(f"Launching scipy and mnist jobs", flush=True)
    workers_mnist = mnist_pool.map_async(mnist_lr_sweep_partial, lr_list, 1)
    workers_scipy = scipy_pool.map_async(conv.f, scipy_data)

    # wait on async processes
    mnist_pool.close()
    mnist_pool.join()
    scipy_pool.close()
    scipy_pool.join()

    # set the event to signal the end of computation
    time.sleep(10)
    print(f"Shutting down procs", flush=True)
    end_ev.set()

    # wait on the monitoring processes and the post-processing process
    for proc in procs:
        proc.join()
    post_proc.join()
    q.close()

    for result in workers_mnist.get():
        print(
            f"Final test for learning rate {result[0]}: loss: {result[1]} accuracy: {result[2]}", flush=True
        )

examples/multiprocessing/torch-scipy-telemetry/telemetry_full.py
DragonHPC/dragon
"""Common imports, functions and constants shared among all tests.
"""

import unittest
import os
import sys
import time
import operator
import gc

import test.support
from test import support

import threading

# Dragon must be imported before multiprocessing
import dragon

import multiprocessing

from multiprocessing import util
from multiprocessing.connection import wait

# Skip tests if _multiprocessing wasn't built.
_multiprocessing = test.support.import_module("_multiprocessing")

# Skip tests if sem_open implementation is broken.
support.skip_if_broken_multiprocessing_synchronize()


def latin(s):
    return s.encode("latin")


def close_queue(queue):
    if isinstance(queue, multiprocessing.queues.Queue):
        queue.close()
        queue.join_thread()


def join_process(process):
    """Since multiprocessing.Process has the same API than threading.Thread
    (join() and is_alive(), the support function can be reused
    """
    support.join_thread(process)


#
# Constants
#

LOG_LEVEL = util.SUBWARNING
# LOG_LEVEL = logging.DEBUG

DELTA = 0.1
CHECK_TIMINGS = False  # making true makes tests take a lot longer
# and can sometimes cause some non-serious
# failures because some calls block a bit
# longer than expected
if CHECK_TIMINGS:
    TIMEOUT1, TIMEOUT2, TIMEOUT3 = 0.82, 0.35, 1.4
else:
    TIMEOUT1, TIMEOUT2, TIMEOUT3 = 0.1, 0.1, 0.1

HAVE_GETVALUE = not getattr(_multiprocessing, "HAVE_BROKEN_SEM_GETVALUE", False)

WIN32 = sys.platform == "win32"

# To speed up tests when using the forkserver, we can preload these:
PRELOAD = ["__main__", "test.test_multiprocessing_forkserver"]


def wait_for_handle(handle, timeout):
    if timeout is not None and timeout < 0.0:
        timeout = None
    return wait([handle], timeout)


try:
    MAXFD = os.sysconf("SC_OPEN_MAX")
except:
    MAXFD = 256


def check_enough_semaphores():
    """Check that the system supports enough semaphores to run the test."""
    # minimum number of semaphores available according to POSIX
    nsems_min = 256
    try:
        nsems = os.sysconf("SC_SEM_NSEMS_MAX")
    except (AttributeError, ValueError):
        # sysconf not available or setting not available
        return
    if nsems == -1 or nsems >= nsems_min:
        return
    raise unittest.SkipTest(
        "The OS doesn't support enough semaphores " "to run the test (required: %d)." % nsems_min
    )


#
# Creates a wrapper for a function which records the time it takes to finish
#


class TimingWrapper(object):
    def __init__(self, func):
        self.func = func
        self.elapsed = None

    def __call__(self, *args, **kwds):
        t = time.monotonic()
        try:
            return self.func(*args, **kwds)
        finally:
            self.elapsed = time.monotonic() - t


#
# Base class for test cases
#


class BaseTestCase(object):
    """Base class for test cases"""

    ALLOWED_TYPES = ("processes", "manager", "threads")

    def assertTimingAlmostEqual(self, a, b):
        if CHECK_TIMINGS:
            self.assertAlmostEqual(a, b, 1)

    def assertReturnsIfImplemented(self, value, func, *args):
        try:
            res = func(*args)
        except NotImplementedError:
            pass
        else:
            return self.assertEqual(value, res)


# DRAGON I am commenting this out, as we are clearly not on Windows

# For the sanity of Windows users, rather than crashing or freezing in
# multiple ways.
#    def __reduce__(self, *args):
#        raise NotImplementedError("shouldn't try to pickle a test case")

#    __reduce_ex__ = __reduce__


#
# Return the value of a semaphore
#
def get_value(self):
    try:
        return self.get_value()
    except AttributeError:
        try:
            return self._Semaphore__value
        except AttributeError:
            try:
                return self._value
            except AttributeError:
                raise NotImplementedError


#
# Mixins
#


class BaseMixin(object):
    """Base class to inherit the correct part of the multiprocessing API"""

    @classmethod
    def setUpClass(cls):
        cls.dangling = (multiprocessing.process._dangling.copy(), threading._dangling.copy())

    @classmethod
    def tearDownClass(cls):
        # bpo-26762: Some multiprocessing objects like Pool create reference
        # cycles. Trigger a garbage collection to break these cycles.
        test.support.gc_collect()

        processes = set(multiprocessing.process._dangling) - set(cls.dangling[0])
        if processes:
            test.support.environment_altered = True
            support.print_warning(f"Dangling processes: {processes}")
        processes = None

        threads = set(threading._dangling) - set(cls.dangling[1])
        if threads:
            test.support.environment_altered = True
            support.print_warning(f"Dangling threads: {threads}")
        threads = None


class ProcessesMixin(BaseMixin):
    """Process based tests"""

    TYPE = "processes"
    Process = multiprocessing.Process
    connection = multiprocessing.connection
    current_process = staticmethod(multiprocessing.current_process)
    parent_process = staticmethod(multiprocessing.parent_process)
    active_children = staticmethod(multiprocessing.active_children)
    Pool = staticmethod(multiprocessing.Pool)
    Pipe = staticmethod(multiprocessing.Pipe)
    Queue = staticmethod(multiprocessing.Queue)
    JoinableQueue = staticmethod(multiprocessing.JoinableQueue)
    Lock = staticmethod(multiprocessing.Lock)
    RLock = staticmethod(multiprocessing.RLock)
    Semaphore = staticmethod(multiprocessing.Semaphore)
    BoundedSemaphore = staticmethod(multiprocessing.BoundedSemaphore)
    Condition = staticmethod(multiprocessing.Condition)
    Event = staticmethod(multiprocessing.Event)
    Barrier = staticmethod(multiprocessing.Barrier)
    Value = staticmethod(multiprocessing.Value)
    Array = staticmethod(multiprocessing.Array)
    RawValue = staticmethod(multiprocessing.RawValue)
    RawArray = staticmethod(multiprocessing.RawArray)


class ManagerMixin(BaseMixin):
    TYPE = "manager"
    Process = multiprocessing.Process
    Queue = property(operator.attrgetter("manager.Queue"))
    JoinableQueue = property(operator.attrgetter("manager.JoinableQueue"))
    Lock = property(operator.attrgetter("manager.Lock"))
    RLock = property(operator.attrgetter("manager.RLock"))
    Semaphore = property(operator.attrgetter("manager.Semaphore"))
    BoundedSemaphore = property(operator.attrgetter("manager.BoundedSemaphore"))
    Condition = property(operator.attrgetter("manager.Condition"))
    Event = property(operator.attrgetter("manager.Event"))
    Barrier = property(operator.attrgetter("manager.Barrier"))
    Value = property(operator.attrgetter("manager.Value"))
    Array = property(operator.attrgetter("manager.Array"))
    list = property(operator.attrgetter("manager.list"))
    dict = property(operator.attrgetter("manager.dict"))
    Namespace = property(operator.attrgetter("manager.Namespace"))

    @classmethod
    def Pool(cls, *args, **kwds):
        return cls.manager.Pool(*args, **kwds)

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.manager = multiprocessing.Manager()

    @classmethod
    def tearDownClass(cls):
        # only the manager process should be returned by active_children()
        # but this can take a bit on slow machines, so wait a few seconds
        # if there are other children too (see #17395)
        start_time = time.monotonic()
        t = 0.01
        while len(multiprocessing.active_children()) > 1:
            time.sleep(t)
            t *= 2
            dt = time.monotonic() - start_time
            if dt >= 5.0:
                test.support.environment_altered = True
                support.print_warning(
                    f"multiprocessing.Manager still has "
                    f"{multiprocessing.active_children()} "
                    f"active children after {dt} seconds"
                )
                break

        gc.collect()  # do garbage collection
        if cls.manager._number_of_objects() != 0:
            # This is not really an error since some tests do not
            # ensure that all processes which hold a reference to a
            # managed object have been joined.
            test.support.environment_altered = True
            support.print_warning("Shared objects which still exist " "at manager shutdown:")
            support.print_warning(cls.manager._debug_info())
        cls.manager.shutdown()
        cls.manager.join()
        cls.manager = None

        super().tearDownClass()


class ThreadsMixin(BaseMixin):
    TYPE = "threads"


# DRAGON need to comment these out, as the dummy class is not present in our current multiprocessing implementation
#     Process = multiprocessing.dummy.Process
#     connection = multiprocessing.dummy.connection
#     current_process = staticmethod(multiprocessing.dummy.current_process)
#     active_children = staticmethod(multiprocessing.dummy.active_children)
#     Pool = staticmethod(multiprocessing.dummy.Pool)
#     Pipe = staticmethod(multiprocessing.dummy.Pipe)
#     Queue = staticmethod(multiprocessing.dummy.Queue)
#     JoinableQueue = staticmethod(multiprocessing.dummy.JoinableQueue)
#     Lock = staticmethod(multiprocessing.dummy.Lock)
#     RLock = staticmethod(multiprocessing.dummy.RLock)
#     Semaphore = staticmethod(multiprocessing.dummy.Semaphore)
#     BoundedSemaphore = staticmethod(multiprocessing.dummy.BoundedSemaphore)
#     Condition = staticmethod(multiprocessing.dummy.Condition)
#     Event = staticmethod(multiprocessing.dummy.Event)
#     Barrier = staticmethod(multiprocessing.dummy.Barrier)
#     Value = staticmethod(multiprocessing.dummy.Value)
#     Array = staticmethod(multiprocessing.dummy.Array)

# DRAGON We set start_method, dangling and old_start_method explicitely,
# to avoid programmatic module creation as done in the standard multiprocessing tests
start_method = "dragon"

dangling = [None, None]
old_start_method = [None]

# DRAGON this was the init code in setUpModule()
def setUpModule():
    """Prepare env for test"""

    multiprocessing.set_forkserver_preload(PRELOAD)
    multiprocessing.process._cleanup()
    dangling[0] = multiprocessing.process._dangling.copy()
    dangling[1] = threading._dangling.copy()
    old_start_method[0] = multiprocessing.get_start_method(allow_none=True)
    try:
        multiprocessing.set_start_method(start_method, force=True)
    except ValueError:
        raise unittest.SkipTest(start_method + " start method not supported")

    if sys.platform.startswith("linux"):
        try:
            lock = multiprocessing.RLock()
        except OSError:
            raise unittest.SkipTest("OSError raises on RLock creation, " "see issue 3111!")
    check_enough_semaphores()
    util.get_temp_dir()  # creates temp directory
    multiprocessing.get_logger().setLevel(LOG_LEVEL)


# DRAGON this was the destroy code in tearDownModule()
def tearDownModule():
    """Clean up the environmen"""

    need_sleep = False

    # bpo-26762: Some multiprocessing objects like Pool create reference
    # cycles. Trigger a garbage collection to break these cycles.
    test.support.gc_collect()

    multiprocessing.set_start_method(old_start_method[0], force=True)
    # pause a bit so we don't get warning about dangling threads/processes
    processes = set(multiprocessing.process._dangling) - set(dangling[0])
    if processes:
        need_sleep = True
        test.support.environment_altered = True
        support.print_warning(f"Dangling processes: {processes}")
    processes = None

    threads = set(threading._dangling) - set(dangling[1])
    if threads:
        need_sleep = True
        test.support.environment_altered = True
        support.print_warning(f"Dangling threads: {threads}")
    threads = None

    # Sleep 500 ms to give time to child processes to complete.
    if need_sleep:
        time.sleep(0.5)

    multiprocessing.util._cleanup_tests()

examples/multiprocessing/unittests/common.py
DragonHPC/dragon
import dragon
import multiprocessing, sys


def foo():
    print("123")


# Because "if __name__ == '__main__'" is missing this will not work
# correctly on Windows.  However, we should get a RuntimeError rather
# than the Windows equivalent of a fork bomb.

if len(sys.argv) > 1:
    multiprocessing.set_start_method(sys.argv[1])
else:
    multiprocessing.set_start_method("spawn")

p = multiprocessing.Process(target=foo)
p.start()
p.join()
sys.exit(p.exitcode)

examples/multiprocessing/unittests/mp_fork_bomb.py
DragonHPC/dragon
"""Test multiprocessing Array()
"""
import unittest
import array

try:
    from ctypes import c_int
except ImportError:
    Structure = object
    c_int = c_double = c_longlong = None

import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing import util

from common import (
    BaseTestCase,
    ProcessesMixin,
    setUpModule,
    tearDownModule,
)


class WithProcessesTestArray(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    @classmethod
    def f(cls, seq):
        for i in range(1, len(seq)):
            seq[i] += seq[i - 1]

    @unittest.skipIf(c_int is None, "requires _ctypes")
    def test_array(self, raw=False):
        seq = [680, 626, 934, 821, 150, 233, 548, 982, 714, 831]
        if raw:
            arr = self.RawArray("i", seq)
        else:
            arr = self.Array("i", seq)

        self.assertEqual(len(arr), len(seq))
        self.assertEqual(arr[3], seq[3])
        self.assertEqual(list(arr[2:7]), list(seq[2:7]))

        arr[4:8] = seq[4:8] = array.array("i", [1, 2, 3, 4])

        self.assertEqual(list(arr[:]), seq)

        self.f(seq)

        p = self.Process(target=self.f, args=(arr,))
        p.daemon = True
        p.start()
        p.join()

        self.assertEqual(list(arr[:]), seq)

    @unittest.skipIf(c_int is None, "requires _ctypes")
    def test_array_from_size(self):
        size = 10
        # Test for zeroing (see issue #11675).
        # The repetition below strengthens the test by increasing the chances
        # of previously allocated non-zero memory being used for the new array
        # on the 2nd and 3rd loops.
        for _ in range(3):
            arr = self.Array("i", size)
            self.assertEqual(len(arr), size)
            self.assertEqual(list(arr), [0] * size)
            arr[:] = range(10)
            self.assertEqual(list(arr), list(range(10)))
            del arr

    @unittest.skipIf(c_int is None, "requires _ctypes")
    def test_rawarray(self):
        self.test_array(raw=True)

    @unittest.skipIf(c_int is None, "requires _ctypes")
    def test_getobj_getlock_obj(self):
        arr1 = self.Array("i", list(range(10)))
        lock1 = arr1.get_lock()
        obj1 = arr1.get_obj()

        arr2 = self.Array("i", list(range(10)), lock=None)
        lock2 = arr2.get_lock()
        obj2 = arr2.get_obj()

        lock = self.Lock()
        arr3 = self.Array("i", list(range(10)), lock=lock)
        lock3 = arr3.get_lock()
        obj3 = arr3.get_obj()
        self.assertEqual(lock, lock3)

        arr4 = self.Array("i", range(10), lock=False)
        self.assertFalse(hasattr(arr4, "get_lock"))
        self.assertFalse(hasattr(arr4, "get_obj"))
        self.assertRaises(AttributeError, self.Array, "i", range(10), lock="notalock")

        arr5 = self.RawArray("i", range(10))
        self.assertFalse(hasattr(arr5, "get_lock"))
        self.assertFalse(hasattr(arr5, "get_obj"))


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_array.py
DragonHPC/dragon
""" Test mulitprocessing barrier
"""
import unittest
import time
import os
import struct
import weakref

from test.support import hashlib_helper

import threading

import dragon  # DRAGON import before multiprocessing

import multiprocessing

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    setUpModule,
    tearDownModule,
    close_queue,
)

#
# Tests for Barrier - adapted from tests in test/lock_tests.py
#

# Many of the tests for threading.Barrier use a list as an atomic
# counter: a value is appended to increment the counter, and the
# length of the list gives the value.  We use the class DummyList
# for the same purpose.


class _DummyList(object):
    def __init__(self):
        wrapper = multiprocessing.heap.BufferWrapper(struct.calcsize("i"))
        lock = multiprocessing.Lock()
        self.__setstate__((wrapper, lock))
        self._lengthbuf[0] = 0

    def __setstate__(self, state):
        (self._wrapper, self._lock) = state
        self._lengthbuf = self._wrapper.create_memoryview().cast("i")

    def __getstate__(self):
        return (self._wrapper, self._lock)

    def append(self, _):
        with self._lock:
            self._lengthbuf[0] += 1

    def __len__(self):
        with self._lock:
            return self._lengthbuf[0]


def _wait():
    # A crude wait/yield function not relying on synchronization primitives.
    time.sleep(0.01)


class Bunch(object):
    """
    A bunch of threads.
    """

    def __init__(self, namespace, f, args, n, wait_before_exit=False):
        """
        Construct a bunch of `n` threads running the same function `f`.
        If `wait_before_exit` is True, the threads won't terminate until
        do_finish() is called.
        """
        self.f = f
        self.args = args
        self.n = n
        self.started = namespace.DummyList()
        self.finished = namespace.DummyList()
        self._can_exit = namespace.Event()
        if not wait_before_exit:
            self._can_exit.set()

        threads = []
        for i in range(n):
            p = namespace.Process(target=self.task)
            p.daemon = True
            p.start()
            threads.append(p)

        def finalize(threads):
            for p in threads:
                p.join()

        self._finalizer = weakref.finalize(self, finalize, threads)

    def task(self):
        pid = os.getpid()
        self.started.append(pid)
        try:
            self.f(*self.args)
        finally:
            self.finished.append(pid)
            self._can_exit.wait(30)
            assert self._can_exit.is_set()
    def wait_for_started(self):
        while len(self.started) < self.n:
            _wait()

    def wait_for_finished(self):
        while len(self.finished) < self.n:
           _wait()

    def do_finish(self):
        self._can_exit.set()

    def close(self):
        self._finalizer()


class AppendTrue(object):
    def __init__(self, obj):
        self.obj = obj

    def __call__(self):
        self.obj.append(True)

class WithProcessesTestBarrier(BaseTestCase, ProcessesMixin, unittest.TestCase):
    """
    Tests for Barrier objects.
    """

    N = 5
    defaultTimeout = 30.0  # XXX Slow Windows buildbots need generous timeout

    def setUp(self):
        self.barrier = self.Barrier(self.N, timeout=self.defaultTimeout)

    def tearDown(self):
        self.barrier.abort()
        self.barrier = None

    def DummyList(self):
        if self.TYPE == "threads":
            return []
        elif self.TYPE == "manager":
            return self.manager.list()
        else:
            return _DummyList()

    def run_threads(self, f, args):
        b = Bunch(self, f, args, self.N - 1)
        try:
            f(*args)
            b.wait_for_finished()
        finally:
            b.close()

    @classmethod
    def multipass(cls, barrier, results, n):
        m = barrier.parties
        assert m == cls.N
        for i in range(n):
            results[0].append(True)
            assert len(results[1]) == i * m
            barrier.wait()
            results[1].append(True)
            assert len(results[0]) == (i + 1) * m
            barrier.wait()
        try:
            assert barrier.n_waiting == 0
        except NotImplementedError:
            pass
        assert not barrier.broken

    def test_barrier(self, passes=1):
        """
        Test that a barrier is passed in lockstep
        """
        results = [self.DummyList(), self.DummyList()]
        self.run_threads(self.multipass, (self.barrier, results, passes))

    def test_barrier_10(self):
        """
        Test that a barrier works for 10 consecutive runs
        """
        return self.test_barrier(10)

    @classmethod
    def _test_wait_return_f(cls, barrier, queue):
        res = barrier.wait()
        queue.put(res)

    def test_wait_return(self):
        """
        test the return value from barrier.wait
        """
        queue = self.Queue()
        self.run_threads(self._test_wait_return_f, (self.barrier, queue))
        results = [queue.get() for i in range(self.N)]
        self.assertEqual(results.count(0), 1)
        close_queue(queue)

    @classmethod
    def _test_action_f(cls, barrier, results):
        barrier.wait()
        if len(results) != 1:
            raise RuntimeError

    def test_action(self):
        """
        Test the 'action' callback
        """
        results = self.DummyList()
        barrier = self.Barrier(self.N, action=AppendTrue(results))
        self.run_threads(self._test_action_f, (barrier, results))
        self.assertEqual(len(results), 1)

    @classmethod
    def _test_abort_f(cls, barrier, results1, results2):
        try:
            i = barrier.wait()
            if i == cls.N // 2:
                raise RuntimeError
            barrier.wait()
            results1.append(True)
        except threading.BrokenBarrierError:
            results2.append(True)
        except RuntimeError:
            barrier.abort()

    def test_abort(self):
        """
        Test that an abort will put the barrier in a broken state
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        self.run_threads(self._test_abort_f, (self.barrier, results1, results2))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertTrue(self.barrier.broken)

    @classmethod
    def _test_reset_f(cls, barrier, results1, results2, results3):
        i = barrier.wait()
        if i == cls.N // 2:
            # Wait until the other threads are all in the barrier.
            while barrier.n_waiting < cls.N - 1:
                time.sleep(0.001)
            
            barrier.reset()
        else:
            try:
                barrier.wait()
                results1.append(True)
            except threading.BrokenBarrierError:
                results2.append(True)
        # Now, pass the barrier again
        barrier.wait()
        results3.append(True)

    def test_reset(self):
        """
        Test that a 'reset' on a barrier frees the waiting threads
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        results3 = self.DummyList()
        self.run_threads(self._test_reset_f, (self.barrier, results1, results2, results3))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertEqual(len(results3), self.N)

    @classmethod
    def _test_abort_and_reset_f(cls, barrier, barrier2, results1, results2, results3):
        try:
            i = barrier.wait()
            if i == cls.N // 2:
                raise RuntimeError
            barrier.wait()
            results1.append(True)
        except threading.BrokenBarrierError:
            results2.append(True)
        except RuntimeError:
            barrier.abort()
        # Synchronize and reset the barrier.  Must synchronize first so
        # that everyone has left it when we reset, and after so that no
        # one enters it before the reset.
        if barrier2.wait() == cls.N // 2:
            barrier.reset()
        barrier2.wait()
        barrier.wait()
        results3.append(True)

    def test_abort_and_reset(self):
        """
        Test that a barrier can be reset after being broken.
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        results3 = self.DummyList()
        barrier2 = self.Barrier(self.N)

        self.run_threads(self._test_abort_and_reset_f, (self.barrier, barrier2, results1, results2, results3))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertEqual(len(results3), self.N)

    @classmethod
    def _test_timeout_f(cls, barrier, results):
        i = barrier.wait()
        if i == cls.N // 2:
            # One thread is late!
            time.sleep(1.0)
        try:
            barrier.wait(0.5)
        except threading.BrokenBarrierError:
            results.append(True)

    def test_timeout(self):
        """
        Test wait(timeout)
        """
        results = self.DummyList()
        self.run_threads(self._test_timeout_f, (self.barrier, results))
        self.assertEqual(len(results), self.barrier.parties)

    @classmethod
    def _test_default_timeout_f(cls, barrier, results):
        i = barrier.wait(cls.defaultTimeout)
        if i == cls.N // 2:
            # One thread is later than the default timeout
            time.sleep(1.0)
        try:
            barrier.wait()
        except threading.BrokenBarrierError:
            results.append(True)

    def test_default_timeout(self):
        """
        Test the barrier's default timeout
        """
        barrier = self.Barrier(self.N, timeout=0.5)
        results = self.DummyList()
        self.run_threads(self._test_default_timeout_f, (barrier, results))
        self.assertEqual(len(results), barrier.parties)

    def test_single_thread(self):
        b = self.Barrier(1)
        b.wait()
        b.wait()

    @classmethod
    def _test_thousand_f(cls, barrier, passes, conn, lock):
        for i in range(passes):
            barrier.wait()
            with lock:
                conn.send(i)

    def test_thousand(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        passes = 1000
        lock = self.Lock()
        conn, child_conn = self.Pipe(False)
        for j in range(self.N):
            p = self.Process(target=self._test_thousand_f, args=(self.barrier, passes, child_conn, lock))
            p.start()
            self.addCleanup(p.join)

        for i in range(passes):
            for j in range(self.N):
                self.assertEqual(conn.recv(), i)


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestBarrier(BaseTestCase, ManagerMixin, unittest.TestCase):
    """
    Tests for Barrier objects.
    """

    N = 5
    defaultTimeout = 30.0  # XXX Slow Windows buildbots need generous timeout

    def setUp(self):
        self.barrier = self.Barrier(self.N, timeout=self.defaultTimeout)

    def tearDown(self):
        self.barrier.abort()
        self.barrier = None

    def DummyList(self):
        if self.TYPE == "threads":
            return []
        elif self.TYPE == "manager":
            return self.manager.list()
        else:
            return _DummyList()

    def run_threads(self, f, args):
        b = Bunch(self, f, args, self.N - 1)
        try:
            f(*args)
            b.wait_for_finished()
        finally:
            b.close()

    @classmethod
    def multipass(cls, barrier, results, n):
        m = barrier.parties
        assert m == cls.N
        for i in range(n):
            results[0].append(True)
            assert len(results[1]) == i * m
            barrier.wait()
            results[1].append(True)
            assert len(results[0]) == (i + 1) * m
            barrier.wait()
        try:
            assert barrier.n_waiting == 0
        except NotImplementedError:
            pass
        assert not barrier.broken

    def test_barrier(self, passes=1):
        """
        Test that a barrier is passed in lockstep
        """
        results = [self.DummyList(), self.DummyList()]
        self.run_threads(self.multipass, (self.barrier, results, passes))

    def test_barrier_10(self):
        """
        Test that a barrier works for 10 consecutive runs
        """
        return self.test_barrier(10)

    @classmethod
    def _test_wait_return_f(cls, barrier, queue):
        res = barrier.wait()
        queue.put(res)

    def test_wait_return(self):
        """
        test the return value from barrier.wait
        """
        queue = self.Queue()
        self.run_threads(self._test_wait_return_f, (self.barrier, queue))
        results = [queue.get() for i in range(self.N)]
        self.assertEqual(results.count(0), 1)
        close_queue(queue)

    @classmethod
    def _test_action_f(cls, barrier, results):
        barrier.wait()
        if len(results) != 1:
            raise RuntimeError

    def test_action(self):
        """
        Test the 'action' callback
        """
        results = self.DummyList()
        barrier = self.Barrier(self.N, action=AppendTrue(results))
        self.run_threads(self._test_action_f, (barrier, results))
        self.assertEqual(len(results), 1)

    @classmethod
    def _test_abort_f(cls, barrier, results1, results2):
        try:
            i = barrier.wait()
            if i == cls.N // 2:
                raise RuntimeError
            barrier.wait()
            results1.append(True)
        except threading.BrokenBarrierError:
            results2.append(True)
        except RuntimeError:
            barrier.abort()

    def test_abort(self):
        """
        Test that an abort will put the barrier in a broken state
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        self.run_threads(self._test_abort_f, (self.barrier, results1, results2))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertTrue(self.barrier.broken)

    @classmethod
    def _test_reset_f(cls, barrier, results1, results2, results3):
        i = barrier.wait()
        if i == cls.N // 2:
            # Wait until the other threads are all in the barrier.
            while barrier.n_waiting < cls.N - 1:
                time.sleep(0.001)
            barrier.reset()
        else:
            try:
                barrier.wait()
                results1.append(True)
            except threading.BrokenBarrierError:
                results2.append(True)
        # Now, pass the barrier again
        barrier.wait()
        results3.append(True)

    def test_reset(self):
        """
        Test that a 'reset' on a barrier frees the waiting threads
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        results3 = self.DummyList()
        self.run_threads(self._test_reset_f, (self.barrier, results1, results2, results3))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertEqual(len(results3), self.N)

    @classmethod
    def _test_abort_and_reset_f(cls, barrier, barrier2, results1, results2, results3):
        try:
            i = barrier.wait()
            if i == cls.N // 2:
                raise RuntimeError
            barrier.wait()
            results1.append(True)
        except threading.BrokenBarrierError:
            results2.append(True)
        except RuntimeError:
            barrier.abort()
        # Synchronize and reset the barrier.  Must synchronize first so
        # that everyone has left it when we reset, and after so that no
        # one enters it before the reset.
        if barrier2.wait() == cls.N // 2:
            barrier.reset()
        barrier2.wait()
        barrier.wait()
        results3.append(True)

    def test_abort_and_reset(self):
        """
        Test that a barrier can be reset after being broken.
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        results3 = self.DummyList()
        barrier2 = self.Barrier(self.N)

        self.run_threads(self._test_abort_and_reset_f, (self.barrier, barrier2, results1, results2, results3))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertEqual(len(results3), self.N)

    @classmethod
    def _test_timeout_f(cls, barrier, results):
        i = barrier.wait()
        if i == cls.N // 2:
            # One thread is late!
            time.sleep(1.0)
        try:
            barrier.wait(0.5)
        except threading.BrokenBarrierError:
            results.append(True)

    def test_timeout(self):
        """
        Test wait(timeout)
        """
        results = self.DummyList()
        self.run_threads(self._test_timeout_f, (self.barrier, results))
        self.assertEqual(len(results), self.barrier.parties)

    @classmethod
    def _test_default_timeout_f(cls, barrier, results):
        i = barrier.wait(cls.defaultTimeout)
        if i == cls.N // 2:
            # One thread is later than the default timeout
            time.sleep(1.0)
        try:
            barrier.wait()
        except threading.BrokenBarrierError:
            results.append(True)

    def test_default_timeout(self):
        """
        Test the barrier's default timeout
        """
        barrier = self.Barrier(self.N, timeout=0.5)
        results = self.DummyList()
        self.run_threads(self._test_default_timeout_f, (barrier, results))
        self.assertEqual(len(results), barrier.parties)

    def test_single_thread(self):
        b = self.Barrier(1)
        b.wait()
        b.wait()

    @classmethod
    def _test_thousand_f(cls, barrier, passes, conn, lock):
        for i in range(passes):
            barrier.wait()
            with lock:
                conn.send(i)

    def test_thousand(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        passes = 1000
        lock = self.Lock()
        conn, child_conn = self.Pipe(False)
        for j in range(self.N):
            p = self.Process(target=self._test_thousand_f, args=(self.barrier, passes, child_conn, lock))
            p.start()
            self.addCleanup(p.join)

        for i in range(passes):
            for j in range(self.N):
                self.assertEqual(conn.recv(), i)


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestBarrier(BaseTestCase, ThreadsMixin, unittest.TestCase):
    """
    Tests for Barrier objects.
    """

    N = 5
    defaultTimeout = 30.0  # XXX Slow Windows buildbots need generous timeout

    def setUp(self):
        self.barrier = self.Barrier(self.N, timeout=self.defaultTimeout)

    def tearDown(self):
        self.barrier.abort()
        self.barrier = None

    def DummyList(self):
        if self.TYPE == "threads":
            return []
        elif self.TYPE == "manager":
            return self.manager.list()
        else:
            return _DummyList()

    def run_threads(self, f, args):
        b = Bunch(self, f, args, self.N - 1)
        try:
            f(*args)
            b.wait_for_finished()
        finally:
            b.close()

    @classmethod
    def multipass(cls, barrier, results, n):
        m = barrier.parties
        assert m == cls.N
        for i in range(n):
            results[0].append(True)
            assert len(results[1]) == i * m
            barrier.wait()
            results[1].append(True)
            assert len(results[0]) == (i + 1) * m
            barrier.wait()
        try:
            assert barrier.n_waiting == 0
        except NotImplementedError:
            pass
        assert not barrier.broken

    def test_barrier(self, passes=1):
        """
        Test that a barrier is passed in lockstep
        """
        results = [self.DummyList(), self.DummyList()]
        self.run_threads(self.multipass, (self.barrier, results, passes))

    def test_barrier_10(self):
        """
        Test that a barrier works for 10 consecutive runs
        """
        return self.test_barrier(10)

    @classmethod
    def _test_wait_return_f(cls, barrier, queue):
        res = barrier.wait()
        queue.put(res)

    def test_wait_return(self):
        """
        test the return value from barrier.wait
        """
        queue = self.Queue()
        self.run_threads(self._test_wait_return_f, (self.barrier, queue))
        results = [queue.get() for i in range(self.N)]
        self.assertEqual(results.count(0), 1)
        close_queue(queue)

    @classmethod
    def _test_action_f(cls, barrier, results):
        barrier.wait()
        if len(results) != 1:
            raise RuntimeError

    def test_action(self):
        """
        Test the 'action' callback
        """
        results = self.DummyList()
        barrier = self.Barrier(self.N, action=AppendTrue(results))
        self.run_threads(self._test_action_f, (barrier, results))
        self.assertEqual(len(results), 1)

    @classmethod
    def _test_abort_f(cls, barrier, results1, results2):
        try:
            i = barrier.wait()
            if i == cls.N // 2:
                raise RuntimeError
            barrier.wait()
            results1.append(True)
        except threading.BrokenBarrierError:
            results2.append(True)
        except RuntimeError:
            barrier.abort()

    def test_abort(self):
        """
        Test that an abort will put the barrier in a broken state
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        self.run_threads(self._test_abort_f, (self.barrier, results1, results2))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertTrue(self.barrier.broken)

    @classmethod
    def _test_reset_f(cls, barrier, results1, results2, results3):
        i = barrier.wait()
        if i == cls.N // 2:
            # Wait until the other threads are all in the barrier.
            while barrier.n_waiting < cls.N - 1:
                time.sleep(0.001)
            barrier.reset()
        else:
            try:
                barrier.wait()
                results1.append(True)
            except threading.BrokenBarrierError:
                results2.append(True)
        # Now, pass the barrier again
        barrier.wait()
        results3.append(True)

    def test_reset(self):
        """
        Test that a 'reset' on a barrier frees the waiting threads
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        results3 = self.DummyList()
        self.run_threads(self._test_reset_f, (self.barrier, results1, results2, results3))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertEqual(len(results3), self.N)

    @classmethod
    def _test_abort_and_reset_f(cls, barrier, barrier2, results1, results2, results3):
        try:
            i = barrier.wait()
            if i == cls.N // 2:
                raise RuntimeError
            barrier.wait()
            results1.append(True)
        except threading.BrokenBarrierError:
            results2.append(True)
        except RuntimeError:
            barrier.abort()
        # Synchronize and reset the barrier.  Must synchronize first so
        # that everyone has left it when we reset, and after so that no
        # one enters it before the reset.
        if barrier2.wait() == cls.N // 2:
            barrier.reset()
        barrier2.wait()
        barrier.wait()
        results3.append(True)

    def test_abort_and_reset(self):
        """
        Test that a barrier can be reset after being broken.
        """
        results1 = self.DummyList()
        results2 = self.DummyList()
        results3 = self.DummyList()
        barrier2 = self.Barrier(self.N)

        self.run_threads(self._test_abort_and_reset_f, (self.barrier, barrier2, results1, results2, results3))
        self.assertEqual(len(results1), 0)
        self.assertEqual(len(results2), self.N - 1)
        self.assertEqual(len(results3), self.N)

    @classmethod
    def _test_timeout_f(cls, barrier, results):
        i = barrier.wait()
        if i == cls.N // 2:
            # One thread is late!
            time.sleep(1.0)
        try:
            barrier.wait(0.5)
        except threading.BrokenBarrierError:
            results.append(True)

    def test_timeout(self):
        """
        Test wait(timeout)
        """
        results = self.DummyList()
        self.run_threads(self._test_timeout_f, (self.barrier, results))
        self.assertEqual(len(results), self.barrier.parties)

    @classmethod
    def _test_default_timeout_f(cls, barrier, results):
        i = barrier.wait(cls.defaultTimeout)
        if i == cls.N // 2:
            # One thread is later than the default timeout
            time.sleep(1.0)
        try:
            barrier.wait()
        except threading.BrokenBarrierError:
            results.append(True)

    def test_default_timeout(self):
        """
        Test the barrier's default timeout
        """
        barrier = self.Barrier(self.N, timeout=0.5)
        results = self.DummyList()
        self.run_threads(self._test_default_timeout_f, (barrier, results))
        self.assertEqual(len(results), barrier.parties)

    def test_single_thread(self):
        b = self.Barrier(1)
        b.wait()
        b.wait()

    @classmethod
    def _test_thousand_f(cls, barrier, passes, conn, lock):
        for i in range(passes):
            barrier.wait()
            with lock:
                conn.send(i)

    def test_thousand(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        passes = 1000
        lock = self.Lock()
        conn, child_conn = self.Pipe(False)
        for j in range(self.N):
            p = self.Process(target=self._test_thousand_f, args=(self.barrier, passes, child_conn, lock))
            p.start()
            self.addCleanup(p.join)

        for i in range(passes):
            for j in range(self.N):
                self.assertEqual(conn.recv(), i)


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_barrier.py
DragonHPC/dragon
""" Test mulitprocessing.Condition
"""
import unittest
import time
import sys
import os
import signal

from test import support

import threading

import dragon  # DRAGON import before multiprocessing

import multiprocessing

try:
    from multiprocessing.sharedctypes import Value, copy

    HAS_SHAREDCTYPES = True
except ImportError:
    HAS_SHAREDCTYPES = False

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    TimingWrapper,
    setUpModule,
    tearDownModule,
    get_value,
    join_process,
    DELTA,
    TIMEOUT1,
)


class WithProcessesTestCondition(BaseTestCase, ProcessesMixin, unittest.TestCase):
    @classmethod
    def f(cls, cond, sleeping, woken, timeout=None):
        cond.acquire()
        sleeping.release()
        cond.wait(timeout)
        woken.release()
        cond.release()

    def assertReachesEventually(self, func, value):
        for i in range(10):
            try:
                if func() == value:
                    break
            except NotImplementedError:
                break
            time.sleep(DELTA)
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(value, func)

    def check_invariant(self, cond):
        # this is only supposed to succeed when there are no sleepers
        if self.TYPE == "processes":
            try:
                sleepers = cond._sleeping_count.get_value() - cond._woken_count.get_value()
                self.assertEqual(sleepers, 0)
                self.assertEqual(cond._wait_semaphore.get_value(), 0)
            except NotImplementedError:
                pass

    def test_notify(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        p = self.Process(target=self.f, args=(cond, sleeping, woken))
        p.daemon = True
        p.start()
        self.addCleanup(p.join)

        p = threading.Thread(target=self.f, args=(cond, sleeping, woken))
        p.daemon = True
        p.start()
        self.addCleanup(p.join)

        # wait for both children to start sleeping
        sleeping.acquire()
        sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake up one process/thread
        cond.acquire()
        cond.notify()
        cond.release()

        # check one process/thread has woken up
        time.sleep(DELTA)
        self.assertReachesEventually(lambda: get_value(woken), 1)
        # self.assertReturnsIfImplemented(1, get_value, woken) # DRAGON: test not deterministic

        # wake up another
        cond.acquire()
        cond.notify()
        cond.release()

        # check other has woken up
        time.sleep(DELTA)
        self.assertReachesEventually(lambda: get_value(woken), 2)
        # self.assertReturnsIfImplemented(2, get_value, woken) # DRAGON: test not deterministic

        # check state is not mucked up
        self.check_invariant(cond)
        p.join()

    def test_notify_all(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        # start some threads/processes which will timeout
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken, TIMEOUT1))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken, TIMEOUT1))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them all to sleep
        for i in range(6):
            sleeping.acquire()

        # check they have all timed out
        for i in range(6):
            woken.acquire()
        self.assertReturnsIfImplemented(0, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)

        # start some more threads/processes
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them to all sleep
        for i in range(6):
            sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake them all up
        cond.acquire()
        cond.notify_all()
        cond.release()

        # check they have all woken
        self.assertReachesEventually(lambda: get_value(woken), 6)

        # check state is not mucked up
        self.check_invariant(cond)

    def test_notify_n(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        # start some threads/processes
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them to all sleep
        for i in range(6):
            sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake some of them up
        cond.acquire()
        cond.notify(n=2)
        cond.release()

        # check 2 have woken
        self.assertReachesEventually(lambda: get_value(woken), 2)

        # wake the rest of them
        cond.acquire()
        cond.notify(n=4)
        cond.release()

        self.assertReachesEventually(lambda: get_value(woken), 6)

        # doesn't do anything more
        cond.acquire()
        cond.notify(n=3)
        cond.release()

        self.assertReturnsIfImplemented(6, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)

    def test_timeout(self):
        cond = self.Condition()
        wait = TimingWrapper(cond.wait)
        cond.acquire()
        res = wait(TIMEOUT1)
        cond.release()
        self.assertEqual(res, False)
        self.assertTimingAlmostEqual(wait.elapsed, TIMEOUT1)

    @classmethod
    def _test_waitfor_f(cls, cond, state):
        with cond:
            state.value = 0
            cond.notify()
            result = cond.wait_for(lambda: state.value == 4)
            if not result or state.value != 4:
                sys.exit(1)

    @unittest.skip("DRAGON: sharedctypes & Value not implemented")
    @unittest.skipUnless(HAS_SHAREDCTYPES, "needs sharedctypes")
    def test_waitfor(self):
        # based on test in test/lock_tests.py
        cond = self.Condition()
        state = self.Value("i", -1)

        p = self.Process(target=self._test_waitfor_f, args=(cond, state))
        p.daemon = True
        p.start()

        with cond:
            result = cond.wait_for(lambda: state.value == 0)
            self.assertTrue(result)
            self.assertEqual(state.value, 0)

        for i in range(4):
            time.sleep(0.01)
            with cond:
                state.value += 1
                cond.notify()

        join_process(p)
        self.assertEqual(p.exitcode, 0)

    @classmethod
    def _test_waitfor_timeout_f(cls, cond, state, success, sem):
        sem.release()
        with cond:
            expected = 0.1
            dt = time.monotonic()
            result = cond.wait_for(lambda: state.value == 4, timeout=expected)
            dt = time.monotonic() - dt
            # borrow logic in assertTimeout() from test/lock_tests.py
            if not result and expected * 0.6 < dt < expected * 10.0:
                success.value = True

    @unittest.skip("DRAGON: sharedctypes & Value not implemented")
    @unittest.skipUnless(HAS_SHAREDCTYPES, "needs sharedctypes")
    def test_waitfor_timeout(self):
        # based on test in test/lock_tests.py
        cond = self.Condition()
        state = self.Value("i", 0)
        success = self.Value("i", False)
        sem = self.Semaphore(0)

        p = self.Process(target=self._test_waitfor_timeout_f, args=(cond, state, success, sem))
        p.daemon = True
        p.start()
        self.assertTrue(sem.acquire(timeout=support.LONG_TIMEOUT))

        # Only increment 3 times, so state == 4 is never reached.
        for i in range(3):
            time.sleep(0.01)
            with cond:
                state.value += 1
                cond.notify()

        join_process(p)
        self.assertTrue(success.value)

    @classmethod
    def _test_wait_result(cls, c, pid):
        with c:
            c.notify()
        time.sleep(1)
        if pid is not None:
            os.kill(pid, signal.SIGINT)

    def test_wait_result(self):
        if isinstance(self, ProcessesMixin) and sys.platform != "win32":
            pid = os.getpid()
        else:
            pid = None

        c = self.Condition()
        with c:
            self.assertFalse(c.wait(0))
            self.assertFalse(c.wait(0.1))

            p = self.Process(target=self._test_wait_result, args=(c, pid))
            p.start()

            self.assertTrue(c.wait(60))
            if pid is not None:
                self.assertRaises(KeyboardInterrupt, c.wait, 60)

            p.join()


@unittest.skip("DRAGON: Manager not implemented")
class WithManagerTestCondition(BaseTestCase, ManagerMixin, unittest.TestCase):
    @classmethod
    def f(cls, cond, sleeping, woken, timeout=None):
        cond.acquire()
        sleeping.release()
        cond.wait(timeout)
        woken.release()
        cond.release()

    def assertReachesEventually(self, func, value):
        for i in range(10):
            try:
                if func() == value:
                    break
            except NotImplementedError:
                break
            time.sleep(DELTA)
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(value, func)

    def check_invariant(self, cond):
        # this is only supposed to succeed when there are no sleepers
        if self.TYPE == "processes":
            try:
                sleepers = cond._sleeping_count.get_value() - cond._woken_count.get_value()
                self.assertEqual(sleepers, 0)
                self.assertEqual(cond._wait_semaphore.get_value(), 0)
            except NotImplementedError:
                pass

    def test_notify(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        p = self.Process(target=self.f, args=(cond, sleeping, woken))
        p.daemon = True
        p.start()
        self.addCleanup(p.join)

        p = threading.Thread(target=self.f, args=(cond, sleeping, woken))
        p.daemon = True
        p.start()
        self.addCleanup(p.join)

        # wait for both children to start sleeping
        sleeping.acquire()
        sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake up one process/thread
        cond.acquire()
        cond.notify()
        cond.release()

        # check one process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(1, get_value, woken)

        # wake up another
        cond.acquire()
        cond.notify()
        cond.release()

        # check other has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(2, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)
        p.join()

    def test_notify_all(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        # start some threads/processes which will timeout
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken, TIMEOUT1))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken, TIMEOUT1))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them all to sleep
        for i in range(6):
            sleeping.acquire()

        # check they have all timed out
        for i in range(6):
            woken.acquire()
        self.assertReturnsIfImplemented(0, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)

        # start some more threads/processes
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them to all sleep
        for i in range(6):
            sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake them all up
        cond.acquire()
        cond.notify_all()
        cond.release()

        # check they have all woken
        self.assertReachesEventually(lambda: get_value(woken), 6)

        # check state is not mucked up
        self.check_invariant(cond)

    def test_notify_n(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        # start some threads/processes
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them to all sleep
        for i in range(6):
            sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake some of them up
        cond.acquire()
        cond.notify(n=2)
        cond.release()

        # check 2 have woken
        self.assertReachesEventually(lambda: get_value(woken), 2)

        # wake the rest of them
        cond.acquire()
        cond.notify(n=4)
        cond.release()

        self.assertReachesEventually(lambda: get_value(woken), 6)

        # doesn't do anything more
        cond.acquire()
        cond.notify(n=3)
        cond.release()

        self.assertReturnsIfImplemented(6, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)

    def test_timeout(self):
        cond = self.Condition()
        wait = TimingWrapper(cond.wait)
        cond.acquire()
        res = wait(TIMEOUT1)
        cond.release()
        self.assertEqual(res, False)
        self.assertTimingAlmostEqual(wait.elapsed, TIMEOUT1)

    @classmethod
    def _test_waitfor_f(cls, cond, state):
        with cond:
            state.value = 0
            cond.notify()
            result = cond.wait_for(lambda: state.value == 4)
            if not result or state.value != 4:
                sys.exit(1)

    @unittest.skip("DRAGON: hangs")
    @unittest.skipUnless(HAS_SHAREDCTYPES, "needs sharedctypes")
    def test_waitfor(self):
        # based on test in test/lock_tests.py
        cond = self.Condition()
        state = self.Value("i", -1)

        p = self.Process(target=self._test_waitfor_f, args=(cond, state))
        p.daemon = True
        p.start()

        with cond:
            result = cond.wait_for(lambda: state.value == 0)
            self.assertTrue(result)
            self.assertEqual(state.value, 0)

        for i in range(4):
            time.sleep(0.01)
            with cond:
                state.value += 1
                cond.notify()

        join_process(p)
        self.assertEqual(p.exitcode, 0)

    @classmethod
    def _test_waitfor_timeout_f(cls, cond, state, success, sem):
        sem.release()
        with cond:
            expected = 0.1
            dt = time.monotonic()
            result = cond.wait_for(lambda: state.value == 4, timeout=expected)
            dt = time.monotonic() - dt
            # borrow logic in assertTimeout() from test/lock_tests.py
            if not result and expected * 0.6 < dt < expected * 10.0:
                success.value = True

    @unittest.skip("DRAGON: hangs")
    @unittest.skipUnless(HAS_SHAREDCTYPES, "needs sharedctypes")
    def test_waitfor_timeout(self):
        # based on test in test/lock_tests.py
        cond = self.Condition()
        state = self.Value("i", 0)
        success = self.Value("i", False)
        sem = self.Semaphore(0)

        p = self.Process(target=self._test_waitfor_timeout_f, args=(cond, state, success, sem))
        p.daemon = True
        p.start()
        self.assertTrue(sem.acquire(timeout=support.LONG_TIMEOUT))

        # Only increment 3 times, so state == 4 is never reached.
        for i in range(3):
            time.sleep(0.01)
            with cond:
                state.value += 1
                cond.notify()

        join_process(p)
        self.assertTrue(success.value)

    @classmethod
    def _test_wait_result(cls, c, pid):
        with c:
            c.notify()
        time.sleep(1)
        if pid is not None:
            os.kill(pid, signal.SIGINT)

    def test_wait_result(self):
        if isinstance(self, ProcessesMixin) and sys.platform != "win32":
            pid = os.getpid()
        else:
            pid = None

        c = self.Condition()
        with c:
            self.assertFalse(c.wait(0))
            self.assertFalse(c.wait(0.1))

            p = self.Process(target=self._test_wait_result, args=(c, pid))
            p.start()

            self.assertTrue(c.wait(60))
            if pid is not None:
                self.assertRaises(KeyboardInterrupt, c.wait, 60)

            p.join()


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestCondition(BaseTestCase, ThreadsMixin, unittest.TestCase):
    @classmethod
    def f(cls, cond, sleeping, woken, timeout=None):
        cond.acquire()
        sleeping.release()
        cond.wait(timeout)
        woken.release()
        cond.release()

    def assertReachesEventually(self, func, value):
        for i in range(10):
            try:
                if func() == value:
                    break
            except NotImplementedError:
                break
            time.sleep(DELTA)
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(value, func)

    def check_invariant(self, cond):
        # this is only supposed to succeed when there are no sleepers
        if self.TYPE == "processes":
            try:
                sleepers = cond._sleeping_count.get_value() - cond._woken_count.get_value()
                self.assertEqual(sleepers, 0)
                self.assertEqual(cond._wait_semaphore.get_value(), 0)
            except NotImplementedError:
                pass

    def test_notify(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        p = self.Process(target=self.f, args=(cond, sleeping, woken))
        p.daemon = True
        p.start()
        self.addCleanup(p.join)

        p = threading.Thread(target=self.f, args=(cond, sleeping, woken))
        p.daemon = True
        p.start()
        self.addCleanup(p.join)

        # wait for both children to start sleeping
        sleeping.acquire()
        sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake up one process/thread
        cond.acquire()
        cond.notify()
        cond.release()

        # check one process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(1, get_value, woken)

        # wake up another
        cond.acquire()
        cond.notify()
        cond.release()

        # check other has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(2, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)
        p.join()

    def test_notify_all(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        # start some threads/processes which will timeout
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken, TIMEOUT1))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken, TIMEOUT1))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them all to sleep
        for i in range(6):
            sleeping.acquire()

        # check they have all timed out
        for i in range(6):
            woken.acquire()
        self.assertReturnsIfImplemented(0, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)

        # start some more threads/processes
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them to all sleep
        for i in range(6):
            sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake them all up
        cond.acquire()
        cond.notify_all()
        cond.release()

        # check they have all woken
        self.assertReachesEventually(lambda: get_value(woken), 6)

        # check state is not mucked up
        self.check_invariant(cond)

    def test_notify_n(self):
        cond = self.Condition()
        sleeping = self.Semaphore(0)
        woken = self.Semaphore(0)

        # start some threads/processes
        for i in range(3):
            p = self.Process(target=self.f, args=(cond, sleeping, woken))
            p.daemon = True
            p.start()
            self.addCleanup(p.join)

            t = threading.Thread(target=self.f, args=(cond, sleeping, woken))
            t.daemon = True
            t.start()
            self.addCleanup(t.join)

        # wait for them to all sleep
        for i in range(6):
            sleeping.acquire()

        # check no process/thread has woken up
        time.sleep(DELTA)
        self.assertReturnsIfImplemented(0, get_value, woken)

        # wake some of them up
        cond.acquire()
        cond.notify(n=2)
        cond.release()

        # check 2 have woken
        self.assertReachesEventually(lambda: get_value(woken), 2)

        # wake the rest of them
        cond.acquire()
        cond.notify(n=4)
        cond.release()

        self.assertReachesEventually(lambda: get_value(woken), 6)

        # doesn't do anything more
        cond.acquire()
        cond.notify(n=3)
        cond.release()

        self.assertReturnsIfImplemented(6, get_value, woken)

        # check state is not mucked up
        self.check_invariant(cond)

    def test_timeout(self):
        cond = self.Condition()
        wait = TimingWrapper(cond.wait)
        cond.acquire()
        res = wait(TIMEOUT1)
        cond.release()
        self.assertEqual(res, False)
        self.assertTimingAlmostEqual(wait.elapsed, TIMEOUT1)

    @classmethod
    def _test_waitfor_f(cls, cond, state):
        with cond:
            state.value = 0
            cond.notify()
            result = cond.wait_for(lambda: state.value == 4)
            if not result or state.value != 4:
                sys.exit(1)

    @unittest.skip("DRAGON: hangs")
    @unittest.skipUnless(HAS_SHAREDCTYPES, "needs sharedctypes")
    def test_waitfor(self):
        # based on test in test/lock_tests.py
        cond = self.Condition()
        state = self.Value("i", -1)

        p = self.Process(target=self._test_waitfor_f, args=(cond, state))
        p.daemon = True
        p.start()

        with cond:
            result = cond.wait_for(lambda: state.value == 0)
            self.assertTrue(result)
            self.assertEqual(state.value, 0)

        for i in range(4):
            time.sleep(0.01)
            with cond:
                state.value += 1
                cond.notify()

        join_process(p)
        self.assertEqual(p.exitcode, 0)

    @classmethod
    def _test_waitfor_timeout_f(cls, cond, state, success, sem):
        sem.release()
        with cond:
            expected = 0.1
            dt = time.monotonic()
            result = cond.wait_for(lambda: state.value == 4, timeout=expected)
            dt = time.monotonic() - dt
            # borrow logic in assertTimeout() from test/lock_tests.py
            if not result and expected * 0.6 < dt < expected * 10.0:
                success.value = True

    @unittest.skip("DRAGON: hangs")
    @unittest.skipUnless(HAS_SHAREDCTYPES, "needs sharedctypes")
    def test_waitfor_timeout(self):
        # based on test in test/lock_tests.py
        cond = self.Condition()
        state = self.Value("i", 0)
        success = self.Value("i", False)
        sem = self.Semaphore(0)

        p = self.Process(target=self._test_waitfor_timeout_f, args=(cond, state, success, sem))
        p.daemon = True
        p.start()
        self.assertTrue(sem.acquire(timeout=support.LONG_TIMEOUT))

        # Only increment 3 times, so state == 4 is never reached.
        for i in range(3):
            time.sleep(0.01)
            with cond:
                state.value += 1
                cond.notify()

        join_process(p)
        self.assertTrue(success.value)

    @classmethod
    def _test_wait_result(cls, c, pid):
        with c:
            c.notify()
        time.sleep(1)
        if pid is not None:
            os.kill(pid, signal.SIGINT)

    def test_wait_result(self):
        if isinstance(self, ProcessesMixin) and sys.platform != "win32":
            pid = os.getpid()
        else:
            pid = None

        c = self.Condition()
        with c:
            self.assertFalse(c.wait(0))
            self.assertFalse(c.wait(0.1))

            p = self.Process(target=self._test_wait_result, args=(c, pid))
            p.start()

            self.assertTrue(c.wait(60))
            if pid is not None:
                self.assertRaises(KeyboardInterrupt, c.wait, 60)

            p.join()


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_condition.py
DragonHPC/dragon
"""Test multiprocessing connection, listener, client, pickling
"""
import unittest
import sys
import time
import array
import os
import errno
import socket

try:
    import msvcrt
except ImportError:
    msvcrt = None

import test.support
from test import support
from test.support import hashlib_helper
from test.support import socket_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing

try:
    from multiprocessing import reduction

    HAS_REDUCTION = reduction.HAVE_SEND_HANDLE
except ImportError:
    HAS_REDUCTION = False

from common import (
    BaseTestCase,
    ProcessesMixin,
    ThreadsMixin,
    TimingWrapper,
    TIMEOUT1,
    MAXFD,
    latin,
    setUpModule,
    tearDownModule,
)

#
#
#

SENTINEL = latin("")


class WithProcessesTestConnection(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    @classmethod
    def _echo(cls, conn):
        for msg in iter(conn.recv_bytes, SENTINEL):
            conn.send_bytes(msg)
        conn.close()

    @unittest.skip("bug filed PE-40894")
    def test_connection(self):
        conn, child_conn = self.Pipe()

        p = self.Process(target=self._echo, args=(child_conn,))
        p.daemon = True
        p.start()

        seq = [1, 2.25, None]
        msg = latin("hello world")
        longmsg = msg * 10
        arr = array.array("i", list(range(4)))

        if self.TYPE == "processes":
            self.assertEqual(type(conn.fileno()), int)

        self.assertEqual(conn.send(seq), None)
        self.assertEqual(conn.recv(), seq)

        self.assertEqual(conn.send_bytes(msg), None)
        self.assertEqual(conn.recv_bytes(), msg)

        if self.TYPE == "processes":
            buffer = array.array("i", [0] * 10)
            expected = list(arr) + [0] * (10 - len(arr))
            self.assertEqual(conn.send_bytes(arr), None)
            self.assertEqual(conn.recv_bytes_into(buffer), len(arr) * buffer.itemsize)
            self.assertEqual(list(buffer), expected)

            buffer = array.array("i", [0] * 10)
            expected = [0] * 3 + list(arr) + [0] * (10 - 3 - len(arr))
            self.assertEqual(conn.send_bytes(arr), None)
            self.assertEqual(conn.recv_bytes_into(buffer, 3 * buffer.itemsize), len(arr) * buffer.itemsize)
            self.assertEqual(list(buffer), expected)

            buffer = bytearray(latin(" " * 40))
            self.assertEqual(conn.send_bytes(longmsg), None)
            try:
                res = conn.recv_bytes_into(buffer)
            except multiprocessing.BufferTooShort as e:
                self.assertEqual(e.args, (longmsg,))
            else:
                self.fail("expected BufferTooShort, got %s" % res)

        poll = TimingWrapper(conn.poll)

        self.assertEqual(poll(), False)
        self.assertTimingAlmostEqual(poll.elapsed, 0)

        self.assertEqual(poll(-1), False)
        self.assertTimingAlmostEqual(poll.elapsed, 0)

        self.assertEqual(poll(TIMEOUT1), False)
        self.assertTimingAlmostEqual(poll.elapsed, TIMEOUT1)

        conn.send(None)
        time.sleep(0.1)

        self.assertEqual(poll(TIMEOUT1), True)
        self.assertTimingAlmostEqual(poll.elapsed, 0)

        self.assertEqual(conn.recv(), None)

        really_big_msg = latin("X") * (1024 * 1024 * 16)  # 16Mb
        conn.send_bytes(really_big_msg)
        self.assertEqual(conn.recv_bytes(), really_big_msg)

        conn.send_bytes(SENTINEL)  # tell child to quit
        child_conn.close()

        if self.TYPE == "processes":
            self.assertEqual(conn.readable, True)
            self.assertEqual(conn.writable, True)
            self.assertRaises(EOFError, conn.recv)
            self.assertRaises(EOFError, conn.recv_bytes)

        p.join()

    def test_duplex_false(self):
        reader, writer = self.Pipe(duplex=False)
        self.assertEqual(writer.send(1), None)
        self.assertEqual(reader.recv(), 1)
        if self.TYPE == "processes":
            self.assertEqual(reader.readable, True)
            self.assertEqual(reader.writable, False)
            self.assertEqual(writer.readable, False)
            self.assertEqual(writer.writable, True)
            self.assertRaises(OSError, reader.send, 2)
            self.assertRaises(OSError, writer.recv)
            self.assertRaises(OSError, writer.poll)

    def test_spawn_close(self):
        # We test that a pipe connection can be closed by parent
        # process immediately after child is spawned.  On Windows this
        # would have sometimes failed on old versions because
        # child_conn would be closed before the child got a chance to
        # duplicate it.
        conn, child_conn = self.Pipe()

        p = self.Process(target=self._echo, args=(child_conn,))
        p.daemon = True
        p.start()
        child_conn.close()  # this might complete before child initializes

        msg = latin("hello")
        conn.send_bytes(msg)
        self.assertEqual(conn.recv_bytes(), msg)

        conn.send_bytes(SENTINEL)
        conn.close()
        p.join()

    def test_sendbytes(self):
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        msg = latin("abcdefghijklmnopqrstuvwxyz")
        a, b = self.Pipe()

        a.send_bytes(msg)
        self.assertEqual(b.recv_bytes(), msg)

        a.send_bytes(msg, 5)
        self.assertEqual(b.recv_bytes(), msg[5:])

        a.send_bytes(msg, 7, 8)
        self.assertEqual(b.recv_bytes(), msg[7 : 7 + 8])

        a.send_bytes(msg, 26)
        self.assertEqual(b.recv_bytes(), latin(""))

        a.send_bytes(msg, 26, 0)
        self.assertEqual(b.recv_bytes(), latin(""))

        self.assertRaises(ValueError, a.send_bytes, msg, 27)

        self.assertRaises(ValueError, a.send_bytes, msg, 22, 5)

        self.assertRaises(ValueError, a.send_bytes, msg, 26, 1)

        self.assertRaises(ValueError, a.send_bytes, msg, -1)

        self.assertRaises(ValueError, a.send_bytes, msg, 4, -1)

    @classmethod
    def _is_fd_assigned(cls, fd):
        try:
            os.fstat(fd)
        except OSError as e:
            if e.errno == errno.EBADF:
                return False
            raise
        else:
            return True

    @classmethod
    def _writefd(cls, conn, data, create_dummy_fds=False):
        if create_dummy_fds:
            for i in range(0, 256):
                if not cls._is_fd_assigned(i):
                    os.dup2(conn.fileno(), i)
        fd = reduction.recv_handle(conn)
        if msvcrt:
            fd = msvcrt.open_osfhandle(fd, os.O_WRONLY)
        os.write(fd, data)
        os.close(fd)

    @unittest.skip("bug filed PE-40895")
    @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction")
    def test_fd_transfer(self):
        if self.TYPE != "processes":
            self.skipTest("only makes sense with processes")
        conn, child_conn = self.Pipe(duplex=True)

        p = self.Process(target=self._writefd, args=(child_conn, b"foo"))
        p.daemon = True
        p.start()
        self.addCleanup(test.support.unlink, test.support.TESTFN)
        with open(test.support.TESTFN, "wb") as f:
            fd = f.fileno()
            if msvcrt:
                fd = msvcrt.get_osfhandle(fd)
            reduction.send_handle(conn, fd, p.pid)
        p.join()
        with open(test.support.TESTFN, "rb") as f:
            self.assertEqual(f.read(), b"foo")

    @unittest.skip("bug filed PE-40896")
    @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction")
    @unittest.skipIf(sys.platform == "win32", "test semantics don't make sense on Windows")
    @unittest.skipIf(MAXFD <= 256, "largest assignable fd number is too small")
    @unittest.skipUnless(hasattr(os, "dup2"), "test needs os.dup2()")
    def test_large_fd_transfer(self):
        # With fd > 256 (issue #11657)
        if self.TYPE != "processes":
            self.skipTest("only makes sense with processes")
        conn, child_conn = self.Pipe(duplex=True)

        p = self.Process(target=self._writefd, args=(child_conn, b"bar", True))
        p.daemon = True
        p.start()
        self.addCleanup(test.support.unlink, test.support.TESTFN)
        with open(test.support.TESTFN, "wb") as f:
            fd = f.fileno()
            for newfd in range(256, MAXFD):
                if not self._is_fd_assigned(newfd):
                    break
            else:
                self.fail("could not find an unassigned large file descriptor")
            os.dup2(fd, newfd)
            try:
                reduction.send_handle(conn, newfd, p.pid)
            finally:
                os.close(newfd)
        p.join()
        with open(test.support.TESTFN, "rb") as f:
            self.assertEqual(f.read(), b"bar")

    @classmethod
    def _send_data_without_fd(self, conn):
        os.write(conn.fileno(), b"\0")

    @unittest.skip("bug filed PE-40896")
    @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction")
    @unittest.skipIf(sys.platform == "win32", "doesn't make sense on Windows")
    def test_missing_fd_transfer(self):
        # Check that exception is raised when received data is not
        # accompanied by a file descriptor in ancillary data.
        if self.TYPE != "processes":
            self.skipTest("only makes sense with processes")
        conn, child_conn = self.Pipe(duplex=True)

        p = self.Process(target=self._send_data_without_fd, args=(child_conn,))
        p.daemon = True
        p.start()
        self.assertRaises(RuntimeError, reduction.recv_handle, conn)
        p.join()

    def test_context(self):
        a, b = self.Pipe()

        with a, b:
            a.send(1729)
            self.assertEqual(b.recv(), 1729)
            if self.TYPE == "processes":
                self.assertFalse(a.closed)
                self.assertFalse(b.closed)

        if self.TYPE == "processes":
            self.assertTrue(a.closed)
            self.assertTrue(b.closed)
            self.assertRaises(OSError, a.recv)
            self.assertRaises(OSError, b.recv)

#
# Test of sending connection and socket objects between processes
#


@unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction")
@hashlib_helper.requires_hashdigest("md5")
class WithProcessesTestPicklingConnections(BaseTestCase, ProcessesMixin, unittest.TestCase):

    ALLOWED_TYPES = ("processes",)

    @classmethod
    def tearDownClass(cls):
        from multiprocessing import resource_sharer

        resource_sharer.stop(timeout=support.LONG_TIMEOUT)

    @classmethod
    def _listener(cls, conn, families):
        for fam in families:
            l = cls.connection.Listener(family=fam)
            conn.send(l.address)
            new_conn = l.accept()
            conn.send(new_conn)
            new_conn.close()
            l.close()

        l = socket.create_server((socket_helper.HOST, 0))
        conn.send(l.getsockname())
        new_conn, addr = l.accept()
        conn.send(new_conn)
        new_conn.close()
        l.close()

        conn.recv()

    @classmethod
    def _remote(cls, conn):
        for (address, msg) in iter(conn.recv, None):
            client = cls.connection.Client(address)
            client.send(msg.upper())
            client.close()

        address, msg = conn.recv()
        client = socket.socket()
        client.connect(address)
        client.sendall(msg.upper())
        client.close()

        conn.close()

    @unittest.skip(f"DRAGON: Broken")
    def test_pickling(self):
        families = self.connection.families

        lconn, lconn0 = self.Pipe()
        lp = self.Process(target=self._listener, args=(lconn0, families))
        lp.daemon = True
        lp.start()
        lconn0.close()

        rconn, rconn0 = self.Pipe()
        rp = self.Process(target=self._remote, args=(rconn0,))
        rp.daemon = True
        rp.start()
        rconn0.close()

        for fam in families:
            msg = ("This connection uses family %s" % fam).encode("ascii")
            address = lconn.recv()
            rconn.send((address, msg))
            new_conn = lconn.recv()
            self.assertEqual(new_conn.recv(), msg.upper())

        rconn.send(None)

        msg = latin("This connection uses a normal socket")
        address = lconn.recv()
        rconn.send((address, msg))
        new_conn = lconn.recv()
        buf = []
        while True:
            s = new_conn.recv(100)
            if not s:
                break
            buf.append(s)
        buf = b"".join(buf)
        self.assertEqual(buf, msg.upper())
        new_conn.close()

        lconn.send(None)

        rconn.close()
        lconn.close()

        lp.join()
        rp.join()

    @classmethod
    def child_access(cls, conn):
        w = conn.recv()
        w.send("all is well")
        w.close()

        r = conn.recv()
        msg = r.recv()
        conn.send(msg * 2)

        conn.close()

    def test_access(self):
        # On Windows, if we do not specify a destination pid when
        # using DupHandle then we need to be careful to use the
        # correct access flags for DuplicateHandle(), or else
        # DupHandle.detach() will raise PermissionError.  For example,
        # for a read only pipe handle we should use
        # access=FILE_GENERIC_READ.  (Unfortunately
        # DUPLICATE_SAME_ACCESS does not work.)
        conn, child_conn = self.Pipe()
        p = self.Process(target=self.child_access, args=(child_conn,))
        p.daemon = True
        p.start()
        child_conn.close()

        r, w = self.Pipe(duplex=False)
        conn.send(w)
        w.close()
        self.assertEqual(r.recv(), "all is well")
        r.close()

        r, w = self.Pipe(duplex=False)
        conn.send(r)
        r.close()
        w.send("foobar")
        w.close()
        self.assertEqual(conn.recv(), "foobar" * 2)

        p.join()


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestConnection(BaseTestCase, ThreadsMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    @classmethod
    def _echo(cls, conn):
        for msg in iter(conn.recv_bytes, SENTINEL):
            conn.send_bytes(msg)
        conn.close()

    def test_connection(self):
        conn, child_conn = self.Pipe()

        p = self.Process(target=self._echo, args=(child_conn,))
        p.daemon = True
        p.start()

        seq = [1, 2.25, None]
        msg = latin("hello world")
        longmsg = msg * 10
        arr = array.array("i", list(range(4)))

        if self.TYPE == "processes":
            self.assertEqual(type(conn.fileno()), int)

        self.assertEqual(conn.send(seq), None)
        self.assertEqual(conn.recv(), seq)

        self.assertEqual(conn.send_bytes(msg), None)
        self.assertEqual(conn.recv_bytes(), msg)

        if self.TYPE == "processes":
            buffer = array.array("i", [0] * 10)
            expected = list(arr) + [0] * (10 - len(arr))
            self.assertEqual(conn.send_bytes(arr), None)
            self.assertEqual(conn.recv_bytes_into(buffer), len(arr) * buffer.itemsize)
            self.assertEqual(list(buffer), expected)

            buffer = array.array("i", [0] * 10)
            expected = [0] * 3 + list(arr) + [0] * (10 - 3 - len(arr))
            self.assertEqual(conn.send_bytes(arr), None)
            self.assertEqual(conn.recv_bytes_into(buffer, 3 * buffer.itemsize), len(arr) * buffer.itemsize)
            self.assertEqual(list(buffer), expected)

            buffer = bytearray(latin(" " * 40))
            self.assertEqual(conn.send_bytes(longmsg), None)
            try:
                res = conn.recv_bytes_into(buffer)
            except multiprocessing.BufferTooShort as e:
                self.assertEqual(e.args, (longmsg,))
            else:
                self.fail("expected BufferTooShort, got %s" % res)

        poll = TimingWrapper(conn.poll)

        self.assertEqual(poll(), False)
        self.assertTimingAlmostEqual(poll.elapsed, 0)

        self.assertEqual(poll(-1), False)
        self.assertTimingAlmostEqual(poll.elapsed, 0)

        self.assertEqual(poll(TIMEOUT1), False)
        self.assertTimingAlmostEqual(poll.elapsed, TIMEOUT1)

        conn.send(None)
        time.sleep(0.1)

        self.assertEqual(poll(TIMEOUT1), True)
        self.assertTimingAlmostEqual(poll.elapsed, 0)

        self.assertEqual(conn.recv(), None)

        really_big_msg = latin("X") * (1024 * 1024 * 16)  # 16Mb
        conn.send_bytes(really_big_msg)
        self.assertEqual(conn.recv_bytes(), really_big_msg)

        conn.send_bytes(SENTINEL)  # tell child to quit
        child_conn.close()

        if self.TYPE == "processes":
            self.assertEqual(conn.readable, True)
            self.assertEqual(conn.writable, True)
            self.assertRaises(EOFError, conn.recv)
            self.assertRaises(EOFError, conn.recv_bytes)

        p.join()

    def test_duplex_false(self):
        reader, writer = self.Pipe(duplex=False)
        self.assertEqual(writer.send(1), None)
        self.assertEqual(reader.recv(), 1)
        if self.TYPE == "processes":
            self.assertEqual(reader.readable, True)
            self.assertEqual(reader.writable, False)
            self.assertEqual(writer.readable, False)
            self.assertEqual(writer.writable, True)
            self.assertRaises(OSError, reader.send, 2)
            self.assertRaises(OSError, writer.recv)
            self.assertRaises(OSError, writer.poll)

    def test_spawn_close(self):
        # We test that a pipe connection can be closed by parent
        # process immediately after child is spawned.  On Windows this
        # would have sometimes failed on old versions because
        # child_conn would be closed before the child got a chance to
        # duplicate it.
        conn, child_conn = self.Pipe()

        p = self.Process(target=self._echo, args=(child_conn,))
        p.daemon = True
        p.start()
        child_conn.close()  # this might complete before child initializes

        msg = latin("hello")
        conn.send_bytes(msg)
        self.assertEqual(conn.recv_bytes(), msg)

        conn.send_bytes(SENTINEL)
        conn.close()
        p.join()

    def test_sendbytes(self):
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        msg = latin("abcdefghijklmnopqrstuvwxyz")
        a, b = self.Pipe()

        a.send_bytes(msg)
        self.assertEqual(b.recv_bytes(), msg)

        a.send_bytes(msg, 5)
        self.assertEqual(b.recv_bytes(), msg[5:])

        a.send_bytes(msg, 7, 8)
        self.assertEqual(b.recv_bytes(), msg[7 : 7 + 8])

        a.send_bytes(msg, 26)
        self.assertEqual(b.recv_bytes(), latin(""))

        a.send_bytes(msg, 26, 0)
        self.assertEqual(b.recv_bytes(), latin(""))

        self.assertRaises(ValueError, a.send_bytes, msg, 27)

        self.assertRaises(ValueError, a.send_bytes, msg, 22, 5)

        self.assertRaises(ValueError, a.send_bytes, msg, 26, 1)

        self.assertRaises(ValueError, a.send_bytes, msg, -1)

        self.assertRaises(ValueError, a.send_bytes, msg, 4, -1)

    @classmethod
    def _is_fd_assigned(cls, fd):
        try:
            os.fstat(fd)
        except OSError as e:
            if e.errno == errno.EBADF:
                return False
            raise
        else:
            return True

    @classmethod
    def _writefd(cls, conn, data, create_dummy_fds=False):
        if create_dummy_fds:
            for i in range(0, 256):
                if not cls._is_fd_assigned(i):
                    os.dup2(conn.fileno(), i)
        fd = reduction.recv_handle(conn)
        if msvcrt:
            fd = msvcrt.open_osfhandle(fd, os.O_WRONLY)
        os.write(fd, data)
        os.close(fd)

    @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction")
    def test_fd_transfer(self):
        if self.TYPE != "processes":
            self.skipTest("only makes sense with processes")
        conn, child_conn = self.Pipe(duplex=True)

        p = self.Process(target=self._writefd, args=(child_conn, b"foo"))
        p.daemon = True
        p.start()
        self.addCleanup(test.support.unlink, test.support.TESTFN)
        with open(test.support.TESTFN, "wb") as f:
            fd = f.fileno()
            if msvcrt:
                fd = msvcrt.get_osfhandle(fd)
            reduction.send_handle(conn, fd, p.pid)
        p.join()
        with open(test.support.TESTFN, "rb") as f:
            self.assertEqual(f.read(), b"foo")

    @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction")
    @unittest.skipIf(sys.platform == "win32", "test semantics don't make sense on Windows")
    @unittest.skipIf(MAXFD <= 256, "largest assignable fd number is too small")
    @unittest.skipUnless(hasattr(os, "dup2"), "test needs os.dup2()")
    def test_large_fd_transfer(self):
        # With fd > 256 (issue #11657)
        if self.TYPE != "processes":
            self.skipTest("only makes sense with processes")
        conn, child_conn = self.Pipe(duplex=True)

        p = self.Process(target=self._writefd, args=(child_conn, b"bar", True))
        p.daemon = True
        p.start()
        self.addCleanup(test.support.unlink, test.support.TESTFN)
        with open(test.support.TESTFN, "wb") as f:
            fd = f.fileno()
            for newfd in range(256, MAXFD):
                if not self._is_fd_assigned(newfd):
                    break
            else:
                self.fail("could not find an unassigned large file descriptor")
            os.dup2(fd, newfd)
            try:
                reduction.send_handle(conn, newfd, p.pid)
            finally:
                os.close(newfd)
        p.join()
        with open(test.support.TESTFN, "rb") as f:
            self.assertEqual(f.read(), b"bar")

    @classmethod
    def _send_data_without_fd(self, conn):
        os.write(conn.fileno(), b"\0")

    @unittest.skipUnless(HAS_REDUCTION, "test needs multiprocessing.reduction")
    @unittest.skipIf(sys.platform == "win32", "doesn't make sense on Windows")
    def test_missing_fd_transfer(self):
        # Check that exception is raised when received data is not
        # accompanied by a file descriptor in ancillary data.
        if self.TYPE != "processes":
            self.skipTest("only makes sense with processes")
        conn, child_conn = self.Pipe(duplex=True)

        p = self.Process(target=self._send_data_without_fd, args=(child_conn,))
        p.daemon = True
        p.start()
        self.assertRaises(RuntimeError, reduction.recv_handle, conn)
        p.join()

    def test_context(self):
        a, b = self.Pipe()

        with a, b:
            a.send(1729)
            self.assertEqual(b.recv(), 1729)
            if self.TYPE == "processes":
                self.assertFalse(a.closed)
                self.assertFalse(b.closed)

        if self.TYPE == "processes":
            self.assertTrue(a.closed)
            self.assertTrue(b.closed)
            self.assertRaises(OSError, a.recv)
            self.assertRaises(OSError, b.recv)


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_connection.py
DragonHPC/dragon
""" Test mulitprocessing.Condition
"""
import unittest

from test.support import hashlib_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing

from common import (
    BaseTestCase,
    ManagerMixin,
    setUpModule,
    tearDownModule,
)


#
#
#


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestContainers(BaseTestCase, ManagerMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('manager',)

    def test_list(self):
        a = self.list(list(range(10)))
        self.assertEqual(a[:], list(range(10)))

        b = self.list()
        self.assertEqual(b[:], [])

        b.extend(list(range(5)))
        self.assertEqual(b[:], list(range(5)))

        self.assertEqual(b[2], 2)
        self.assertEqual(b[2:10], [2, 3, 4])

        b *= 2
        self.assertEqual(b[:], [0, 1, 2, 3, 4, 0, 1, 2, 3, 4])

        self.assertEqual(b + [5, 6], [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6])

        self.assertEqual(a[:], list(range(10)))

        d = [a, b]
        e = self.list(d)
        self.assertEqual(
            [element[:] for element in e], [[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], [0, 1, 2, 3, 4, 0, 1, 2, 3, 4]]
        )

        f = self.list([a])
        a.append("hello")
        self.assertEqual(f[0][:], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, "hello"])

    def test_list_iter(self):
        a = self.list(list(range(10)))
        it = iter(a)
        self.assertEqual(list(it), list(range(10)))
        self.assertEqual(list(it), [])  # exhausted
        # list modified during iteration
        it = iter(a)
        a[0] = 100
        self.assertEqual(next(it), 100)

    def test_list_proxy_in_list(self):
        a = self.list([self.list(range(3)) for _i in range(3)])
        self.assertEqual([inner[:] for inner in a], [[0, 1, 2]] * 3)

        a[0][-1] = 55
        self.assertEqual(a[0][:], [0, 1, 55])
        for i in range(1, 3):
            self.assertEqual(a[i][:], [0, 1, 2])

        self.assertEqual(a[1].pop(), 2)
        self.assertEqual(len(a[1]), 2)
        for i in range(0, 3, 2):
            self.assertEqual(len(a[i]), 3)

        del a

        b = self.list()
        b.append(b)
        del b

    def test_dict(self):
        d = self.dict()
        indices = list(range(65, 70))
        for i in indices:
            d[i] = chr(i)
        self.assertEqual(d.copy(), dict((i, chr(i)) for i in indices))
        self.assertEqual(sorted(d.keys()), indices)
        self.assertEqual(sorted(d.values()), [chr(i) for i in indices])
        self.assertEqual(sorted(d.items()), [(i, chr(i)) for i in indices])

    def test_dict_iter(self):
        d = self.dict()
        indices = list(range(65, 70))
        for i in indices:
            d[i] = chr(i)
        it = iter(d)
        self.assertEqual(list(it), indices)
        self.assertEqual(list(it), [])  # exhausted
        # dictionary changed size during iteration
        it = iter(d)
        d.clear()
        self.assertRaises(RuntimeError, next, it)

    def test_dict_proxy_nested(self):
        pets = self.dict(ferrets=2, hamsters=4)
        supplies = self.dict(water=10, feed=3)
        d = self.dict(pets=pets, supplies=supplies)

        self.assertEqual(supplies["water"], 10)
        self.assertEqual(d["supplies"]["water"], 10)

        d["supplies"]["blankets"] = 5
        self.assertEqual(supplies["blankets"], 5)
        self.assertEqual(d["supplies"]["blankets"], 5)

        d["supplies"]["water"] = 7
        self.assertEqual(supplies["water"], 7)
        self.assertEqual(d["supplies"]["water"], 7)

        del pets
        del supplies
        self.assertEqual(d["pets"]["ferrets"], 2)
        d["supplies"]["blankets"] = 11
        self.assertEqual(d["supplies"]["blankets"], 11)

        pets = d["pets"]
        supplies = d["supplies"]
        supplies["water"] = 7
        self.assertEqual(supplies["water"], 7)
        self.assertEqual(d["supplies"]["water"], 7)

        d.clear()
        self.assertEqual(len(d), 0)
        self.assertEqual(supplies["water"], 7)
        self.assertEqual(pets["hamsters"], 4)

        l = self.list([pets, supplies])
        l[0]["marmots"] = 1
        self.assertEqual(pets["marmots"], 1)
        self.assertEqual(l[0]["marmots"], 1)

        del pets
        del supplies
        self.assertEqual(l[0]["marmots"], 1)

        outer = self.list([[88, 99], l])
        self.assertIsInstance(outer[0], list)  # Not a ListProxy
        self.assertEqual(outer[-1][-1]["feed"], 3)

    def test_namespace(self):
        n = self.Namespace()
        n.name = "Bob"
        n.job = "Builder"
        n._hidden = "hidden"
        self.assertEqual((n.name, n.job), ("Bob", "Builder"))
        del n.job
        self.assertEqual(str(n), "Namespace(name='Bob')")
        self.assertTrue(hasattr(n, "name"))
        self.assertTrue(not hasattr(n, "job"))


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_containers.py
DragonHPC/dragon
""" Test mulitprocessing event
"""
import unittest
import time
from test.support import hashlib_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    TimingWrapper,
    setUpModule,
    tearDownModule,
    TIMEOUT1,
    TIMEOUT2,
)


class WithProcessesTestEvent(BaseTestCase, ProcessesMixin, unittest.TestCase):
    @classmethod
    def _test_event(cls, event):
        time.sleep(TIMEOUT2)
        event.set()

    def test_event(self):
        event = self.Event()
        wait = TimingWrapper(event.wait)

        # Removed temporarily, due to API shear, this does not
        # work with threading._Event objects. is_set == isSet
        self.assertEqual(event.is_set(), False)

        # Removed, threading.Event.wait() will return the value of the __flag
        # instead of None. API Shear with the semaphore backed mp.Event
        self.assertEqual(wait(0.0), False)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        self.assertEqual(wait(TIMEOUT1), False)
        self.assertTimingAlmostEqual(wait.elapsed, TIMEOUT1)

        event.set()

        # See note above on the API differences
        self.assertEqual(event.is_set(), True)
        self.assertEqual(wait(), True)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        self.assertEqual(wait(TIMEOUT1), True)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        # self.assertEqual(event.is_set(), True)

        event.clear()

        # self.assertEqual(event.is_set(), False)

        p = self.Process(target=self._test_event, args=(event,))
        p.daemon = True
        p.start()
        self.assertEqual(wait(), True)
        p.join()


#


@hashlib_helper.requires_hashdigest("md5")
@unittest.skip("DRAGON: not implemented")
class WithManagerTestEvent(BaseTestCase, ManagerMixin, unittest.TestCase):
    @classmethod
    def _test_event(cls, event):
        time.sleep(TIMEOUT2)
        event.set()

    def test_event(self):
        event = self.Event()
        wait = TimingWrapper(event.wait)

        # Removed temporarily, due to API shear, this does not
        # work with threading._Event objects. is_set == isSet
        self.assertEqual(event.is_set(), False)

        # Removed, threading.Event.wait() will return the value of the __flag
        # instead of None. API Shear with the semaphore backed mp.Event
        self.assertEqual(wait(0.0), False)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        self.assertEqual(wait(TIMEOUT1), False)
        self.assertTimingAlmostEqual(wait.elapsed, TIMEOUT1)

        event.set()

        # See note above on the API differences
        self.assertEqual(event.is_set(), True)
        self.assertEqual(wait(), True)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        self.assertEqual(wait(TIMEOUT1), True)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        # self.assertEqual(event.is_set(), True)

        event.clear()

        # self.assertEqual(event.is_set(), False)

        p = self.Process(target=self._test_event, args=(event,))
        p.daemon = True
        p.start()
        self.assertEqual(wait(), True)
        p.join()


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestEvent(BaseTestCase, ThreadsMixin, unittest.TestCase):
    @classmethod
    def _test_event(cls, event):
        time.sleep(TIMEOUT2)
        event.set()

    def test_event(self):
        event = self.Event()
        wait = TimingWrapper(event.wait)

        # Removed temporarily, due to API shear, this does not
        # work with threading._Event objects. is_set == isSet
        self.assertEqual(event.is_set(), False)

        # Removed, threading.Event.wait() will return the value of the __flag
        # instead of None. API Shear with the semaphore backed mp.Event
        self.assertEqual(wait(0.0), False)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        self.assertEqual(wait(TIMEOUT1), False)
        self.assertTimingAlmostEqual(wait.elapsed, TIMEOUT1)

        event.set()

        # See note above on the API differences
        self.assertEqual(event.is_set(), True)
        self.assertEqual(wait(), True)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        self.assertEqual(wait(TIMEOUT1), True)
        self.assertTimingAlmostEqual(wait.elapsed, 0.0)
        # self.assertEqual(event.is_set(), True)

        event.clear()

        # self.assertEqual(event.is_set(), False)

        p = self.Process(target=self._test_event, args=(event,))
        p.daemon = True
        p.start()
        self.assertEqual(wait(), True)
        p.join()


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_event.py
DragonHPC/dragon
"""Test Finalize, Import * and Logging
"""
import sys
import time
import os
import gc
import random

import unittest
import test.support

import threading

import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing import util

import logging

try:
    from ctypes import c_int
except ImportError:
    Structure = object
    c_int = c_double = c_longlong = None

try:
    from multiprocessing import reduction

    HAS_REDUCTION = reduction.HAVE_SEND_HANDLE
except ImportError:
    HAS_REDUCTION = False

from common import (
    BaseTestCase,
    ProcessesMixin,
    LOG_LEVEL,
    setUpModule,
    tearDownModule,
)

#
#
#


@unittest.skip("DRAGON: not implemented")
class WithProcessesTestFinalize(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    def setUp(self):
        self.registry_backup = util._finalizer_registry.copy()
        util._finalizer_registry.clear()

    def tearDown(self):
        self.assertFalse(util._finalizer_registry)
        util._finalizer_registry.update(self.registry_backup)

    @classmethod
    def _test_finalize(cls, conn):
        class Foo(object):
            pass

        a = Foo()
        util.Finalize(a, conn.send, args=("a",))
        del a  # triggers callback for a

        b = Foo()
        close_b = util.Finalize(b, conn.send, args=("b",))
        close_b()  # triggers callback for b
        close_b()  # does nothing because callback has already been called
        del b  # does nothing because callback has already been called

        c = Foo()
        util.Finalize(c, conn.send, args=("c",))

        d10 = Foo()
        util.Finalize(d10, conn.send, args=("d10",), exitpriority=1)

        d01 = Foo()
        util.Finalize(d01, conn.send, args=("d01",), exitpriority=0)
        d02 = Foo()
        util.Finalize(d02, conn.send, args=("d02",), exitpriority=0)
        d03 = Foo()
        util.Finalize(d03, conn.send, args=("d03",), exitpriority=0)

        util.Finalize(None, conn.send, args=("e",), exitpriority=-10)

        util.Finalize(None, conn.send, args=("STOP",), exitpriority=-100)

        # call multiprocessing's cleanup function then exit process without
        # garbage collecting locals
        util._exit_function()
        conn.close()
        os._exit(0)

    def test_finalize(self):
        conn, child_conn = self.Pipe()

        p = self.Process(target=self._test_finalize, args=(child_conn,))
        p.daemon = True
        p.start()
        p.join()

        result = [obj for obj in iter(conn.recv, "STOP")]
        self.assertEqual(result, ["a", "b", "d10", "d03", "d02", "d01", "e"])

    def test_thread_safety(self):
        # bpo-24484: _run_finalizers() should be thread-safe
        def cb():
            pass

        class Foo(object):
            def __init__(self):
                self.ref = self  # create reference cycle
                # insert finalizer at random key
                util.Finalize(self, cb, exitpriority=random.randint(1, 100))

        finish = False
        exc = None

        def run_finalizers():
            nonlocal exc
            while not finish:
                time.sleep(random.random() * 1e-1)
                try:
                    # A GC run will eventually happen during this,
                    # collecting stale Foo's and mutating the registry
                    util._run_finalizers()
                except Exception as e:
                    exc = e

        def make_finalizers():
            nonlocal exc
            d = {}
            while not finish:
                try:
                    # Old Foo's get gradually replaced and later
                    # collected by the GC (because of the cyclic ref)
                    d[random.getrandbits(5)] = {Foo() for i in range(10)}
                except Exception as e:
                    exc = e
                    d.clear()

        old_interval = sys.getswitchinterval()
        old_threshold = gc.get_threshold()
        try:
            sys.setswitchinterval(1e-6)
            gc.set_threshold(5, 5, 5)
            threads = [threading.Thread(target=run_finalizers), threading.Thread(target=make_finalizers)]
            with test.support.start_threads(threads):
                time.sleep(4.0)  # Wait a bit to trigger race condition
                finish = True
            if exc is not None:
                raise exc
        finally:
            sys.setswitchinterval(old_interval)
            gc.set_threshold(*old_threshold)
            gc.collect()  # Collect remaining Foo's


#
# Test that from ... import * works for each module
#


class WithProcessesTestImportStar(unittest.TestCase):
    def get_module_names(self):
        import glob

        folder = os.path.dirname(multiprocessing.__file__)
        pattern = os.path.join(glob.escape(folder), "*.py")
        files = glob.glob(pattern)
        modules = [os.path.splitext(os.path.split(f)[1])[0] for f in files]
        modules = ["multiprocessing." + m for m in modules]
        modules.remove("multiprocessing.__init__")
        modules.append("multiprocessing")
        return modules

    def test_import(self):
        modules = self.get_module_names()
        if sys.platform == "win32":
            modules.remove("multiprocessing.popen_fork")
            modules.remove("multiprocessing.popen_forkserver")
            modules.remove("multiprocessing.popen_spawn_posix")
        else:
            modules.remove("multiprocessing.popen_spawn_win32")
            if not HAS_REDUCTION:
                modules.remove("multiprocessing.popen_forkserver")

        if c_int is None:
            # This module requires _ctypes
            modules.remove("multiprocessing.sharedctypes")

        for name in modules:
            __import__(name)
            mod = sys.modules[name]
            self.assertTrue(hasattr(mod, "__all__"), name)

            for attr in mod.__all__:
                self.assertTrue(hasattr(mod, attr), "%r does not have attribute %r" % (mod, attr))


#
# Quick test that logging works -- does not test logging output
#


class WithProcessesTestLogging(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    def test_enable_logging(self):
        logger = multiprocessing.get_logger()
        logger.setLevel(util.SUBWARNING)
        self.assertTrue(logger is not None)
        logger.debug("this will not be printed")
        logger.info("nor will this")
        logger.setLevel(LOG_LEVEL)

    @classmethod
    def _test_level(cls, conn):
        logger = multiprocessing.get_logger()
        conn.send(logger.getEffectiveLevel())

    @unittest.skip(f"DRAGON: Logging currently incomplete PE-41692")
    def test_level(self):
        LEVEL1 = 32
        LEVEL2 = 37

        logger = multiprocessing.get_logger()
        root_logger = logging.getLogger()
        root_level = root_logger.level

        reader, writer = multiprocessing.Pipe(duplex=False)

        logger.setLevel(LEVEL1)
        p = self.Process(target=self._test_level, args=(writer,))
        p.start()
        self.assertEqual(LEVEL1, reader.recv())
        p.join()
        p.close()

        logger.setLevel(logging.NOTSET)
        root_logger.setLevel(LEVEL2)
        p = self.Process(target=self._test_level, args=(writer,))
        p.start()
        self.assertEqual(LEVEL2, reader.recv())
        p.join()
        p.close()

        root_logger.setLevel(root_level)
        logger.setLevel(level=LOG_LEVEL)


# class _TestLoggingProcessName(BaseTestCase):
#
#     def handle(self, record):
#         assert record.processName == multiprocessing.current_process().name
#         self.__handled = True
#
#     def test_logging(self):
#         handler = logging.Handler()
#         handler.handle = self.handle
#         self.__handled = False
#         # Bypass getLogger() and side-effects
#         logger = logging.getLoggerClass()(
#                 'multiprocessing.test.TestLoggingProcessName')
#         logger.addHandler(handler)
#         logger.propagate = False
#
#         logger.warn('foo')
#         assert self.__handled


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_finalize.py
DragonHPC/dragon
"""Test multiprocessing connection, listener, client, pickling
"""
import unittest
import time
import socket

from test import support
from test.support import hashlib_helper
from test.support import socket_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing import util

try:
    from multiprocessing import reduction

    HAS_REDUCTION = reduction.HAVE_SEND_HANDLE
except ImportError:
    HAS_REDUCTION = False

from common import (
    BaseTestCase,
    ProcessesMixin,
    ThreadsMixin,
    latin,
    setUpModule,
    tearDownModule,
)


#
#
#

SENTINEL = latin("")


@unittest.skip(f"DRAGON: not implemented")
class WithProcessesTestListener(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    def test_multiple_bind(self):
        for family in self.connection.families:
            l = self.connection.Listener(family=family)
            self.addCleanup(l.close)
            self.assertRaises(OSError, self.connection.Listener, l.address, family)

    def test_context(self):
        with self.connection.Listener() as l:
            with self.connection.Client(l.address) as c:
                with l.accept() as d:
                    c.send(1729)
                    self.assertEqual(d.recv(), 1729)

        if self.TYPE == "processes":
            self.assertRaises(OSError, l.accept)

    @unittest.skipUnless(util.abstract_sockets_supported, "test needs abstract socket support")
    def test_abstract_socket(self):
        with self.connection.Listener("\0something") as listener:
            with self.connection.Client(listener.address) as client:
                with listener.accept() as d:
                    client.send(1729)
                    self.assertEqual(d.recv(), 1729)

        if self.TYPE == "processes":
            self.assertRaises(OSError, listener.accept)


@unittest.skip(f"DRAGON: not implemented")
class WithProcessesTestListenerClient(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    @classmethod
    def _test(cls, address):
        conn = cls.connection.Client(address)
        conn.send("hello")
        conn.close()

    def test_listener_client(self):
        for family in self.connection.families:
            l = self.connection.Listener(family=family)
            p = self.Process(target=self._test, args=(l.address,))
            p.daemon = True
            p.start()
            conn = l.accept()
            self.assertEqual(conn.recv(), "hello")
            p.join()
            l.close()

    def test_issue14725(self):
        l = self.connection.Listener()
        p = self.Process(target=self._test, args=(l.address,))
        p.daemon = True
        p.start()
        time.sleep(1)
        # On Windows the client process should by now have connected,
        # written data and closed the pipe handle by now.  This causes
        # ConnectNamdedPipe() to fail with ERROR_NO_DATA.  See Issue
        # 14725.
        conn = l.accept()
        self.assertEqual(conn.recv(), "hello")
        conn.close()
        p.join()
        l.close()

    def test_issue16955(self):
        for fam in self.connection.families:
            l = self.connection.Listener(family=fam)
            c = self.connection.Client(l.address)
            a = l.accept()
            a.send_bytes(b"hello")
            self.assertTrue(c.poll(1))
            a.close()
            c.close()
            l.close()


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestListenerClient(BaseTestCase, ThreadsMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    @classmethod
    def _test(cls, address):
        conn = cls.connection.Client(address)
        conn.send("hello")
        conn.close()

    def test_listener_client(self):
        for family in self.connection.families:
            l = self.connection.Listener(family=family)
            p = self.Process(target=self._test, args=(l.address,))
            p.daemon = True
            p.start()
            conn = l.accept()
            self.assertEqual(conn.recv(), "hello")
            p.join()
            l.close()

    def test_issue14725(self):
        l = self.connection.Listener()
        p = self.Process(target=self._test, args=(l.address,))
        p.daemon = True
        p.start()
        time.sleep(1)
        # On Windows the client process should by now have connected,
        # written data and closed the pipe handle by now.  This causes
        # ConnectNamdedPipe() to fail with ERROR_NO_DATA.  See Issue
        # 14725.
        conn = l.accept()
        self.assertEqual(conn.recv(), "hello")
        conn.close()
        p.join()
        l.close()

    def test_issue16955(self):
        for fam in self.connection.families:
            l = self.connection.Listener(family=fam)
            c = self.connection.Client(l.address)
            a = l.accept()
            a.send_bytes(b"hello")
            self.assertTrue(c.poll(1))
            a.close()
            c.close()
            l.close()


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_listener.py
DragonHPC/dragon
""" Test multiprocessing lock
"""
import unittest
from test.support import hashlib_helper

import threading

import dragon  # DRAGON import before multiprocessing

import multiprocessing

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    setUpModule,
    tearDownModule,
)

#
#
#


class WithProcessesTestLock(BaseTestCase, ProcessesMixin, unittest.TestCase):
    def test_lock(self):
        lock = self.Lock()
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(False), False)
        self.assertEqual(lock.release(), None)
        self.assertRaises((ValueError, threading.ThreadError), lock.release)

    def test_rlock(self):
        lock = self.RLock()
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.release(), None)
        self.assertEqual(lock.release(), None)
        self.assertEqual(lock.release(), None)
        self.assertRaises((AssertionError, RuntimeError), lock.release)

    def test_lock_context(self):
        with self.Lock():
            pass


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestLock(BaseTestCase, ManagerMixin, unittest.TestCase):
    def test_lock(self):
        lock = self.Lock()
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(False), False)
        self.assertEqual(lock.release(), None)
        self.assertRaises((ValueError, threading.ThreadError), lock.release)

    def test_rlock(self):
        lock = self.RLock()
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.release(), None)
        self.assertEqual(lock.release(), None)
        self.assertEqual(lock.release(), None)
        self.assertRaises((AssertionError, RuntimeError), lock.release)

    def test_lock_context(self):
        with self.Lock():
            pass


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestLock(BaseTestCase, ThreadsMixin, unittest.TestCase):
    def test_lock(self):
        lock = self.Lock()
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(False), False)
        self.assertEqual(lock.release(), None)
        self.assertRaises((ValueError, threading.ThreadError), lock.release)

    def test_rlock(self):
        lock = self.RLock()
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.acquire(), True)
        self.assertEqual(lock.release(), None)
        self.assertEqual(lock.release(), None)
        self.assertEqual(lock.release(), None)
        self.assertRaises((AssertionError, RuntimeError), lock.release)

    def test_lock_context(self):
        with self.Lock():
            pass


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_lock.py
DragonHPC/dragon
"""Test multiprocessing manager restart
"""
import unittest
import time
import os
import errno
import signal
import queue as pyqueue

from test.support import hashlib_helper
from test.support import socket_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing.managers import BaseManager

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    setUpModule,
    tearDownModule,
)

_queue = pyqueue.Queue()


def get_queue():
    return _queue


class QueueManager(BaseManager):
    """manager class used by server process"""


QueueManager.register("get_queue", callable=get_queue)


class QueueManager2(BaseManager):
    """manager class which specifies the same interface as QueueManager"""


QueueManager2.register("get_queue")

SERIALIZER = "xmlrpclib"


@hashlib_helper.requires_hashdigest("md5")
@unittest.skip("DRAGON: not implemented")
class WithProcessesTestManagerRestart(BaseTestCase, ProcessesMixin, unittest.TestCase):
    @classmethod
    def _putter(cls, address, authkey):
        manager = QueueManager(address=address, authkey=authkey, serializer=SERIALIZER)
        manager.connect()
        queue = manager.get_queue()
        queue.put("hello world")

    def test_rapid_restart(self):
        authkey = os.urandom(32)
        manager = QueueManager(address=(socket_helper.HOST, 0), authkey=authkey, serializer=SERIALIZER)
        try:
            srvr = manager.get_server()
            addr = srvr.address
            # Close the connection.Listener socket which gets opened as a part
            # of manager.get_server(). It's not needed for the test.
            srvr.listener.close()
            manager.start()

            p = self.Process(target=self._putter, args=(manager.address, authkey))
            p.start()
            p.join()
            queue = manager.get_queue()
            self.assertEqual(queue.get(), "hello world")
            del queue
        finally:
            if hasattr(manager, "shutdown"):
                manager.shutdown()

        manager = QueueManager(address=addr, authkey=authkey, serializer=SERIALIZER)
        try:
            manager.start()
            self.addCleanup(manager.shutdown)
        except OSError as e:
            if e.errno != errno.EADDRINUSE:
                raise
            # Retry after some time, in case the old socket was lingering
            # (sporadic failure on buildbots)
            time.sleep(1.0)
            manager = QueueManager(address=addr, authkey=authkey, serializer=SERIALIZER)
            if hasattr(manager, "shutdown"):
                self.addCleanup(manager.shutdown)


#
# Test of creating a customized manager class
#

from multiprocessing.managers import BaseManager, BaseProxy, RemoteError


class FooBar(object):
    def f(self):
        return "f()"

    def g(self):
        raise ValueError

    def _h(self):
        return "_h()"


def baz():
    for i in range(10):
        yield i * i


class IteratorProxy(BaseProxy):
    _exposed_ = ("__next__",)

    def __iter__(self):
        return self

    def __next__(self):
        return self._callmethod("__next__")


class MyManager(BaseManager):
    pass


MyManager.register("Foo", callable=FooBar)
MyManager.register("Bar", callable=FooBar, exposed=("f", "_h"))
MyManager.register("baz", callable=baz, proxytype=IteratorProxy)


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestMyManager(BaseTestCase, ManagerMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('manager',)

    def test_mymanager(self):
        manager = MyManager()
        manager.start()
        self.common(manager)
        manager.shutdown()

        # bpo-30356: BaseManager._finalize_manager() sends SIGTERM
        # to the manager process if it takes longer than 1 second to stop,
        # which happens on slow buildbots.
        self.assertIn(manager._process.exitcode, (0, -signal.SIGTERM))

    def test_mymanager_context(self):
        with MyManager() as manager:
            self.common(manager)
        # bpo-30356: BaseManager._finalize_manager() sends SIGTERM
        # to the manager process if it takes longer than 1 second to stop,
        # which happens on slow buildbots.
        self.assertIn(manager._process.exitcode, (0, -signal.SIGTERM))

    def test_mymanager_context_prestarted(self):
        manager = MyManager()
        manager.start()
        with manager:
            self.common(manager)
        self.assertEqual(manager._process.exitcode, 0)

    def common(self, manager):
        foo = manager.Foo()
        bar = manager.Bar()
        baz = manager.baz()

        foo_methods = [name for name in ("f", "g", "_h") if hasattr(foo, name)]
        bar_methods = [name for name in ("f", "g", "_h") if hasattr(bar, name)]

        self.assertEqual(foo_methods, ["f", "g"])
        self.assertEqual(bar_methods, ["f", "_h"])

        self.assertEqual(foo.f(), "f()")
        self.assertRaises(ValueError, foo.g)
        self.assertEqual(foo._callmethod("f"), "f()")
        self.assertRaises(RemoteError, foo._callmethod, "_h")

        self.assertEqual(bar.f(), "f()")
        self.assertEqual(bar._h(), "_h()")
        self.assertEqual(bar._callmethod("f"), "f()")
        self.assertEqual(bar._callmethod("_h"), "_h()")

        self.assertEqual(list(baz), [i * i for i in range(10)])


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestRemoteManager(BaseTestCase, ManagerMixin, unittest.TestCase):

    ALLOWED_TYPES = ("manager",)
    values = [
        "hello world",
        None,
        True,
        2.25,
        "hall\xe5 v\xe4rlden",
        "\u043f\u0440\u0438\u0432\u0456\u0442 \u0441\u0432\u0456\u0442",
        b"hall\xe5 v\xe4rlden",
    ]
    result = values[:]

    @classmethod
    def _putter(cls, address, authkey):
        manager = QueueManager2(address=address, authkey=authkey, serializer=SERIALIZER)
        manager.connect()
        queue = manager.get_queue()
        # Note that xmlrpclib will deserialize object as a list not a tuple
        queue.put(tuple(cls.values))

    def test_remote(self):
        authkey = os.urandom(32)

        manager = QueueManager(address=(socket_helper.HOST, 0), authkey=authkey, serializer=SERIALIZER)
        manager.start()
        self.addCleanup(manager.shutdown)

        p = self.Process(target=self._putter, args=(manager.address, authkey))
        p.daemon = True
        p.start()

        manager2 = QueueManager2(address=manager.address, authkey=authkey, serializer=SERIALIZER)
        manager2.connect()
        queue = manager2.get_queue()

        self.assertEqual(queue.get(), self.result)

        # Because we are using xmlrpclib for serialization instead of
        # pickle this will cause a serialization error.
        self.assertRaises(Exception, queue.put, time.sleep)

        # Make queue finalizer run before the server is stopped
        del queue


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestManagerRestart(BaseTestCase, ManagerMixin, unittest.TestCase):
    @classmethod
    def _putter(cls, address, authkey):
        manager = QueueManager(address=address, authkey=authkey, serializer=SERIALIZER)
        manager.connect()
        queue = manager.get_queue()
        queue.put("hello world")

    def test_rapid_restart(self):
        authkey = os.urandom(32)
        manager = QueueManager(address=(socket_helper.HOST, 0), authkey=authkey, serializer=SERIALIZER)
        try:
            srvr = manager.get_server()
            addr = srvr.address
            # Close the connection.Listener socket which gets opened as a part
            # of manager.get_server(). It's not needed for the test.
            srvr.listener.close()
            manager.start()

            p = self.Process(target=self._putter, args=(manager.address, authkey))
            p.start()
            p.join()
            queue = manager.get_queue()
            self.assertEqual(queue.get(), "hello world")
            del queue
        finally:
            if hasattr(manager, "shutdown"):
                manager.shutdown()

        manager = QueueManager(address=addr, authkey=authkey, serializer=SERIALIZER)
        try:
            manager.start()
            self.addCleanup(manager.shutdown)
        except OSError as e:
            if e.errno != errno.EADDRINUSE:
                raise
            # Retry after some time, in case the old socket was lingering
            # (sporadic failure on buildbots)
            time.sleep(1.0)
            manager = QueueManager(address=addr, authkey=authkey, serializer=SERIALIZER)
            if hasattr(manager, "shutdown"):
                self.addCleanup(manager.shutdown)


@unittest.skip("DRAGON: Threads not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithThreadsTestManagerRestart(BaseTestCase, ThreadsMixin, unittest.TestCase):
    @classmethod
    def _putter(cls, address, authkey):
        manager = QueueManager(address=address, authkey=authkey, serializer=SERIALIZER)
        manager.connect()
        queue = manager.get_queue()
        queue.put("hello world")

    def test_rapid_restart(self):
        authkey = os.urandom(32)
        manager = QueueManager(address=(socket_helper.HOST, 0), authkey=authkey, serializer=SERIALIZER)
        try:
            srvr = manager.get_server()
            addr = srvr.address
            # Close the connection.Listener socket which gets opened as a part
            # of manager.get_server(). It's not needed for the test.
            srvr.listener.close()
            manager.start()

            p = self.Process(target=self._putter, args=(manager.address, authkey))
            p.start()
            p.join()
            queue = manager.get_queue()
            self.assertEqual(queue.get(), "hello world")
            del queue
        finally:
            if hasattr(manager, "shutdown"):
                manager.shutdown()

        manager = QueueManager(address=addr, authkey=authkey, serializer=SERIALIZER)
        try:
            manager.start()
            self.addCleanup(manager.shutdown)
        except OSError as e:
            if e.errno != errno.EADDRINUSE:
                raise
            # Retry after some time, in case the old socket was lingering
            # (sporadic failure on buildbots)
            time.sleep(1.0)
            manager = QueueManager(address=addr, authkey=authkey, serializer=SERIALIZER)
            if hasattr(manager, "shutdown"):
                self.addCleanup(manager.shutdown)


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_manager.py
DragonHPC/dragon
"""Various unittests not associated with a MixIn.
"""
import unittest
import unittest.mock
import os
import sys
import time
import random
import socket
import io
import subprocess
import signal
import warnings
import weakref
import errno
import gc

import queue as pyqueue

import unittest
import unittest.mock
import test.support
import test.support.script_helper
from test import support
from test.support import hashlib_helper
from test.support import socket_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing import util

try:
    from multiprocessing import reduction

    HAS_REDUCTION = reduction.HAVE_SEND_HANDLE
except ImportError:
    HAS_REDUCTION = False

from common import (
    join_process,
    setUpModule,
    tearDownModule,
    WIN32,
    PRELOAD,
)

if os.name == "posix":
    from multiprocessing import resource_tracker

    def _resource_unlink(name, rtype):
        resource_tracker._CLEANUP_FUNCS[rtype](name)


#
# Test to verify handle verification, see issue 3321
#


class TestInvalidHandle(unittest.TestCase):
    @unittest.skipIf(WIN32, "skipped on Windows")
    def test_invalid_handles(self):
        conn = multiprocessing.connection.Connection(44977608)
        # check that poll() doesn't crash
        try:
            conn.poll()
        except (ValueError, OSError):
            pass
        finally:
            # Hack private attribute _handle to avoid printing an error
            # in conn.__del__
            conn._handle = None
        self.assertRaises((ValueError, OSError), multiprocessing.connection.Connection, -1)


@hashlib_helper.requires_hashdigest("md5")
class OtherTest(unittest.TestCase):
    # TODO: add more tests for deliver/answer challenge.
    def test_deliver_challenge_auth_failure(self):
        class _FakeConnection(object):
            def recv_bytes(self, size):
                return b"something bogus"

            def send_bytes(self, data):
                pass

        self.assertRaises(
            multiprocessing.AuthenticationError,
            multiprocessing.connection.deliver_challenge,
            _FakeConnection(),
            b"abc",
        )

    def test_answer_challenge_auth_failure(self):
        class _FakeConnection(object):
            def __init__(self):
                self.count = 0

            def recv_bytes(self, size):
                self.count += 1
                if self.count == 1:
                    return multiprocessing.connection.CHALLENGE
                elif self.count == 2:
                    return b"something bogus"
                return b""

            def send_bytes(self, data):
                pass

        self.assertRaises(
            multiprocessing.AuthenticationError,
            multiprocessing.connection.answer_challenge,
            _FakeConnection(),
            b"abc",
        )


#
# Test Manager.start()/Pool.__init__() initializer feature - see issue 5585
#


def initializer(ns):
    ns.test += 1


@hashlib_helper.requires_hashdigest("md5")
@unittest.skip("CIRRUS-1473 DragonPool nor DragonPatchedPool pass this.")
class TestInitializers(unittest.TestCase):
    def setUp(self):
        self.mgr = multiprocessing.Manager()
        self.ns = self.mgr.Namespace()
        self.ns.test = 0

    def tearDown(self):
        self.mgr.shutdown()
        self.mgr.join()

    def test_manager_initializer(self):
        m = multiprocessing.managers.SyncManager()
        self.assertRaises(TypeError, m.start, 1)
        m.start(initializer, (self.ns,))
        self.assertEqual(self.ns.test, 1)
        m.shutdown()
        m.join()

    def test_pool_initializer(self):
        self.assertRaises(TypeError, multiprocessing.Pool, initializer=1)
        p = multiprocessing.Pool(1, initializer, (self.ns,))
        p.close()
        p.join()
        self.assertEqual(self.ns.test, 1)


#
# Issue 5155, 5313, 5331: Test process in processes
# Verifies os.close(sys.stdin.fileno) vs. sys.stdin.close() behavior
#


def _this_sub_process(q):
    try:
        item = q.get(block=False)
    except pyqueue.Empty:
        pass


def _test_process():
    queue = multiprocessing.Queue()
    subProc = multiprocessing.Process(target=_this_sub_process, args=(queue,))
    subProc.daemon = True
    subProc.start()
    subProc.join()


def _afunc(x):
    return x * x


def pool_in_process():
    pool = multiprocessing.Pool(processes=4)
    x = pool.map(_afunc, [1, 2, 3, 4, 5, 6, 7])
    pool.close()
    pool.join()


class _file_like(object):
    def __init__(self, delegate):
        self._delegate = delegate
        self._pid = None

    @property
    def cache(self):
        pid = os.getpid()
        # There are no race conditions since fork keeps only the running thread
        if pid != self._pid:
            self._pid = pid
            self._cache = []
        return self._cache

    def write(self, data):
        self.cache.append(data)

    def flush(self):
        self._delegate.write("".join(self.cache))
        self._cache = []


class TestStdinBadfiledescriptor(unittest.TestCase):
    def test_queue_in_process(self):
        proc = multiprocessing.Process(target=_test_process)
        proc.start()
        proc.join()

    @unittest.skip("bug filed CIRRUS-1473.")
    def test_pool_in_process(self):
        p = multiprocessing.Process(target=pool_in_process)
        p.start()
        p.join()

    def test_flushing(self):
        sio = io.StringIO()
        flike = _file_like(sio)
        flike.write("foo")
        proc = multiprocessing.Process(target=lambda: flike.flush())
        flike.flush()
        assert sio.getvalue() == "foo"


class TestWait(unittest.TestCase):
    @classmethod
    def _child_test_wait(cls, w, slow):
        # DRAGON os.getpid() is not applicable in dragon, as it doesn't generalize to multinode.
        p = multiprocessing.current_process()  # DRAGON
        for i in range(10):
            if slow:
                time.sleep(random.random() * 0.1)
            # w.send((i, os.getpid())) DRAGON
            w.send((i, p.pid))  # DRAGON
        w.close()

    def test_wait(self, slow=False):
        from multiprocessing.connection import wait

        readers = []
        procs = []
        messages = []

        for i in range(4):
            r, w = multiprocessing.Pipe(duplex=False)
            p = multiprocessing.Process(target=self._child_test_wait, args=(w, slow))
            p.daemon = True
            p.start()
            w.close()
            readers.append(r)
            procs.append(p)
            self.addCleanup(p.join)

        while readers:
            for r in wait(readers):
                try:
                    msg = r.recv()
                except EOFError:
                    readers.remove(r)
                    r.close()
                else:
                    messages.append(msg)

        messages.sort()
        expected = sorted((i, p.pid) for i in range(10) for p in procs)
        self.assertEqual(messages, expected)

    @classmethod
    def _child_test_wait_socket(cls, address, slow):
        s = socket.socket()
        s.connect(address)
        for i in range(10):
            if slow:
                time.sleep(random.random() * 0.1)
            s.sendall(("%s\n" % i).encode("ascii"))
        s.close()

    def test_wait_socket(self, slow=False):
        from multiprocessing.connection import wait

        l = socket.create_server((socket_helper.HOST, 0))
        addr = l.getsockname()
        readers = []
        procs = []
        dic = {}

        for i in range(4):
            p = multiprocessing.Process(target=self._child_test_wait_socket, args=(addr, slow))
            p.daemon = True
            p.start()
            procs.append(p)
            self.addCleanup(p.join)

        for i in range(4):
            r, _ = l.accept()
            readers.append(r)
            dic[r] = []
        l.close()

        while readers:
            for r in wait(readers):
                msg = r.recv(32)
                if not msg:
                    readers.remove(r)
                    r.close()
                else:
                    dic[r].append(msg)

        expected = "".join("%s\n" % i for i in range(10)).encode("ascii")
        for v in dic.values():
            self.assertEqual(b"".join(v), expected)

    def test_wait_slow(self):
        self.test_wait(True)

    def test_wait_socket_slow(self):
        self.test_wait_socket(True)

    def test_wait_timeout(self):
        from multiprocessing.connection import wait

        expected = 5
        a, b = multiprocessing.Pipe()

        start = time.monotonic()
        res = wait([a, b], expected)
        delta = time.monotonic() - start

        self.assertEqual(res, [])
        self.assertLess(delta, expected * 2)
        self.assertGreater(delta, expected * 0.5)

        b.send(None)

        start = time.monotonic()
        res = wait([a, b], 20)
        delta = time.monotonic() - start

        self.assertEqual(res, [a])
        self.assertLess(delta, 0.4)

    @classmethod
    def signal_and_sleep(cls, sem, period):
        sem.release()
        time.sleep(period)

    @unittest.skip("bug filed PE-40903")
    def test_wait_integer(self):
        from multiprocessing.connection import wait

        expected = 3
        sorted_ = lambda l: sorted(l, key=lambda x: id(x))
        sem = multiprocessing.Semaphore(0)
        a, b = multiprocessing.Pipe()
        p = multiprocessing.Process(target=self.signal_and_sleep, args=(sem, expected))

        p.start()
        self.assertIsInstance(p.sentinel, int)
        self.assertTrue(sem.acquire(timeout=20))

        start = time.monotonic()
        res = wait([a, p.sentinel, b], expected + 20)
        delta = time.monotonic() - start

        self.assertEqual(res, [p.sentinel])
        self.assertLess(delta, expected + 2)
        self.assertGreater(delta, expected - 2)

        a.send(None)

        start = time.monotonic()
        res = wait([a, p.sentinel, b], 20)
        delta = time.monotonic() - start

        self.assertEqual(sorted_(res), sorted_([p.sentinel, b]))
        self.assertLess(delta, 0.4)

        b.send(None)

        start = time.monotonic()
        res = wait([a, p.sentinel, b], 20)
        delta = time.monotonic() - start

        self.assertEqual(sorted_(res), sorted_([a, p.sentinel, b]))
        self.assertLess(delta, 0.4)

        p.terminate()
        p.join()

    def test_neg_timeout(self):
        from multiprocessing.connection import wait

        a, b = multiprocessing.Pipe()
        t = time.monotonic()
        res = wait([a], timeout=-1)
        t = time.monotonic() - t
        self.assertEqual(res, [])
        self.assertLess(t, 1)
        a.close()
        b.close()


#
# Issue 14151: Test invalid family on invalid environment
#

@unittest.skip(f"DRAGON: Not Implemented")
class TestInvalidFamily(unittest.TestCase):
    @unittest.skipIf(WIN32, "skipped on Windows")
    def test_invalid_family(self):
        with self.assertRaises(ValueError):
            multiprocessing.connection.Listener(r"\\.\test")

    @unittest.skipUnless(WIN32, "skipped on non-Windows platforms")
    def test_invalid_family_win32(self):
        with self.assertRaises(ValueError):
            multiprocessing.connection.Listener("/var/test.pipe")


#
# Issue 12098: check sys.flags of child matches that for parent
#


class TestFlags(unittest.TestCase):
    @classmethod
    def run_in_grandchild(cls, conn):
        conn.send(tuple(sys.flags))

    @classmethod
    def run_in_child(cls):
        import json

        r, w = multiprocessing.Pipe(duplex=False)
        p = multiprocessing.Process(target=cls.run_in_grandchild, args=(w,))
        p.start()
        grandchild_flags = r.recv()
        p.join()
        r.close()
        w.close()
        flags = (tuple(sys.flags), grandchild_flags)
        print(json.dumps(flags))

    def test_flags(self):
        import json

        # start child process using unusual flags
        prog = "from test._test_multiprocessing import TestFlags; " + "TestFlags.run_in_child()"
        data = subprocess.check_output([sys.executable, "-E", "-S", "-O", "-c", prog])
        child_flags, grandchild_flags = json.loads(data.decode("ascii"))
        self.assertEqual(child_flags, grandchild_flags)


#
# Test interaction with socket timeouts - see Issue #6056
#

@unittest.skip(f"Dragon: Not Implemented")
class TestTimeouts(unittest.TestCase):
    @classmethod
    def _test_timeout(cls, child, address):
        time.sleep(1)
        child.send(123)
        child.close()
        conn = multiprocessing.connection.Client(address)
        conn.send(456)
        conn.close()

    def test_timeout(self):
        old_timeout = socket.getdefaulttimeout()
        try:
            socket.setdefaulttimeout(0.1)
            parent, child = multiprocessing.Pipe(duplex=True)
            l = multiprocessing.connection.Listener(family="AF_INET")
            p = multiprocessing.Process(target=self._test_timeout, args=(child, l.address))
            p.start()
            child.close()
            self.assertEqual(parent.recv(), 123)
            parent.close()
            conn = l.accept()
            self.assertEqual(conn.recv(), 456)
            conn.close()
            l.close()
            join_process(p)
        finally:
            socket.setdefaulttimeout(old_timeout)


#
# Test what happens with no "if __name__ == '__main__'"
#


class TestNoForkBomb(unittest.TestCase):
    @unittest.skip("bug filed PE-40903")
    def test_noforkbomb(self):
        sm = multiprocessing.get_start_method()
        name = os.path.join(os.path.dirname(__file__), "mp_fork_bomb.py")
        if sm != "fork":
            rc, out, err = test.support.script_helper.assert_python_failure(name, sm)
            self.assertEqual(out, b"")
            self.assertIn(b"RuntimeError", err)
        else:
            rc, out, err = test.support.script_helper.assert_python_ok(name, sm)
            self.assertEqual(out.rstrip(), b"123")
            self.assertEqual(err, b"")


#
# Issue #17555: ForkAwareThreadLock
#


class TestForkAwareThreadLock(unittest.TestCase):
    # We recursively start processes.  Issue #17555 meant that the
    # after fork registry would get duplicate entries for the same
    # lock.  The size of the registry at generation n was ~2**n.

    @classmethod
    def child(cls, n, conn):
        if n > 1:
            p = multiprocessing.Process(target=cls.child, args=(n - 1, conn))
            p.start()
            conn.close()
            join_process(p)
        else:
            conn.send(len(util._afterfork_registry))
        conn.close()

    def test_lock(self):
        r, w = multiprocessing.Pipe(False)
        l = util.ForkAwareThreadLock()
        old_size = len(util._afterfork_registry)
        p = multiprocessing.Process(target=self.child, args=(5, w))
        p.start()
        w.close()
        new_size = r.recv()
        join_process(p)
        self.assertLessEqual(new_size, old_size)


#
# Check that non-forked child processes do not inherit unneeded fds/handles
#


class TestCloseFds(unittest.TestCase):
    def get_high_socket_fd(self):
        if WIN32:
            # The child process will not have any socket handles, so
            # calling socket.fromfd() should produce WSAENOTSOCK even
            # if there is a handle of the same number.
            return socket.socket().detach()
        else:
            # We want to produce a socket with an fd high enough that a
            # freshly created child process will not have any fds as high.
            fd = socket.socket().detach()
            to_close = []
            while fd < 50:
                to_close.append(fd)
                fd = os.dup(fd)
            for x in to_close:
                os.close(x)
            return fd

    def close(self, fd):
        if WIN32:
            socket.socket(socket.AF_INET, socket.SOCK_STREAM, fileno=fd).close()
        else:
            os.close(fd)

    @classmethod
    def _test_closefds(cls, conn, fd):
        try:
            s = socket.fromfd(fd, socket.AF_INET, socket.SOCK_STREAM)
        except Exception as e:
            conn.send(e)
        else:
            s.close()
            conn.send(None)

    def test_closefd(self):
        if not HAS_REDUCTION:
            raise unittest.SkipTest("requires fd pickling")

        reader, writer = multiprocessing.Pipe()
        fd = self.get_high_socket_fd()
        try:
            p = multiprocessing.Process(target=self._test_closefds, args=(writer, fd))
            p.start()
            writer.close()
            e = reader.recv()
            join_process(p)
        finally:
            self.close(fd)
            writer.close()
            reader.close()

        if multiprocessing.get_start_method() == "fork":
            self.assertIs(e, None)
        else:
            WSAENOTSOCK = 10038
            self.assertIsInstance(e, OSError)
            self.assertTrue(e.errno == errno.EBADF or e.winerror == WSAENOTSOCK, e)


#
# Issue #17097: EINTR should be ignored by recv(), send(), accept() etc
#


class TestIgnoreEINTR(unittest.TestCase):

    # Sending CONN_MAX_SIZE bytes into a multiprocessing pipe must block
    CONN_MAX_SIZE = max(support.PIPE_MAX_SIZE, support.SOCK_MAX_SIZE)

    @classmethod
    def _test_ignore(cls, conn):
        def handler(signum, frame):
            pass

        signal.signal(signal.SIGUSR1, handler)
        conn.send("ready")
        x = conn.recv()
        conn.send(x)
        conn.send_bytes(b"x" * cls.CONN_MAX_SIZE)

    @unittest.skip("DRAGON: not implemented")
    @unittest.skipUnless(hasattr(signal, "SIGUSR1"), "requires SIGUSR1")
    def test_ignore(self):
        conn, child_conn = multiprocessing.Pipe()
        try:
            p = multiprocessing.Process(target=self._test_ignore, args=(child_conn,))
            p.daemon = True
            p.start()
            child_conn.close()
            self.assertEqual(conn.recv(), "ready")
            time.sleep(0.1)
            os.kill(p.pid, signal.SIGUSR1)
            time.sleep(0.1)
            conn.send(1234)
            self.assertEqual(conn.recv(), 1234)
            time.sleep(0.1)
            os.kill(p.pid, signal.SIGUSR1)
            self.assertEqual(conn.recv_bytes(), b"x" * self.CONN_MAX_SIZE)
            time.sleep(0.1)
            p.join()
        finally:
            conn.close()

    @classmethod
    def _test_ignore_listener(cls, conn):
        def handler(signum, frame):
            pass

        signal.signal(signal.SIGUSR1, handler)
        with multiprocessing.connection.Listener() as l:
            conn.send(l.address)
            a = l.accept()
            a.send("welcome")

    @unittest.skip("bug filed PE-40903")
    @unittest.skipUnless(hasattr(signal, "SIGUSR1"), "requires SIGUSR1")
    def test_ignore_listener(self):
        conn, child_conn = multiprocessing.Pipe()
        try:
            p = multiprocessing.Process(target=self._test_ignore_listener, args=(child_conn,))
            p.daemon = True
            p.start()
            child_conn.close()
            address = conn.recv()
            time.sleep(0.1)
            os.kill(p.pid, signal.SIGUSR1)
            time.sleep(0.1)
            client = multiprocessing.connection.Client(address)
            self.assertEqual(client.recv(), "welcome")
            p.join()
        finally:
            conn.close()


class TestStartMethod(unittest.TestCase):
    @classmethod
    def _check_context(cls, conn):
        conn.send(multiprocessing.get_start_method())

    def check_context(self, ctx):
        r, w = ctx.Pipe(duplex=False)
        p = ctx.Process(target=self._check_context, args=(w,))
        p.start()
        w.close()
        child_method = r.recv()
        r.close()
        p.join()
        self.assertEqual(child_method, ctx.get_start_method())

    @unittest.skip("bug filed PE-40903")
    def test_context(self):
        for method in ("fork", "spawn", "forkserver"):
            try:
                ctx = multiprocessing.get_context(method)
            except ValueError:
                continue
            self.assertEqual(ctx.get_start_method(), method)
            self.assertIs(ctx.get_context(), ctx)
            self.assertRaises(ValueError, ctx.set_start_method, "spawn")
            self.assertRaises(ValueError, ctx.set_start_method, None)
            self.check_context(ctx)

    @unittest.skip("bug filed PE-40903")
    def test_set_get(self):
        multiprocessing.set_forkserver_preload(PRELOAD)
        count = 0
        old_method = multiprocessing.get_start_method()
        try:
            for method in ("fork", "spawn", "forkserver"):
                try:
                    multiprocessing.set_start_method(method, force=True)
                except ValueError:
                    continue
                self.assertEqual(multiprocessing.get_start_method(), method)
                ctx = multiprocessing.get_context()
                self.assertEqual(ctx.get_start_method(), method)
                self.assertTrue(type(ctx).__name__.lower().startswith(method))
                self.assertTrue(ctx.Process.__name__.lower().startswith(method))
                self.check_context(multiprocessing)
                count += 1
        finally:
            multiprocessing.set_start_method(old_method, force=True)
        self.assertGreaterEqual(count, 1)

    def test_get_all(self):
        methods = multiprocessing.get_all_start_methods()
        if sys.platform == "win32":
            self.assertEqual(methods, ["spawn"])
        else:
            # DRAGON Need to modify this, because we have another start method
            self.assertTrue("dragon" in methods)  # DRAGON
            methods.remove("dragon")  # DRAGON
            self.assertTrue(
                methods == ["fork", "spawn"]
                or methods == ["spawn", "fork"]
                or methods == ["fork", "spawn", "forkserver"]
                or methods == ["spawn", "fork", "forkserver"]
            )

    def test_preload_resources(self):
        if multiprocessing.get_start_method() != "forkserver":
            self.skipTest("test only relevant for 'forkserver' method")
        name = os.path.join(os.path.dirname(__file__), "mp_preload.py")
        rc, out, err = test.support.script_helper.assert_python_ok(name)
        out = out.decode()
        err = err.decode()
        if out.rstrip() != "ok" or err != "":
            print(out)
            print(err)
            self.fail("failed spawning forkserver or grandchild")

@unittest.skip(f"DRAGON: Not Implemented PE-45417")
@unittest.skipIf(sys.platform == "win32", "test semantics don't make sense on Windows")
class TestResourceTracker(unittest.TestCase):
    def test_resource_tracker(self):
        #
        # Check that killing process does not leak named semaphores
        #
        cmd = """if 1:
            import time, os, tempfile
            import multiprocessing as mp
            from multiprocessing import resource_tracker
            from multiprocessing.shared_memory import SharedMemory

            mp.set_start_method("spawn")
            rand = tempfile._RandomNameSequence()


            def create_and_register_resource(rtype):
                if rtype == "semaphore":
                    lock = mp.Lock()
                    return lock, lock._semlock.name
                elif rtype == "shared_memory":
                    sm = SharedMemory(create=True, size=10)
                    return sm, sm._name
                else:
                    raise ValueError(
                        "Resource type {{}} not understood".format(rtype))


            resource1, rname1 = create_and_register_resource("{rtype}")
            resource2, rname2 = create_and_register_resource("{rtype}")

            os.write({w}, rname1.encode("ascii") + b"\\n")
            os.write({w}, rname2.encode("ascii") + b"\\n")

            time.sleep(10)
        """
        for rtype in resource_tracker._CLEANUP_FUNCS:
            with self.subTest(rtype=rtype):
                if rtype == "noop":
                    # Artefact resource type used by the resource_tracker
                    continue
                r, w = os.pipe()
                p = subprocess.Popen(
                    [sys.executable, "-E", "-c", cmd.format(w=w, rtype=rtype)],
                    pass_fds=[w],
                    stderr=subprocess.PIPE,
                )
                os.close(w)
                with open(r, "rb", closefd=True) as f:
                    name1 = f.readline().rstrip().decode("ascii")
                    name2 = f.readline().rstrip().decode("ascii")
                _resource_unlink(name1, rtype)
                p.terminate()
                p.wait()

                deadline = time.monotonic() + support.LONG_TIMEOUT
                while time.monotonic() < deadline:
                    time.sleep(0.5)
                    try:
                        _resource_unlink(name2, rtype)
                    except OSError as e:
                        # docs say it should be ENOENT, but OSX seems to give
                        # EINVAL
                        self.assertIn(e.errno, (errno.ENOENT, errno.EINVAL))
                        break
                else:
                    raise AssertionError(
                        f"A {rtype} resource was leaked after a process was " f"abruptly terminated."
                    )
                err = p.stderr.read().decode("utf-8")
                p.stderr.close()
                expected = "resource_tracker: There appear to be 2 leaked {} " "objects".format(rtype)
                self.assertRegex(err, expected)
                self.assertRegex(err, r"resource_tracker: %r: \[Errno" % name1)

    def check_resource_tracker_death(self, signum, should_die):
        # bpo-31310: if the semaphore tracker process has died, it should
        # be restarted implicitly.
        from multiprocessing.resource_tracker import _resource_tracker

        pid = _resource_tracker._pid
        if pid is not None:
            os.kill(pid, signal.SIGKILL)
            support.wait_process(pid, exitcode=-signal.SIGKILL)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            _resource_tracker.ensure_running()
        pid = _resource_tracker._pid

        os.kill(pid, signum)
        time.sleep(1.0)  # give it time to die

        ctx = multiprocessing.get_context("spawn")
        with warnings.catch_warnings(record=True) as all_warn:
            warnings.simplefilter("always")
            sem = ctx.Semaphore()
            sem.acquire()
            sem.release()
            wr = weakref.ref(sem)
            # ensure `sem` gets collected, which triggers communication with
            # the semaphore tracker
            del sem
            gc.collect()
            self.assertIsNone(wr())
            if should_die:
                self.assertEqual(len(all_warn), 1)
                the_warn = all_warn[0]
                self.assertTrue(issubclass(the_warn.category, UserWarning))
                self.assertTrue("resource_tracker: process died" in str(the_warn.message))
            else:
                self.assertEqual(len(all_warn), 0)

    def test_resource_tracker_sigint(self):
        # Catchable signal (ignored by semaphore tracker)
        self.check_resource_tracker_death(signal.SIGINT, False)

    def test_resource_tracker_sigterm(self):
        # Catchable signal (ignored by semaphore tracker)
        self.check_resource_tracker_death(signal.SIGTERM, False)

    def test_resource_tracker_sigkill(self):
        # Uncatchable signal.
        self.check_resource_tracker_death(signal.SIGKILL, True)

    @staticmethod
    def _is_resource_tracker_reused(conn, pid):
        from multiprocessing.resource_tracker import _resource_tracker

        _resource_tracker.ensure_running()
        # The pid should be None in the child process, expect for the fork
        # context. It should not be a new value.
        reused = _resource_tracker._pid in (None, pid)
        reused &= _resource_tracker._check_alive()
        conn.send(reused)

    @unittest.skip("bug filed PE-40903")
    def test_resource_tracker_reused(self):
        from multiprocessing.resource_tracker import _resource_tracker

        _resource_tracker.ensure_running()
        pid = _resource_tracker._pid

        r, w = multiprocessing.Pipe(duplex=False)
        p = multiprocessing.Process(target=self._is_resource_tracker_reused, args=(w, pid))
        p.start()
        is_resource_tracker_reused = r.recv()

        # Clean up
        p.join()
        w.close()
        r.close()

        self.assertTrue(is_resource_tracker_reused)


class TestSimpleQueue(unittest.TestCase):
    @classmethod
    def _test_empty(cls, queue, child_can_start, parent_can_continue):
        child_can_start.wait()
        # issue 30301, could fail under spawn and forkserver
        try:
            queue.put(queue.empty())
            queue.put(queue.empty())
        finally:
            parent_can_continue.set()

    def test_empty(self):
        queue = multiprocessing.SimpleQueue()
        child_can_start = multiprocessing.Event()
        parent_can_continue = multiprocessing.Event()

        proc = multiprocessing.Process(
            target=self._test_empty, args=(queue, child_can_start, parent_can_continue)
        )
        proc.daemon = True
        proc.start()

        self.assertTrue(queue.empty())

        child_can_start.set()
        parent_can_continue.wait()

        self.assertFalse(queue.empty())
        self.assertEqual(queue.get(), True)
        self.assertEqual(queue.get(), False)
        self.assertTrue(queue.empty())

        proc.join()

    def test_close(self):
        queue = multiprocessing.SimpleQueue()
        queue.close()
        # closing a queue twice should not fail
        queue.close()

    # Test specific to CPython since it tests private attributes
    @test.support.cpython_only
    def test_closed(self):
        queue = multiprocessing.SimpleQueue()
        queue.close()
        self.assertTrue(queue._reader.closed)
        self.assertTrue(queue._writer.closed)


class TestPoolNotLeakOnFailure(unittest.TestCase):
    @unittest.skip("CIRRUS-1473 DRAGON: Fails with native pool.")
    def test_release_unused_processes(self):
        # Issue #19675: During pool creation, if we can't create a process,
        # don't leak already created ones.
        will_fail_in = 3
        forked_processes = []

        class FailingForkProcess:
            def __init__(self, **kwargs):
                self.name = "Fake Process"
                self.exitcode = None
                self.state = None
                forked_processes.append(self)

            def start(self):
                nonlocal will_fail_in
                if will_fail_in <= 0:
                    raise OSError("Manually induced OSError")
                will_fail_in -= 1
                self.state = "started"

            def terminate(self):
                self.state = "stopping"

            def join(self):
                if self.state == "stopping":
                    self.state = "stopped"

            def is_alive(self):
                return self.state == "started" or self.state == "stopping"

        with self.assertRaisesRegex(OSError, "Manually induced OSError"):
            p = multiprocessing.pool.Pool(5, context=unittest.mock.MagicMock(Process=FailingForkProcess))
            p.close()
            p.join()
        self.assertFalse(any(process.is_alive() for process in forked_processes))

    @unittest.skip("DRAGON: doesn't have managers yet.")
    @hashlib_helper.requires_hashdigest("md5")
    class TestSyncManagerTypes(unittest.TestCase):
        """Test all the types which can be shared between a parent and a
        child process by using a manager which acts as an intermediary
        between them.

        In the following unit-tests the base type is created in the parent
        process, the @classmethod represents the worker process and the
        shared object is readable and editable between the two.

        # The child.
        @classmethod
        def _test_list(cls, obj):
            assert obj[0] == 5
            assert obj.append(6)

        # The parent.
        def test_list(self):
            o = self.manager.list()
            o.append(5)
            self.run_worker(self._test_list, o)
            assert o[1] == 6
        """

        # DRAGON: commented out so file works. Comment back in, when we have managers
        # manager_class = multiprocessing.managers.SyncManager

        def setUp(self):
            self.manager = self.manager_class()
            self.manager.start()
            self.proc = None

        def tearDown(self):
            if self.proc is not None and self.proc.is_alive():
                self.proc.terminate()
                self.proc.join()
            self.manager.shutdown()
            self.manager = None
            self.proc = None

        @classmethod
        def setUpClass(cls):
            support.reap_children()

        tearDownClass = setUpClass

        def wait_proc_exit(self):
            # Only the manager process should be returned by active_children()
            # but this can take a bit on slow machines, so wait a few seconds
            # if there are other children too (see #17395).
            join_process(self.proc)
            start_time = time.monotonic()
            t = 0.01
            while len(multiprocessing.active_children()) > 1:
                time.sleep(t)
                t *= 2
                dt = time.monotonic() - start_time
                if dt >= 5.0:
                    test.support.environment_altered = True
                    support.print_warning(
                        f"multiprocessing.Manager still has "
                        f"{multiprocessing.active_children()} "
                        f"active children after {dt} seconds"
                    )
                    break

        def run_worker(self, worker, obj):
            self.proc = multiprocessing.Process(target=worker, args=(obj,))
            self.proc.daemon = True
            self.proc.start()
            self.wait_proc_exit()
            self.assertEqual(self.proc.exitcode, 0)

        @classmethod
        def _test_event(cls, obj):
            assert obj.is_set()
            obj.wait()
            obj.clear()
            obj.wait(0.001)

        def test_event(self):
            o = self.manager.Event()
            o.set()
            self.run_worker(self._test_event, o)
            assert not o.is_set()
            o.wait(0.001)

        @classmethod
        def _test_lock(cls, obj):
            obj.acquire()

        def test_lock(self, lname="Lock"):
            o = getattr(self.manager, lname)()
            self.run_worker(self._test_lock, o)
            o.release()
            self.assertRaises(RuntimeError, o.release)  # already released

        @classmethod
        def _test_rlock(cls, obj):
            obj.acquire()
            obj.release()

        def test_rlock(self, lname="Lock"):
            o = getattr(self.manager, lname)()
            self.run_worker(self._test_rlock, o)

        @classmethod
        def _test_semaphore(cls, obj):
            obj.acquire()

        def test_semaphore(self, sname="Semaphore"):
            o = getattr(self.manager, sname)()
            self.run_worker(self._test_semaphore, o)
            o.release()

        def test_bounded_semaphore(self):
            self.test_semaphore(sname="BoundedSemaphore")

        @classmethod
        def _test_condition(cls, obj):
            obj.acquire()
            obj.release()

        def test_condition(self):
            o = self.manager.Condition()
            self.run_worker(self._test_condition, o)

        @classmethod
        def _test_barrier(cls, obj):
            assert obj.parties == 5
            obj.reset()

        def test_barrier(self):
            o = self.manager.Barrier(5)
            self.run_worker(self._test_barrier, o)

        @classmethod
        def _test_pool(cls, obj):
            # TODO: fix https://bugs.python.org/issue35919
            with obj:
                pass

        def test_pool(self):
            o = self.manager.Pool(processes=4)
            self.run_worker(self._test_pool, o)

        @classmethod
        def _test_queue(cls, obj):
            assert obj.qsize() == 2
            assert obj.full()
            assert not obj.empty()
            assert obj.get() == 5
            assert not obj.empty()
            assert obj.get() == 6
            assert obj.empty()

        def test_queue(self, qname="Queue"):
            o = getattr(self.manager, qname)(2)
            o.put(5)
            o.put(6)
            self.run_worker(self._test_queue, o)
            assert o.empty()
            assert not o.full()

        def test_joinable_queue(self):
            self.test_queue("JoinableQueue")

        @classmethod
        def _test_list(cls, obj):
            assert obj[0] == 5
            assert obj.count(5) == 1
            assert obj.index(5) == 0
            obj.sort()
            obj.reverse()
            for x in obj:
                pass
            assert len(obj) == 1
            assert obj.pop(0) == 5

        def test_list(self):
            o = self.manager.list()
            o.append(5)
            self.run_worker(self._test_list, o)
            assert not o
            self.assertEqual(len(o), 0)

        @classmethod
        def _test_dict(cls, obj):
            assert len(obj) == 1
            assert obj["foo"] == 5
            assert obj.get("foo") == 5
            assert list(obj.items()) == [("foo", 5)]
            assert list(obj.keys()) == ["foo"]
            assert list(obj.values()) == [5]
            assert obj.copy() == {"foo": 5}
            assert obj.popitem() == ("foo", 5)

        def test_dict(self):
            o = self.manager.dict()
            o["foo"] = 5
            self.run_worker(self._test_dict, o)
            assert not o
            self.assertEqual(len(o), 0)

        @classmethod
        def _test_value(cls, obj):
            assert obj.value == 1
            assert obj.get() == 1
            obj.set(2)

        def test_value(self):
            o = self.manager.Value("i", 1)
            self.run_worker(self._test_value, o)
            self.assertEqual(o.value, 2)
            self.assertEqual(o.get(), 2)

        @classmethod
        def _test_array(cls, obj):
            assert obj[0] == 0
            assert obj[1] == 1
            assert len(obj) == 2
            assert list(obj) == [0, 1]

        def test_array(self):
            o = self.manager.Array("i", [0, 1])
            self.run_worker(self._test_array, o)

        @classmethod
        def _test_namespace(cls, obj):
            assert obj.x == 0
            assert obj.y == 1

        def test_namespace(self):
            o = self.manager.Namespace()
            o.x = 0
            o.y = 1
            self.run_worker(self._test_namespace, o)


class MiscTestCase(unittest.TestCase):
    def test__all__(self):
        # Just make sure names in blacklist are excluded
        support.check__all__(
            self, multiprocessing, extra=multiprocessing.__all__, blacklist=["SUBDEBUG", "SUBWARNING"]
        )


#

# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_others.py
DragonHPC/dragon
"""Test multiprocessing.Process()
"""
import os
import time
import signal

import unittest


import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing import util
from multiprocessing.connection import wait

from common import (
    BaseTestCase,
    ProcessesMixin,
    ThreadsMixin,
    setUpModule,
    tearDownModule,
)


class WithProcessesTestPoll(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    def test_empty_string(self):
        a, b = self.Pipe()
        self.assertEqual(a.poll(), False)
        b.send_bytes(b"")
        self.assertEqual(a.poll(), True)
        self.assertEqual(a.poll(), True)

    @classmethod
    def _child_strings(cls, conn, strings):
        for s in strings:
            time.sleep(0.1)
            conn.send_bytes(s)
        conn.close()

    def test_strings(self):
        strings = (b"hello", b"", b"a", b"b", b"", b"bye", b"", b"lop")
        a, b = self.Pipe()
        p = self.Process(target=self._child_strings, args=(b, strings))
        p.start()

        for s in strings:
            for i in range(200):
                if a.poll(0.01):
                    break
            x = a.recv_bytes()
            self.assertEqual(s, x)

        p.join()

    @classmethod
    def _child_boundaries(cls, r):
        # Polling may "pull" a message in to the child process, but we
        # don't want it to pull only part of a message, as that would
        # corrupt the pipe for any other processes which might later
        # read from it.
        r.poll(5)

    def test_boundaries(self):
        r, w = self.Pipe(False)
        p = self.Process(target=self._child_boundaries, args=(r,))
        p.start()
        time.sleep(2)
        L = [b"first", b"second"]
        for obj in L:
            w.send_bytes(obj)
        w.close()
        p.join()
        self.assertIn(r.recv_bytes(), L)

    @classmethod
    def _child_dont_merge(cls, b):
        b.send_bytes(b"a")
        b.send_bytes(b"b")
        b.send_bytes(b"cd")

    def test_dont_merge(self):
        a, b = self.Pipe()
        self.assertEqual(a.poll(0.0), False)
        self.assertEqual(a.poll(0.1), False)

        p = self.Process(target=self._child_dont_merge, args=(b,))
        p.start()

        self.assertEqual(a.recv_bytes(), b"a")
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.recv_bytes(), b"b")
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.poll(0.0), True)
        self.assertEqual(a.recv_bytes(), b"cd")

        p.join()


class WithProcessesTestPollEintr(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    @classmethod
    def _killer(cls, pid):
        time.sleep(0.1)
        os.kill(pid, signal.SIGUSR1)

    @unittest.skipUnless(hasattr(signal, "SIGUSR1"), "requires SIGUSR1")
    def test_poll_eintr(self):
        got_signal = [False]

        def record(*args):
            got_signal[0] = True

        pid = os.getpid()
        oldhandler = signal.signal(signal.SIGUSR1, record)
        try:
            killer = self.Process(target=self._killer, args=(pid,))
            killer.start()
            try:
                p = self.Process(target=time.sleep, args=(2,))
                p.start()
                p.join()
            finally:
                killer.join()
            self.assertTrue(got_signal[0])
            self.assertEqual(p.exitcode, 0)
        finally:
            signal.signal(signal.SIGUSR1, oldhandler)


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestPoll(BaseTestCase, ThreadsMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    def test_empty_string(self):
        a, b = self.Pipe()
        self.assertEqual(a.poll(), False)
        b.send_bytes(b"")
        self.assertEqual(a.poll(), True)
        self.assertEqual(a.poll(), True)

    @classmethod
    def _child_strings(cls, conn, strings):
        for s in strings:
            time.sleep(0.1)
            conn.send_bytes(s)
        conn.close()

    def test_strings(self):
        strings = (b"hello", b"", b"a", b"b", b"", b"bye", b"", b"lop")
        a, b = self.Pipe()
        p = self.Process(target=self._child_strings, args=(b, strings))
        p.start()

        for s in strings:
            for i in range(200):
                if a.poll(0.01):
                    break
            x = a.recv_bytes()
            self.assertEqual(s, x)

        p.join()

    @classmethod
    def _child_boundaries(cls, r):
        # Polling may "pull" a message in to the child process, but we
        # don't want it to pull only part of a message, as that would
        # corrupt the pipe for any other processes which might later
        # read from it.
        r.poll(5)

    def test_boundaries(self):
        r, w = self.Pipe(False)
        p = self.Process(target=self._child_boundaries, args=(r,))
        p.start()
        time.sleep(2)
        L = [b"first", b"second"]
        for obj in L:
            w.send_bytes(obj)
        w.close()
        p.join()
        self.assertIn(r.recv_bytes(), L)

    @classmethod
    def _child_dont_merge(cls, b):
        b.send_bytes(b"a")
        b.send_bytes(b"b")
        b.send_bytes(b"cd")

    def test_dont_merge(self):
        a, b = self.Pipe()
        self.assertEqual(a.poll(0.0), False)
        self.assertEqual(a.poll(0.1), False)

        p = self.Process(target=self._child_dont_merge, args=(b,))
        p.start()

        self.assertEqual(a.recv_bytes(), b"a")
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.recv_bytes(), b"b")
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.poll(1.0), True)
        self.assertEqual(a.poll(0.0), True)
        self.assertEqual(a.recv_bytes(), b"cd")

        p.join()


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_poll.py
DragonHPC/dragon
"""Test multiprocessing.Pool()
"""
import sys
import time
import weakref
import gc
import itertools

import unittest
import test.support
import test.support.script_helper
from test.support import hashlib_helper
from test import support

import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing import util
from multiprocessing.managers import RemoteError

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    TimingWrapper,
    DELTA,
    TIMEOUT1,
    TIMEOUT2,
    setUpModule,
    tearDownModule,
)

def sqr(x, wait=0.0):
    time.sleep(wait)
    return x * x


def mul(x, y):
    return x * y


def raise_large_valuerror(wait):
    time.sleep(wait)
    raise ValueError("x" * 1024**2)


def identity(x):
    return x


class CountedObject(object):
    n_instances = 0

    def __new__(cls):
        cls.n_instances += 1
        return object.__new__(cls)

    def __del__(self):
        type(self).n_instances -= 1


class SayWhenError(ValueError):
    pass


def exception_throwing_generator(total, when):
    if when == -1:
        raise SayWhenError("Somebody said when")
    for i in range(total):
        if i == when:
            raise SayWhenError("Somebody said when")
        yield i



class WithProcessesTestPool(BaseTestCase, ProcessesMixin, unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.pool = cls.Pool(4)

    @classmethod
    def tearDownClass(cls):
        cls.pool.terminate()
        cls.pool.join()
        cls.pool = None
        super().tearDownClass()

    def test_apply(self):
        papply = self.pool.apply
        self.assertEqual(papply(sqr, (5,)), sqr(5))
        self.assertEqual(papply(sqr, (), {"x": 3}), sqr(x=3))

    def test_map(self):
        pmap = self.pool.map
        self.assertEqual(pmap(sqr, list(range(10))), list(map(sqr, list(range(10)))))
        self.assertEqual(pmap(sqr, list(range(100)), chunksize=20), list(map(sqr, list(range(100)))))

    def test_starmap(self):
        psmap = self.pool.starmap
        tuples = list(zip(range(10), range(9, -1, -1)))
        self.assertEqual(psmap(mul, tuples), list(itertools.starmap(mul, tuples)))
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(psmap(mul, tuples, chunksize=20), list(itertools.starmap(mul, tuples)))

    def test_starmap_async(self):
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(self.pool.starmap_async(mul, tuples).get(), list(itertools.starmap(mul, tuples)))

    def test_map_async(self):
        self.assertEqual(self.pool.map_async(sqr, list(range(10))).get(), list(map(sqr, list(range(10)))))

    def test_map_async_callbacks(self):
        call_args = self.manager.list() if self.TYPE == "manager" else []
        self.pool.map_async(int, ["1"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(1, len(call_args))
        self.assertEqual([1], call_args[0])
        self.pool.map_async(int, ["a"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(2, len(call_args))
        self.assertIsInstance(call_args[1], ValueError)

    @unittest.skip("bug filed PE-40898. Fails with native pool. Need to handle failure of put differently.")
    def test_map_unplicklable(self):
        # Issue #19425 -- failure to pickle should not cause a hang
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class A(object):
            def __reduce__(self):
                raise RuntimeError("cannot pickle")

        with self.assertRaises(RuntimeError):
            self.pool.map(sqr, [A()] * 10)

    def test_map_chunksize(self):
        try:
            self.pool.map_async(sqr, [], chunksize=1).get(timeout=TIMEOUT1)
        except multiprocessing.TimeoutError:
            self.fail("pool.map_async with chunksize stalled on null list")

    def test_map_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(1, -1), 1)
        # again, make sure it's reentrant
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(1, -1), 1)

        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(10, 3), 1)

        class SpecialIterable:
            def __iter__(self):
                return self

            def __next__(self):
                raise SayWhenError

            def __len__(self):
                return 1

        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, SpecialIterable(), 1)
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, SpecialIterable(), 1)

    def test_async(self):
        res = self.pool.apply_async(
            sqr,
            (
                7,
                TIMEOUT1,
            ),
        )
        get = TimingWrapper(res.get)
        self.assertEqual(get(), 49)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT1)

    def test_async_timeout(self):
        res = self.pool.apply_async(sqr, (6, TIMEOUT2 + 1.0))
        get = TimingWrapper(res.get)
        self.assertRaises(multiprocessing.TimeoutError, get, timeout=TIMEOUT2)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT2)

    def test_imap(self):
        it = self.pool.imap(sqr, list(range(10)))
        self.assertEqual(list(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap(sqr, list(range(10)))
        for i in range(10):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

        it = self.pool.imap(sqr, list(range(1000)), chunksize=100)
        for i in range(1000):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

    def test_imap_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        it = self.pool.imap(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)
        # again, make sure it's reentrant
        it = self.pool.imap(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)

        it = self.pool.imap(sqr, exception_throwing_generator(10, 3), 1)
        for i in range(3):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)

        # SayWhenError seen at start of problematic chunk's results
        it = self.pool.imap(sqr, exception_throwing_generator(20, 7), 2)
        for i in range(6):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)
        it = self.pool.imap(sqr, exception_throwing_generator(20, 7), 4)
        for i in range(4):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)

    def test_imap_unordered(self):
        it = self.pool.imap_unordered(sqr, list(range(10)))
        self.assertEqual(sorted(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap_unordered(sqr, list(range(1000)), chunksize=100)
        self.assertEqual(sorted(it), list(map(sqr, list(range(1000)))))

    def test_imap_unordered_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        it = self.pool.imap_unordered(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)
        # again, make sure it's reentrant
        it = self.pool.imap_unordered(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)

        it = self.pool.imap_unordered(sqr, exception_throwing_generator(10, 3), 1)
        expected_values = list(map(sqr, list(range(10))))
        with self.assertRaises(SayWhenError):
            # imap_unordered makes it difficult to anticipate the SayWhenError
            for i in range(10):
                value = next(it)
                self.assertIn(value, expected_values)
                expected_values.remove(value)

        it = self.pool.imap_unordered(sqr, exception_throwing_generator(20, 7), 2)
        expected_values = list(map(sqr, list(range(20))))
        with self.assertRaises(SayWhenError):
            for i in range(20):
                value = next(it)
                self.assertIn(value, expected_values)
                expected_values.remove(value)

    
    def test_make_pool(self):
        expected_error = RemoteError if self.TYPE == "manager" else ValueError

        self.assertRaises(expected_error, self.Pool, -1)
        self.assertRaises(expected_error, self.Pool, 0)

        if self.TYPE != "manager":
            p = self.Pool(3)
            try:
                self.assertEqual(3, len(p._pool))
            finally:
                p.close()
                p.join()

    def test_terminate(self):
        result = self.pool.map_async(time.sleep, [0.1 for i in range(10000)], chunksize=1)
        self.pool.terminate()
        join = TimingWrapper(self.pool.join)
        join()
        # Sanity check the pool didn't wait for all tasks to finish
        self.assertLess(join.elapsed, 2.0)

    def test_empty_iterable(self):
        # See Issue 12157
        p = self.Pool(1)

        self.assertEqual(p.map(sqr, []), [])
        self.assertEqual(list(p.imap(sqr, [])), [])
        self.assertEqual(list(p.imap_unordered(sqr, [])), [])
        self.assertEqual(p.map_async(sqr, []).get(), [])

        p.close()
        p.join()

    def test_context(self):
        if self.TYPE == "processes":
            L = list(range(10))
            expected = [sqr(i) for i in L]
            with self.Pool(2) as p:
                r = p.map_async(sqr, L)
                self.assertEqual(r.get(), expected)
            p.join()
            self.assertRaises(ValueError, p.map_async, sqr, L)

    @classmethod
    def _test_traceback(cls):
        raise RuntimeError(123)  # some comment

    @unittest.skip("bug filed PE-40898")
    def test_traceback(self):
        # We want ensure that the traceback from the child process is
        # contained in the traceback raised in the main process.
        if self.TYPE == "processes":
            with self.Pool(1) as p:
                try:
                    p.apply(self._test_traceback)
                except Exception as e:
                    exc = e
                else:
                    self.fail("expected RuntimeError")
            p.join()
            self.assertIs(type(exc), RuntimeError)
            self.assertEqual(exc.args, (123,))
            cause = exc.__cause__
            self.assertIs(type(cause), multiprocessing.pool.RemoteTraceback)
            self.assertIn("raise RuntimeError(123) # some comment", cause.tb)

            with test.support.captured_stderr() as f1:
                try:
                    raise exc
                except RuntimeError:
                    sys.excepthook(*sys.exc_info())
            self.assertIn("raise RuntimeError(123) # some comment", f1.getvalue())
            # _helper_reraises_exception should not make the error
            # a remote exception
            with self.Pool(1) as p:
                try:
                    p.map(sqr, exception_throwing_generator(1, -1), 1)
                except Exception as e:
                    exc = e
                else:
                    self.fail("expected SayWhenError")
                self.assertIs(type(exc), SayWhenError)
                self.assertIs(exc.__cause__, None)
            p.join()

    @classmethod
    def _test_wrapped_exception(cls):
        raise RuntimeError("foo")

    @unittest.skip("bug filed CIRRUS-1473.")
    def test_wrapped_exception(self):
        # Issue #20980: Should not wrap exception when using thread pool
        with self.Pool(1) as p:
            with self.assertRaises(RuntimeError):
                p.apply(self._test_wrapped_exception)
        p.join()

    def test_map_no_failfast(self):
        # Issue #23992: the fail-fast behaviour when an exception is raised
        # during map() would make Pool.join() deadlock, because a worker
        # process would fill the result queue (after the result handler thread
        # terminated, hence not draining it anymore).

        t_start = time.monotonic()

        with self.assertRaises(ValueError):
            with self.Pool(2) as p:
                try:
                    p.map(raise_large_valuerror, [0, 1])
                finally:
                    time.sleep(0.5)
                    p.close()
                    p.join()

        # check that we indeed waited for all jobs
        self.assertGreater(time.monotonic() - t_start, 0.9)

    @unittest.skip("bug filed CIRRUS-1473.")
    def test_release_task_refs(self):
        # Issue #29861: task arguments and results should not be kept
        # alive after we are done with them.
        objs = [CountedObject() for i in range(10)]
        refs = [weakref.ref(o) for o in objs]
        self.pool.map(identity, objs)

        del objs
        gc.collect()  # For PyPy or other GCs.
        time.sleep(DELTA)  # let threaded cleanup code run
        self.assertEqual(set(wr() for wr in refs), {None})
        # With a process pool, copies of the objects are returned, check
        # they were released too.
        self.assertEqual(CountedObject.n_instances, 0)

    def test_enter(self):
        if self.TYPE == "manager":
            self.skipTest("test not applicable to manager")

        pool = self.Pool(1)
        with pool:
            pass
            # call pool.terminate()
        # pool is no longer running

        with self.assertRaises(ValueError):
            # bpo-35477: pool.__enter__() fails if the pool is not running
            with pool:
                pass
        pool.join()

    @unittest.skip("bug filed CIRRUS-1473")
    def test_resource_warning(self):
        if self.TYPE == "manager":
            self.skipTest("test not applicable to manager")

        pool = self.Pool(1)
        pool.terminate()
        pool.join()

        # force state to RUN to emit ResourceWarning in __del__()
        pool._state = multiprocessing.pool.RUN

        with support.check_warnings(("unclosed running multiprocessing pool", ResourceWarning)):
            pool = None
            support.gc_collect()


def raising():
    raise KeyError("key")


def unpickleable_result():
    return lambda: 42


class WithProcessesTestPoolWorkerErrors(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', )

    def test_async_error_callback(self):
        p = multiprocessing.Pool(2)

        scratchpad = [None]

        def errback(exc):
            scratchpad[0] = exc

        res = p.apply_async(raising, error_callback=errback)
        self.assertRaises(KeyError, res.get)
        self.assertTrue(scratchpad[0])
        self.assertIsInstance(scratchpad[0], KeyError)

        p.close()
        p.join()

    @unittest.skip(f"DRAGON: Bug filed PE-41882")
    def test_unpickleable_result(self):
        from multiprocessing.pool import MaybeEncodingError

        p = multiprocessing.Pool(2)

        # Make sure we don't lose pool processes because of encoding errors.
        for iteration in range(20):

            scratchpad = [None]

            def errback(exc):
                scratchpad[0] = exc

            res = p.apply_async(unpickleable_result, error_callback=errback)
            self.assertRaises(MaybeEncodingError, res.get)
            wrapped = scratchpad[0]
            self.assertTrue(wrapped)
            self.assertIsInstance(scratchpad[0], MaybeEncodingError)
            self.assertIsNotNone(wrapped.exc)
            self.assertIsNotNone(wrapped.value)

        p.close()
        p.join()


class WithProcessesTestPoolWorkerLifetime(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes', )

    @unittest.skip(f"bug filed CIRRUS-1473")
    def test_pool_worker_lifetime(self):
        p = multiprocessing.Pool(3, maxtasksperchild=10)
        self.assertEqual(3, len(p._pool))
        origworkerpids = [w.pid for w in p._pool]
        # Run many tasks so each worker gets replaced (hopefully)
        results = []
        for i in range(100):
            results.append(p.apply_async(sqr, (i,)))
        # Fetch the results and verify we got the right answers,
        # also ensuring all the tasks have completed.
        for (j, res) in enumerate(results):
            self.assertEqual(res.get(), sqr(j))
        # Refill the pool
        p._repopulate_pool()
        # Wait until all workers are alive
        # (countdown * DELTA = 5 seconds max startup process time)
        countdown = 50
        while countdown and not all(w.is_alive() for w in p._pool):
            countdown -= 1
            time.sleep(DELTA)
        finalworkerpids = [w.pid for w in p._pool]
        # All pids should be assigned.  See issue #7805.
        self.assertNotIn(None, origworkerpids)
        self.assertNotIn(None, finalworkerpids)
        # Finally, check that the worker pids have changed
        self.assertNotEqual(sorted(origworkerpids), sorted(finalworkerpids))
        p.close()
        p.join()

    def test_pool_worker_lifetime_early_close(self):
        # Issue #10332: closing a pool whose workers have limited lifetimes
        # before all the tasks completed would make join() hang.
        p = multiprocessing.Pool(3, maxtasksperchild=1)
        results = []
        for i in range(6):
            results.append(p.apply_async(sqr, (i, 0.3)))
        p.close()
        p.join()
        # check the results
        for (j, res) in enumerate(results):
            self.assertEqual(res.get(), sqr(j))

    def test_worker_finalization_via_atexit_handler_of_multiprocessing(self):
        # tests cases against bpo-38744 and bpo-39360
        cmd = """if 1:
            from multiprocessing import Pool
            problem = None
            class A:
                def __init__(self):
                    self.pool = Pool(processes=1)
            def test():
                global problem
                problem = A()
                problem.pool.map(float, tuple(range(10)))
            if __name__ == "__main__":
                test()
        """
        rc, out, err = test.support.script_helper.assert_python_ok("-c", cmd)
        self.assertEqual(rc, 0)


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestPool(BaseTestCase, ManagerMixin, unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.pool = cls.Pool(4)

    @classmethod
    def tearDownClass(cls):
        cls.pool.terminate()
        cls.pool.join()
        cls.pool = None
        super().tearDownClass()

    def test_apply(self):
        papply = self.pool.apply
        self.assertEqual(papply(sqr, (5,)), sqr(5))
        self.assertEqual(papply(sqr, (), {"x": 3}), sqr(x=3))

    def test_map(self):
        pmap = self.pool.map
        self.assertEqual(pmap(sqr, list(range(10))), list(map(sqr, list(range(10)))))
        self.assertEqual(pmap(sqr, list(range(100)), chunksize=20), list(map(sqr, list(range(100)))))

    def test_starmap(self):
        psmap = self.pool.starmap
        tuples = list(zip(range(10), range(9, -1, -1)))
        self.assertEqual(psmap(mul, tuples), list(itertools.starmap(mul, tuples)))
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(psmap(mul, tuples, chunksize=20), list(itertools.starmap(mul, tuples)))

    def test_starmap_async(self):
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(self.pool.starmap_async(mul, tuples).get(), list(itertools.starmap(mul, tuples)))

    def test_map_async(self):
        self.assertEqual(self.pool.map_async(sqr, list(range(10))).get(), list(map(sqr, list(range(10)))))

    def test_map_async_callbacks(self):
        call_args = self.manager.list() if self.TYPE == "manager" else []
        self.pool.map_async(int, ["1"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(1, len(call_args))
        self.assertEqual([1], call_args[0])
        self.pool.map_async(int, ["a"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(2, len(call_args))
        self.assertIsInstance(call_args[1], ValueError)

    def test_map_unplicklable(self):
        # Issue #19425 -- failure to pickle should not cause a hang
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class A(object):
            def __reduce__(self):
                raise RuntimeError("cannot pickle")

        with self.assertRaises(RuntimeError):
            self.pool.map(sqr, [A()] * 10)

    def test_map_chunksize(self):
        try:
            self.pool.map_async(sqr, [], chunksize=1).get(timeout=TIMEOUT1)
        except multiprocessing.TimeoutError:
            self.fail("pool.map_async with chunksize stalled on null list")

    def test_map_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(1, -1), 1)
        # again, make sure it's reentrant
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(1, -1), 1)

        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(10, 3), 1)

        class SpecialIterable:
            def __iter__(self):
                return self

            def __next__(self):
                raise SayWhenError

            def __len__(self):
                return 1

        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, SpecialIterable(), 1)
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, SpecialIterable(), 1)

    def test_async(self):
        res = self.pool.apply_async(
            sqr,
            (
                7,
                TIMEOUT1,
            ),
        )
        get = TimingWrapper(res.get)
        self.assertEqual(get(), 49)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT1)

    def test_async_timeout(self):
        res = self.pool.apply_async(sqr, (6, TIMEOUT2 + 1.0))
        get = TimingWrapper(res.get)
        self.assertRaises(multiprocessing.TimeoutError, get, timeout=TIMEOUT2)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT2)

    def test_imap(self):
        it = self.pool.imap(sqr, list(range(10)))
        self.assertEqual(list(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap(sqr, list(range(10)))
        for i in range(10):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

        it = self.pool.imap(sqr, list(range(1000)), chunksize=100)
        for i in range(1000):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

    def test_imap_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        it = self.pool.imap(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)
        # again, make sure it's reentrant
        it = self.pool.imap(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)

        it = self.pool.imap(sqr, exception_throwing_generator(10, 3), 1)
        for i in range(3):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)

        # SayWhenError seen at start of problematic chunk's results
        it = self.pool.imap(sqr, exception_throwing_generator(20, 7), 2)
        for i in range(6):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)
        it = self.pool.imap(sqr, exception_throwing_generator(20, 7), 4)
        for i in range(4):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)

    def test_imap_unordered(self):
        it = self.pool.imap_unordered(sqr, list(range(10)))
        self.assertEqual(sorted(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap_unordered(sqr, list(range(1000)), chunksize=100)
        self.assertEqual(sorted(it), list(map(sqr, list(range(1000)))))

    def test_imap_unordered_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        it = self.pool.imap_unordered(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)
        # again, make sure it's reentrant
        it = self.pool.imap_unordered(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)

        it = self.pool.imap_unordered(sqr, exception_throwing_generator(10, 3), 1)
        expected_values = list(map(sqr, list(range(10))))
        with self.assertRaises(SayWhenError):
            # imap_unordered makes it difficult to anticipate the SayWhenError
            for i in range(10):
                value = next(it)
                self.assertIn(value, expected_values)
                expected_values.remove(value)

        it = self.pool.imap_unordered(sqr, exception_throwing_generator(20, 7), 2)
        expected_values = list(map(sqr, list(range(20))))
        with self.assertRaises(SayWhenError):
            for i in range(20):
                value = next(it)
                self.assertIn(value, expected_values)
                expected_values.remove(value)

    def test_make_pool(self):
        expected_error = RemoteError if self.TYPE == "manager" else ValueError

        self.assertRaises(expected_error, self.Pool, -1)
        self.assertRaises(expected_error, self.Pool, 0)

        if self.TYPE != "manager":
            p = self.Pool(3)
            try:
                self.assertEqual(3, len(p._pool))
            finally:
                p.close()
                p.join()

    def test_terminate(self):
        result = self.pool.map_async(time.sleep, [0.1 for i in range(10000)], chunksize=1)
        self.pool.terminate()
        join = TimingWrapper(self.pool.join)
        join()
        # Sanity check the pool didn't wait for all tasks to finish
        self.assertLess(join.elapsed, 2.0)

    def test_empty_iterable(self):
        # See Issue 12157
        p = self.Pool(1)

        self.assertEqual(p.map(sqr, []), [])
        self.assertEqual(list(p.imap(sqr, [])), [])
        self.assertEqual(list(p.imap_unordered(sqr, [])), [])
        self.assertEqual(p.map_async(sqr, []).get(), [])

        p.close()
        p.join()

    def test_context(self):
        if self.TYPE == "processes":
            L = list(range(10))
            expected = [sqr(i) for i in L]
            with self.Pool(2) as p:
                r = p.map_async(sqr, L)
                self.assertEqual(r.get(), expected)
            p.join()
            self.assertRaises(ValueError, p.map_async, sqr, L)

    @classmethod
    def _test_traceback(cls):
        raise RuntimeError(123)  # some comment

    def test_traceback(self):
        # We want ensure that the traceback from the child process is
        # contained in the traceback raised in the main process.
        if self.TYPE == "processes":
            with self.Pool(1) as p:
                try:
                    p.apply(self._test_traceback)
                except Exception as e:
                    exc = e
                else:
                    self.fail("expected RuntimeError")
            p.join()
            self.assertIs(type(exc), RuntimeError)
            self.assertEqual(exc.args, (123,))
            cause = exc.__cause__
            self.assertIs(type(cause), multiprocessing.pool.RemoteTraceback)
            self.assertIn("raise RuntimeError(123) # some comment", cause.tb)

            with test.support.captured_stderr() as f1:
                try:
                    raise exc
                except RuntimeError:
                    sys.excepthook(*sys.exc_info())
            self.assertIn("raise RuntimeError(123) # some comment", f1.getvalue())
            # _helper_reraises_exception should not make the error
            # a remote exception
            with self.Pool(1) as p:
                try:
                    p.map(sqr, exception_throwing_generator(1, -1), 1)
                except Exception as e:
                    exc = e
                else:
                    self.fail("expected SayWhenError")
                self.assertIs(type(exc), SayWhenError)
                self.assertIs(exc.__cause__, None)
            p.join()

    @classmethod
    def _test_wrapped_exception(cls):
        raise RuntimeError("foo")

    def test_wrapped_exception(self):
        # Issue #20980: Should not wrap exception when using thread pool
        with self.Pool(1) as p:
            with self.assertRaises(RuntimeError):
                p.apply(self._test_wrapped_exception)
        p.join()

    def test_map_no_failfast(self):
        # Issue #23992: the fail-fast behaviour when an exception is raised
        # during map() would make Pool.join() deadlock, because a worker
        # process would fill the result queue (after the result handler thread
        # terminated, hence not draining it anymore).

        t_start = time.monotonic()

        with self.assertRaises(ValueError):
            with self.Pool(2) as p:
                try:
                    p.map(raise_large_valuerror, [0, 1])
                finally:
                    time.sleep(0.5)
                    p.close()
                    p.join()

        # check that we indeed waited for all jobs
        self.assertGreater(time.monotonic() - t_start, 0.9)

    def test_release_task_refs(self):
        # Issue #29861: task arguments and results should not be kept
        # alive after we are done with them.
        objs = [CountedObject() for i in range(10)]
        refs = [weakref.ref(o) for o in objs]
        self.pool.map(identity, objs)

        del objs
        gc.collect()  # For PyPy or other GCs.
        time.sleep(DELTA)  # let threaded cleanup code run
        self.assertEqual(set(wr() for wr in refs), {None})
        # With a process pool, copies of the objects are returned, check
        # they were released too.
        self.assertEqual(CountedObject.n_instances, 0)

    def test_enter(self):
        if self.TYPE == "manager":
            self.skipTest("test not applicable to manager")

        pool = self.Pool(1)
        with pool:
            pass
            # call pool.terminate()
        # pool is no longer running

        with self.assertRaises(ValueError):
            # bpo-35477: pool.__enter__() fails if the pool is not running
            with pool:
                pass
        pool.join()

    def test_resource_warning(self):
        if self.TYPE == "manager":
            self.skipTest("test not applicable to manager")

        pool = self.Pool(1)
        pool.terminate()
        pool.join()

        # force state to RUN to emit ResourceWarning in __del__()
        pool._state = multiprocessing.pool.RUN

        with support.check_warnings(("unclosed running multiprocessing pool", ResourceWarning)):
            pool = None
            support.gc_collect()


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestPool(BaseTestCase, ThreadsMixin, unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.pool = cls.Pool(4)

    @classmethod
    def tearDownClass(cls):
        cls.pool.terminate()
        cls.pool.join()
        cls.pool = None
        super().tearDownClass()

    def test_apply(self):
        papply = self.pool.apply
        self.assertEqual(papply(sqr, (5,)), sqr(5))
        self.assertEqual(papply(sqr, (), {"x": 3}), sqr(x=3))

    def test_map(self):
        pmap = self.pool.map
        self.assertEqual(pmap(sqr, list(range(10))), list(map(sqr, list(range(10)))))
        self.assertEqual(pmap(sqr, list(range(100)), chunksize=20), list(map(sqr, list(range(100)))))

    def test_starmap(self):
        psmap = self.pool.starmap
        tuples = list(zip(range(10), range(9, -1, -1)))
        self.assertEqual(psmap(mul, tuples), list(itertools.starmap(mul, tuples)))
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(psmap(mul, tuples, chunksize=20), list(itertools.starmap(mul, tuples)))

    def test_starmap_async(self):
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(self.pool.starmap_async(mul, tuples).get(), list(itertools.starmap(mul, tuples)))

    def test_map_async(self):
        self.assertEqual(self.pool.map_async(sqr, list(range(10))).get(), list(map(sqr, list(range(10)))))

    def test_map_async_callbacks(self):
        call_args = self.manager.list() if self.TYPE == "manager" else []
        self.pool.map_async(int, ["1"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(1, len(call_args))
        self.assertEqual([1], call_args[0])
        self.pool.map_async(int, ["a"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(2, len(call_args))
        self.assertIsInstance(call_args[1], ValueError)

    def test_map_unplicklable(self):
        # Issue #19425 -- failure to pickle should not cause a hang
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class A(object):
            def __reduce__(self):
                raise RuntimeError("cannot pickle")

        with self.assertRaises(RuntimeError):
            self.pool.map(sqr, [A()] * 10)

    def test_map_chunksize(self):
        try:
            self.pool.map_async(sqr, [], chunksize=1).get(timeout=TIMEOUT1)
        except multiprocessing.TimeoutError:
            self.fail("pool.map_async with chunksize stalled on null list")

    def test_map_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(1, -1), 1)
        # again, make sure it's reentrant
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(1, -1), 1)

        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, exception_throwing_generator(10, 3), 1)

        class SpecialIterable:
            def __iter__(self):
                return self

            def __next__(self):
                raise SayWhenError

            def __len__(self):
                return 1

        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, SpecialIterable(), 1)
        with self.assertRaises(SayWhenError):
            self.pool.map(sqr, SpecialIterable(), 1)

    def test_async(self):
        res = self.pool.apply_async(
            sqr,
            (
                7,
                TIMEOUT1,
            ),
        )
        get = TimingWrapper(res.get)
        self.assertEqual(get(), 49)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT1)

    def test_async_timeout(self):
        res = self.pool.apply_async(sqr, (6, TIMEOUT2 + 1.0))
        get = TimingWrapper(res.get)
        self.assertRaises(multiprocessing.TimeoutError, get, timeout=TIMEOUT2)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT2)

    def test_imap(self):
        it = self.pool.imap(sqr, list(range(10)))
        self.assertEqual(list(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap(sqr, list(range(10)))
        for i in range(10):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

        it = self.pool.imap(sqr, list(range(1000)), chunksize=100)
        for i in range(1000):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

    def test_imap_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        it = self.pool.imap(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)
        # again, make sure it's reentrant
        it = self.pool.imap(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)

        it = self.pool.imap(sqr, exception_throwing_generator(10, 3), 1)
        for i in range(3):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)

        # SayWhenError seen at start of problematic chunk's results
        it = self.pool.imap(sqr, exception_throwing_generator(20, 7), 2)
        for i in range(6):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)
        it = self.pool.imap(sqr, exception_throwing_generator(20, 7), 4)
        for i in range(4):
            self.assertEqual(next(it), i * i)
        self.assertRaises(SayWhenError, it.__next__)

    def test_imap_unordered(self):
        it = self.pool.imap_unordered(sqr, list(range(10)))
        self.assertEqual(sorted(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap_unordered(sqr, list(range(1000)), chunksize=100)
        self.assertEqual(sorted(it), list(map(sqr, list(range(1000)))))

    def test_imap_unordered_handle_iterable_exception(self):
        if self.TYPE == "manager":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # SayWhenError seen at the very first of the iterable
        it = self.pool.imap_unordered(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)
        # again, make sure it's reentrant
        it = self.pool.imap_unordered(sqr, exception_throwing_generator(1, -1), 1)
        self.assertRaises(SayWhenError, it.__next__)

        it = self.pool.imap_unordered(sqr, exception_throwing_generator(10, 3), 1)
        expected_values = list(map(sqr, list(range(10))))
        with self.assertRaises(SayWhenError):
            # imap_unordered makes it difficult to anticipate the SayWhenError
            for i in range(10):
                value = next(it)
                self.assertIn(value, expected_values)
                expected_values.remove(value)

        it = self.pool.imap_unordered(sqr, exception_throwing_generator(20, 7), 2)
        expected_values = list(map(sqr, list(range(20))))
        with self.assertRaises(SayWhenError):
            for i in range(20):
                value = next(it)
                self.assertIn(value, expected_values)
                expected_values.remove(value)

    def test_make_pool(self):
        expected_error = RemoteError if self.TYPE == "manager" else ValueError

        self.assertRaises(expected_error, self.Pool, -1)
        self.assertRaises(expected_error, self.Pool, 0)

        if self.TYPE != "manager":
            p = self.Pool(3)
            try:
                self.assertEqual(3, len(p._pool))
            finally:
                p.close()
                p.join()

    def test_terminate(self):
        result = self.pool.map_async(time.sleep, [0.1 for i in range(10000)], chunksize=1)
        self.pool.terminate()
        join = TimingWrapper(self.pool.join)
        join()
        # Sanity check the pool didn't wait for all tasks to finish
        self.assertLess(join.elapsed, 2.0)

    def test_empty_iterable(self):
        # See Issue 12157
        p = self.Pool(1)

        self.assertEqual(p.map(sqr, []), [])
        self.assertEqual(list(p.imap(sqr, [])), [])
        self.assertEqual(list(p.imap_unordered(sqr, [])), [])
        self.assertEqual(p.map_async(sqr, []).get(), [])

        p.close()
        p.join()

    def test_context(self):
        if self.TYPE == "processes":
            L = list(range(10))
            expected = [sqr(i) for i in L]
            with self.Pool(2) as p:
                r = p.map_async(sqr, L)
                self.assertEqual(r.get(), expected)
            p.join()
            self.assertRaises(ValueError, p.map_async, sqr, L)

    @classmethod
    def _test_traceback(cls):
        raise RuntimeError(123)  # some comment

    def test_traceback(self):
        # We want ensure that the traceback from the child process is
        # contained in the traceback raised in the main process.
        if self.TYPE == "processes":
            with self.Pool(1) as p:
                try:
                    p.apply(self._test_traceback)
                except Exception as e:
                    exc = e
                else:
                    self.fail("expected RuntimeError")
            p.join()
            self.assertIs(type(exc), RuntimeError)
            self.assertEqual(exc.args, (123,))
            cause = exc.__cause__
            self.assertIs(type(cause), multiprocessing.pool.RemoteTraceback)
            self.assertIn("raise RuntimeError(123) # some comment", cause.tb)

            with test.support.captured_stderr() as f1:
                try:
                    raise exc
                except RuntimeError:
                    sys.excepthook(*sys.exc_info())
            self.assertIn("raise RuntimeError(123) # some comment", f1.getvalue())
            # _helper_reraises_exception should not make the error
            # a remote exception
            with self.Pool(1) as p:
                try:
                    p.map(sqr, exception_throwing_generator(1, -1), 1)
                except Exception as e:
                    exc = e
                else:
                    self.fail("expected SayWhenError")
                self.assertIs(type(exc), SayWhenError)
                self.assertIs(exc.__cause__, None)
            p.join()

    @classmethod
    def _test_wrapped_exception(cls):
        raise RuntimeError("foo")

    def test_wrapped_exception(self):
        # Issue #20980: Should not wrap exception when using thread pool
        with self.Pool(1) as p:
            with self.assertRaises(RuntimeError):
                p.apply(self._test_wrapped_exception)
        p.join()

    def test_map_no_failfast(self):
        # Issue #23992: the fail-fast behaviour when an exception is raised
        # during map() would make Pool.join() deadlock, because a worker
        # process would fill the result queue (after the result handler thread
        # terminated, hence not draining it anymore).

        t_start = time.monotonic()

        with self.assertRaises(ValueError):
            with self.Pool(2) as p:
                try:
                    p.map(raise_large_valuerror, [0, 1])
                finally:
                    time.sleep(0.5)
                    p.close()
                    p.join()

        # check that we indeed waited for all jobs
        self.assertGreater(time.monotonic() - t_start, 0.9)

    def test_release_task_refs(self):
        # Issue #29861: task arguments and results should not be kept
        # alive after we are done with them.
        objs = [CountedObject() for i in range(10)]
        refs = [weakref.ref(o) for o in objs]
        self.pool.map(identity, objs)

        del objs
        gc.collect()  # For PyPy or other GCs.
        time.sleep(DELTA)  # let threaded cleanup code run
        self.assertEqual(set(wr() for wr in refs), {None})
        # With a process pool, copies of the objects are returned, check
        # they were released too.
        self.assertEqual(CountedObject.n_instances, 0)

    def test_enter(self):
        if self.TYPE == "manager":
            self.skipTest("test not applicable to manager")

        pool = self.Pool(1)
        with pool:
            pass
            # call pool.terminate()
        # pool is no longer running

        with self.assertRaises(ValueError):
            # bpo-35477: pool.__enter__() fails if the pool is not running
            with pool:
                pass
        pool.join()

    def test_resource_warning(self):
        if self.TYPE == "manager":
            self.skipTest("test not applicable to manager")

        pool = self.Pool(1)
        pool.terminate()
        pool.join()

        # force state to RUN to emit ResourceWarning in __del__()
        pool._state = multiprocessing.pool.RUN

        with support.check_warnings(("unclosed running multiprocessing pool", ResourceWarning)):
            pool = None
            support.gc_collect()

# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_pool.py
DragonHPC/dragon
"""Test multiprocessing.Process()
"""
import os
import sys
import time
import weakref
import signal
import io
import gc

import unittest
import test.support
from test import support

import threading

import dragon  # DRAGON import before multiprocessing

import multiprocessing

from common import (
    BaseTestCase,
    ProcessesMixin,
    ThreadsMixin,
    TimingWrapper,
    DELTA,
    join_process,
    close_queue,
    wait_for_handle,
    setUpModule,
    tearDownModule,
)

#
#
#


class DummyCallable:
    """Just a Dummy"""

    def __call__(self, q, c):
        assert isinstance(c, DummyCallable)
        q.put(5)


class WithProcessesTestProcess(BaseTestCase, ProcessesMixin, unittest.TestCase):
    """Test multiprocessing.Process()"""

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    def test_current(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        current = self.current_process()
        authkey = current.authkey

        self.assertTrue(current.is_alive())
        self.assertTrue(not current.daemon)
        self.assertIsInstance(authkey, bytes)
        self.assertTrue(len(authkey) > 0)
        self.assertEqual(current.ident, os.getpid())
        self.assertEqual(current.exitcode, None)

    def test_daemon_argument(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # By default uses the current process's daemon flag.
        proc0 = self.Process(target=self._test)
        self.assertEqual(proc0.daemon, self.current_process().daemon)
        proc1 = self.Process(target=self._test, daemon=True)
        self.assertTrue(proc1.daemon)
        proc2 = self.Process(target=self._test, daemon=False)
        self.assertFalse(proc2.daemon)

    @classmethod
    def _test(cls, q, *args, **kwds):
        current = cls.current_process()
        q.put(args)
        q.put(kwds)
        q.put(current.name)
        if cls.TYPE != "threads":
            q.put(bytes(current.authkey))
            q.put(current.pid)

    def test_parent_process_attributes(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        self.assertIsNone(self.parent_process())

        rconn, wconn = self.Pipe(duplex=False)
        p = self.Process(target=self._test_send_parent_process, args=(wconn,))
        p.start()
        p.join()
        parent_pid, parent_name = rconn.recv()
        self.assertEqual(parent_pid, self.current_process().pid)
        self.assertEqual(parent_pid, os.getpid())
        self.assertEqual(parent_name, self.current_process().name)

    @classmethod
    def _test_send_parent_process(cls, wconn):
        from multiprocessing.process import parent_process

        wconn.send([parent_process().pid, parent_process().name])

    def test_parent_process(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # Launch a child process. Make it launch a grandchild process. Kill the
        # child process and make sure that the grandchild notices the death of
        # its parent (a.k.a the child process).
        rconn, wconn = self.Pipe(duplex=False)
        p = self.Process(target=self._test_create_grandchild_process, args=(wconn,))
        p.start()

        if not rconn.poll(timeout=support.LONG_TIMEOUT):
            raise AssertionError("Could not communicate with child process")
        parent_process_status = rconn.recv()
        self.assertEqual(parent_process_status, "alive")

        p.terminate()
        p.join()

        if not rconn.poll(timeout=support.LONG_TIMEOUT):
            raise AssertionError("Could not communicate with child process")
        parent_process_status = rconn.recv()
        self.assertEqual(parent_process_status, "not alive")

    @classmethod
    def _test_create_grandchild_process(cls, wconn):
        p = cls.Process(target=cls._test_report_parent_status, args=(wconn,))
        p.start()
        time.sleep(300)

    @classmethod
    def _test_report_parent_status(cls, wconn):
        from multiprocessing.process import parent_process

        wconn.send("alive" if parent_process().is_alive() else "not alive")
        parent_process().join(timeout=support.SHORT_TIMEOUT)
        wconn.send("alive" if parent_process().is_alive() else "not alive")

    def test_process(self):
        q = self.Queue(1)
        e = self.Event()
        args = (q, 1, 2)
        kwargs = {"hello": 23, "bye": 2.54}
        name = "SomeProcess"
        p = self.Process(target=self._test, args=args, kwargs=kwargs, name=name)
        p.daemon = True
        current = self.current_process()

        if self.TYPE != "threads":
            self.assertEqual(p.authkey, current.authkey)
        self.assertEqual(p.is_alive(), False)
        self.assertEqual(p.daemon, True)
        self.assertNotIn(p, self.active_children())
        self.assertTrue(type(self.active_children()) is list)
        self.assertEqual(p.exitcode, None)

        p.start()

        self.assertEqual(p.exitcode, None)
        self.assertEqual(p.is_alive(), True)
        self.assertIn(p, self.active_children())

        self.assertEqual(q.get(), args[1:])
        self.assertEqual(q.get(), kwargs)
        self.assertEqual(q.get(), p.name)
        if self.TYPE != "threads":
            self.assertEqual(q.get(), current.authkey)
            self.assertEqual(q.get(), p.pid)

        p.join()

        self.assertEqual(p.exitcode, 0)
        self.assertEqual(p.is_alive(), False)
        self.assertNotIn(p, self.active_children())
        close_queue(q)

    @unittest.skipUnless(threading._HAVE_THREAD_NATIVE_ID, "needs native_id")
    def test_process_mainthread_native_id(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        current_mainthread_native_id = threading.main_thread().native_id

        q = self.Queue(1)
        p = self.Process(target=self._test_process_mainthread_native_id, args=(q,))
        p.start()

        child_mainthread_native_id = q.get()
        p.join()
        close_queue(q)

        self.assertNotEqual(current_mainthread_native_id, child_mainthread_native_id)

    @classmethod
    def _test_process_mainthread_native_id(cls, q):
        mainthread_native_id = threading.main_thread().native_id
        q.put(mainthread_native_id)

    @classmethod
    def _sleep_some(cls):
        time.sleep(100)

    @classmethod
    def _test_sleep(cls, delay):
        time.sleep(delay)

    def _kill_process(self, meth):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        p = self.Process(target=self._sleep_some)
        p.daemon = True
        p.start()

        self.assertEqual(p.is_alive(), True)
        self.assertIn(p, self.active_children())
        self.assertEqual(p.exitcode, None)

        join = TimingWrapper(p.join)

        self.assertEqual(join(0), None)
        self.assertTimingAlmostEqual(join.elapsed, 0.0)
        self.assertEqual(p.is_alive(), True)

        self.assertEqual(join(-1), None)
        self.assertTimingAlmostEqual(join.elapsed, 0.0)
        self.assertEqual(p.is_alive(), True)

        # XXX maybe terminating too soon causes the problems on Gentoo...
        time.sleep(1)

        meth(p)

        if hasattr(signal, "alarm"):
            # On the Gentoo buildbot waitpid() often seems to block forever.
            # We use alarm() to interrupt it if it blocks for too long.
            def handler(*args):
                raise RuntimeError("join took too long: %s" % p)

            old_handler = signal.signal(signal.SIGALRM, handler)
            try:
                signal.alarm(10)
                self.assertEqual(join(), None)
            finally:
                signal.alarm(0)
                signal.signal(signal.SIGALRM, old_handler)
        else:
            self.assertEqual(join(), None)

        self.assertTimingAlmostEqual(join.elapsed, 0.0)

        self.assertEqual(p.is_alive(), False)
        self.assertNotIn(p, self.active_children())

        p.join()

        return p.exitcode

    def test_terminate(self):
        exitcode = self._kill_process(multiprocessing.Process.terminate)
        if os.name != "nt":
            self.assertEqual(exitcode, -signal.SIGTERM)

    def test_kill(self):
        exitcode = self._kill_process(multiprocessing.Process.kill)
        if os.name != "nt":
            self.assertEqual(exitcode, -signal.SIGKILL)

    def test_cpu_count(self):
        try:
            cpus = multiprocessing.cpu_count()
        except NotImplementedError:
            cpus = 1
        self.assertTrue(type(cpus) is int)
        self.assertTrue(cpus >= 1)

    def test_active_children(self):
        self.assertEqual(type(self.active_children()), list)

        p = self.Process(target=time.sleep, args=(DELTA,))
        self.assertNotIn(p, self.active_children())

        p.daemon = True
        p.start()
        self.assertIn(p, self.active_children())

        p.join()
        self.assertNotIn(p, self.active_children())

    @classmethod
    def _test_recursion(cls, wconn, id):
        wconn.send(id)
        if len(id) < 2:
            for i in range(2):
                p = cls.Process(target=cls._test_recursion, args=(wconn, id + [i]))
                p.start()
                p.join()

    @unittest.skip("bug filed PE-40917")
    def test_recursion(self):
        rconn, wconn = self.Pipe(duplex=False)
        self._test_recursion(wconn, [])

        time.sleep(DELTA)
        result = []
        while rconn.poll():
            result.append(rconn.recv())

        expected = [[], [0], [0, 0], [0, 1], [1], [1, 0], [1, 1]]
        self.assertEqual(result, expected)

    @classmethod
    def _test_sentinel(cls, event):
        event.wait(10.0)

    def test_sentinel(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        event = self.Event()
        p = self.Process(target=self._test_sentinel, args=(event,))
        with self.assertRaises(ValueError):
            p.sentinel
        p.start()
        self.addCleanup(p.join)
        sentinel = p.sentinel
        self.assertIsInstance(sentinel, int)
        self.assertFalse(wait_for_handle(sentinel, timeout=0.0))
        event.set()
        p.join()
        self.assertTrue(wait_for_handle(sentinel, timeout=1))

    @classmethod
    def _test_close(cls, rc=0, q=None):
        if q is not None:
            q.get()
        sys.exit(rc)

    def test_close(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        q = self.Queue()
        p = self.Process(target=self._test_close, kwargs={"q": q})
        p.daemon = True
        p.start()
        self.assertEqual(p.is_alive(), True)
        # Child is still alive, cannot close
        with self.assertRaises(ValueError):
            p.close()

        q.put(None)
        p.join()
        self.assertEqual(p.is_alive(), False)
        self.assertEqual(p.exitcode, 0)
        p.close()
        with self.assertRaises(ValueError):
            p.is_alive()
        with self.assertRaises(ValueError):
            p.join()
        with self.assertRaises(ValueError):
            p.terminate()
        p.close()

        wr = weakref.ref(p)
        del p
        gc.collect()
        self.assertIs(wr(), None)

        close_queue(q)

    def test_many_processes(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        sm = multiprocessing.get_start_method()
        N = 5 if sm == "spawn" else 100

        # Try to overwhelm the forkserver loop with events
        procs = [self.Process(target=self._test_sleep, args=(0.01,)) for i in range(N)]
        for p in procs:
            p.start()
        for p in procs:
            join_process(p)
        for p in procs:
            self.assertEqual(p.exitcode, 0)

        procs = [self.Process(target=self._sleep_some) for i in range(N)]
        for p in procs:
            p.start()
        time.sleep(0.001)  # let the children start...
        for p in procs:
            p.terminate()
        for p in procs:
            join_process(p)
        if os.name != "nt":
            exitcodes = [-signal.SIGTERM]
            if sys.platform == "darwin":
                # bpo-31510: On macOS, killing a freshly started process with
                # SIGTERM sometimes kills the process with SIGKILL.
                exitcodes.append(-signal.SIGKILL)
            for p in procs:
                self.assertIn(p.exitcode, exitcodes)

    def test_lose_target_ref(self):
        c = DummyCallable()
        wr = weakref.ref(c)
        q = self.Queue()
        p = self.Process(target=c, args=(q, c))
        del c
        p.start()
        p.join()
        gc.collect()  # For PyPy or other GCs.
        self.assertIs(wr(), None)
        self.assertEqual(q.get(), 5)
        close_queue(q)

    @classmethod
    def _test_child_fd_inflation(self, evt, q):
        q.put(test.support.fd_count())
        evt.wait()

    def test_child_fd_inflation(self):
        # Number of fds in child processes should not grow with the
        # number of running children.
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        sm = multiprocessing.get_start_method()
        if sm == "fork":
            # The fork method by design inherits all fds from the parent,
            # trying to go against it is a lost battle
            self.skipTest("test not appropriate for {}".format(sm))

        N = 5
        evt = self.Event()
        q = self.Queue()

        procs = [self.Process(target=self._test_child_fd_inflation, args=(evt, q)) for i in range(N)]
        for p in procs:
            p.start()

        try:
            fd_counts = [q.get() for i in range(N)]
            self.assertEqual(len(set(fd_counts)), 1, fd_counts)

        finally:
            evt.set()
            for p in procs:
                p.join()
            close_queue(q)

    @classmethod
    def _test_wait_for_threads(self, evt):
        def func1():
            time.sleep(0.5)
            evt.set()

        def func2():
            time.sleep(20)
            evt.clear()

        threading.Thread(target=func1).start()
        threading.Thread(target=func2, daemon=True).start()

    def test_wait_for_threads(self):
        # A child process should wait for non-daemonic threads to end
        # before exiting
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        evt = self.Event()
        proc = self.Process(target=self._test_wait_for_threads, args=(evt,))
        proc.start()
        proc.join()
        self.assertTrue(evt.is_set())

    @classmethod
    def _test_error_on_stdio_flush(self, evt, break_std_streams={}):
        for stream_name, action in break_std_streams.items():
            if action == "close":
                stream = io.StringIO()
                stream.close()
            else:
                assert action == "remove"
                stream = None
            setattr(sys, stream_name, None)
        evt.set()

    def test_error_on_stdio_flush_1(self):
        # Check that Process works with broken standard streams
        streams = [io.StringIO(), None]
        streams[0].close()
        for stream_name in ("stdout", "stderr"):
            for stream in streams:
                old_stream = getattr(sys, stream_name)
                setattr(sys, stream_name, stream)
                try:
                    evt = self.Event()
                    proc = self.Process(target=self._test_error_on_stdio_flush, args=(evt,))
                    proc.start()
                    proc.join()
                    self.assertTrue(evt.is_set())
                    self.assertEqual(proc.exitcode, 0)
                finally:
                    setattr(sys, stream_name, old_stream)

    def test_error_on_stdio_flush_2(self):
        # Same as test_error_on_stdio_flush_1(), but standard streams are
        # broken by the child process
        for stream_name in ("stdout", "stderr"):
            for action in ("close", "remove"):
                old_stream = getattr(sys, stream_name)
                try:
                    evt = self.Event()
                    proc = self.Process(
                        target=self._test_error_on_stdio_flush, args=(evt, {stream_name: action})
                    )
                    proc.start()
                    proc.join()
                    self.assertTrue(evt.is_set())
                    self.assertEqual(proc.exitcode, 0)
                finally:
                    setattr(sys, stream_name, old_stream)

    @classmethod
    def _sleep_and_set_event(self, evt, delay=0.0):
        time.sleep(delay)
        evt.set()

    def check_forkserver_death(self, signum):
        # bpo-31308: if the forkserver process has died, we should still
        # be able to create and run new Process instances (the forkserver
        # is implicitly restarted).
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        sm = multiprocessing.get_start_method()
        # if sm != "forkserver":
        #     # The fork method by design inherits all fds from the parent,
        #     # trying to go against it is a lost battle
        #     self.skipTest("test not appropriate for {}".format(sm))

        from multiprocessing.forkserver import _forkserver

        _forkserver.ensure_running()

        # First process sleeps 500 ms
        delay = 0.5

        evt = self.Event()
        proc = self.Process(target=self._sleep_and_set_event, args=(evt, delay))
        proc.start()

        pid = _forkserver._forkserver_pid
        os.kill(pid, signum)
        # give time to the fork server to die and time to proc to complete
        time.sleep(delay * 2.0)

        evt2 = self.Event()
        proc2 = self.Process(target=self._sleep_and_set_event, args=(evt2,))
        proc2.start()
        proc2.join()
        self.assertTrue(evt2.is_set())
        self.assertEqual(proc2.exitcode, 0)

        proc.join()
        self.assertTrue(evt.is_set())
        self.assertIn(proc.exitcode, (0, 255))

    @unittest.skip("DRAGON: Semlock not implemented")
    def test_forkserver_sigint(self):
        # Catchable signal
        self.check_forkserver_death(signal.SIGINT)

    @unittest.skip("DRAGON: Semlock not implemented")
    def test_forkserver_sigkill(self):
        # Uncatchable signal
        if os.name != "nt":
            self.check_forkserver_death(signal.SIGKILL)


class _UpperCaser(multiprocessing.Process):
    def __init__(self):
        multiprocessing.Process.__init__(self)
        self.child_conn, self.parent_conn = multiprocessing.Pipe()

    def run(self):
        self.parent_conn.close()
        for s in iter(self.child_conn.recv, None):
            self.child_conn.send(s.upper())
        self.child_conn.close()

    def submit(self, s):
        assert type(s) is str
        self.parent_conn.send(s)
        return self.parent_conn.recv()

    def stop(self):
        self.parent_conn.send(None)
        self.parent_conn.close()
        self.child_conn.close()


class TestSubclassingProcess(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ("processes",)

    def test_subclassing(self):
        uppercaser = _UpperCaser()
        uppercaser.daemon = True
        uppercaser.start()
        self.assertEqual(uppercaser.submit("hello"), "HELLO")
        self.assertEqual(uppercaser.submit("world"), "WORLD")
        uppercaser.stop()
        uppercaser.join()

    def test_stderr_flush(self):
        # sys.stderr is flushed at process shutdown (issue #13812)
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        testfn = test.support.TESTFN
        self.addCleanup(test.support.unlink, testfn)
        proc = self.Process(target=self._test_stderr_flush, args=(testfn,))
        proc.start()
        proc.join()
        with open(testfn, "r") as f:
            err = f.read()
            # The whole traceback was printed
            self.assertIn("ZeroDivisionError", err)
            # DRAGON     self.assertIn("test_multiprocessing.py", err)
            self.assertIn("test_process.py", err)
            self.assertIn("1 / 0  # MARKER", err)

    @classmethod
    def _test_stderr_flush(cls, testfn):
        fd = os.open(testfn, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
        sys.stderr = open(fd, "w", closefd=False)
        1 / 0  # MARKER

    @classmethod
    def _test_sys_exit(cls, reason, testfn):
        fd = os.open(testfn, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
        sys.stderr = open(fd, "w", closefd=False)
        sys.exit(reason)

    def test_sys_exit(self):
        # See Issue 13854
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        testfn = test.support.TESTFN
        self.addCleanup(test.support.unlink, testfn)

        for reason in (
            [1, 2, 3],
            "ignore this",
        ):
            p = self.Process(target=self._test_sys_exit, args=(reason, testfn))
            p.daemon = True
            p.start()
            join_process(p)
            self.assertEqual(p.exitcode, 1)

            with open(testfn, "r") as f:
                content = f.read()
            self.assertEqual(content.rstrip(), str(reason))

            os.unlink(testfn)

        cases = [
            ((True,), 1),
            ((False,), 0),
            ((8,), 8),
            ((None,), 0),
            ((), 0),
        ]

        for args, expected in cases:
            with self.subTest(args=args):
                p = self.Process(target=sys.exit, args=args)
                p.daemon = True
                p.start()
                join_process(p)
                self.assertEqual(p.exitcode, expected)


class WithThreadsTestProcess(BaseTestCase, ThreadsMixin, unittest.TestCase):
    """Test multiprocessing.Process()"""

    # DRAGON ALLOWED_TYPES = ('processes', 'threads')

    def test_current(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        current = self.current_process()
        authkey = current.authkey

        self.assertTrue(current.is_alive())
        self.assertTrue(not current.daemon)
        self.assertIsInstance(authkey, bytes)
        self.assertTrue(len(authkey) > 0)
        self.assertEqual(current.ident, os.getpid())
        self.assertEqual(current.exitcode, None)

    def test_daemon_argument(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # By default uses the current process's daemon flag.
        proc0 = self.Process(target=self._test)
        self.assertEqual(proc0.daemon, self.current_process().daemon)
        proc1 = self.Process(target=self._test, daemon=True)
        self.assertTrue(proc1.daemon)
        proc2 = self.Process(target=self._test, daemon=False)
        self.assertFalse(proc2.daemon)

    @classmethod
    def _test(cls, q, *args, **kwds):
        current = cls.current_process()
        q.put(args)
        q.put(kwds)
        q.put(current.name)
        if cls.TYPE != "threads":
            q.put(bytes(current.authkey))
            q.put(current.pid)

    def test_parent_process_attributes(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        self.assertIsNone(self.parent_process())

        rconn, wconn = self.Pipe(duplex=False)
        p = self.Process(target=self._test_send_parent_process, args=(wconn,))
        p.start()
        p.join()
        parent_pid, parent_name = rconn.recv()
        self.assertEqual(parent_pid, self.current_process().pid)
        self.assertEqual(parent_pid, os.getpid())
        self.assertEqual(parent_name, self.current_process().name)

    @classmethod
    def _test_send_parent_process(cls, wconn):
        from multiprocessing.process import parent_process

        wconn.send([parent_process().pid, parent_process().name])

    def test_parent_process(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        # Launch a child process. Make it launch a grandchild process. Kill the
        # child process and make sure that the grandchild notices the death of
        # its parent (a.k.a the child process).
        rconn, wconn = self.Pipe(duplex=False)
        p = self.Process(target=self._test_create_grandchild_process, args=(wconn,))
        p.start()

        if not rconn.poll(timeout=support.LONG_TIMEOUT):
            raise AssertionError("Could not communicate with child process")
        parent_process_status = rconn.recv()
        self.assertEqual(parent_process_status, "alive")

        p.terminate()
        p.join()

        if not rconn.poll(timeout=support.LONG_TIMEOUT):
            raise AssertionError("Could not communicate with child process")
        parent_process_status = rconn.recv()
        self.assertEqual(parent_process_status, "not alive")

    @classmethod
    def _test_create_grandchild_process(cls, wconn):
        p = cls.Process(target=cls._test_report_parent_status, args=(wconn,))
        p.start()
        time.sleep(300)

    @classmethod
    def _test_report_parent_status(cls, wconn):
        from multiprocessing.process import parent_process

        wconn.send("alive" if parent_process().is_alive() else "not alive")
        parent_process().join(timeout=support.SHORT_TIMEOUT)
        wconn.send("alive" if parent_process().is_alive() else "not alive")

    @unittest.skip("bug filed PE-40908")
    def test_process(self):
        q = self.Queue(1)
        e = self.Event()
        args = (q, 1, 2)
        kwargs = {"hello": 23, "bye": 2.54}
        name = "SomeProcess"
        p = self.Process(target=self._test, args=args, kwargs=kwargs, name=name)
        p.daemon = True
        current = self.current_process()

        if self.TYPE != "threads":
            self.assertEqual(p.authkey, current.authkey)
        self.assertEqual(p.is_alive(), False)
        self.assertEqual(p.daemon, True)
        self.assertNotIn(p, self.active_children())
        self.assertTrue(type(self.active_children()) is list)
        self.assertEqual(p.exitcode, None)

        p.start()

        self.assertEqual(p.exitcode, None)
        self.assertEqual(p.is_alive(), True)
        self.assertIn(p, self.active_children())

        self.assertEqual(q.get(), args[1:])
        self.assertEqual(q.get(), kwargs)
        self.assertEqual(q.get(), p.name)
        if self.TYPE != "threads":
            self.assertEqual(q.get(), current.authkey)
            self.assertEqual(q.get(), p.pid)

        p.join()

        self.assertEqual(p.exitcode, 0)
        self.assertEqual(p.is_alive(), False)
        self.assertNotIn(p, self.active_children())
        close_queue(q)

    @unittest.skipUnless(threading._HAVE_THREAD_NATIVE_ID, "needs native_id")
    def test_process_mainthread_native_id(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        current_mainthread_native_id = threading.main_thread().native_id

        q = self.Queue(1)
        p = self.Process(target=self._test_process_mainthread_native_id, args=(q,))
        p.start()

        child_mainthread_native_id = q.get()
        p.join()
        close_queue(q)

        self.assertNotEqual(current_mainthread_native_id, child_mainthread_native_id)

    @classmethod
    def _test_process_mainthread_native_id(cls, q):
        mainthread_native_id = threading.main_thread().native_id
        q.put(mainthread_native_id)

    @classmethod
    def _sleep_some(cls):
        time.sleep(100)

    @classmethod
    def _test_sleep(cls, delay):
        time.sleep(delay)

    def _kill_process(self, meth):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        p = self.Process(target=self._sleep_some)
        p.daemon = True
        p.start()

        self.assertEqual(p.is_alive(), True)
        self.assertIn(p, self.active_children())
        self.assertEqual(p.exitcode, None)

        join = TimingWrapper(p.join)

        self.assertEqual(join(0), None)
        self.assertTimingAlmostEqual(join.elapsed, 0.0)
        self.assertEqual(p.is_alive(), True)

        self.assertEqual(join(-1), None)
        self.assertTimingAlmostEqual(join.elapsed, 0.0)
        self.assertEqual(p.is_alive(), True)

        # XXX maybe terminating too soon causes the problems on Gentoo...
        time.sleep(1)

        meth(p)

        if hasattr(signal, "alarm"):
            # On the Gentoo buildbot waitpid() often seems to block forever.
            # We use alarm() to interrupt it if it blocks for too long.
            def handler(*args):
                raise RuntimeError("join took too long: %s" % p)

            old_handler = signal.signal(signal.SIGALRM, handler)
            try:
                signal.alarm(10)
                self.assertEqual(join(), None)
            finally:
                signal.alarm(0)
                signal.signal(signal.SIGALRM, old_handler)
        else:
            self.assertEqual(join(), None)

        self.assertTimingAlmostEqual(join.elapsed, 0.0)

        self.assertEqual(p.is_alive(), False)
        self.assertNotIn(p, self.active_children())

        p.join()

        return p.exitcode

    def test_terminate(self):
        exitcode = self._kill_process(multiprocessing.Process.terminate)
        if os.name != "nt":
            self.assertEqual(exitcode, -signal.SIGTERM)

    def test_kill(self):
        exitcode = self._kill_process(multiprocessing.Process.kill)
        if os.name != "nt":
            self.assertEqual(exitcode, -signal.SIGKILL)

    def test_cpu_count(self):
        try:
            cpus = multiprocessing.cpu_count()
        except NotImplementedError:
            cpus = 1
        self.assertTrue(type(cpus) is int)
        self.assertTrue(cpus >= 1)

    def test_active_children(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        self.assertEqual(type(self.active_children()), list)

        p = self.Process(target=time.sleep, args=(DELTA,))
        self.assertNotIn(p, self.active_children())

        p.daemon = True
        p.start()
        self.assertIn(p, self.active_children())

        p.join()
        self.assertNotIn(p, self.active_children())

    @classmethod
    def _test_recursion(cls, wconn, id):
        wconn.send(id)
        if len(id) < 2:
            for i in range(2):
                p = cls.Process(target=cls._test_recursion, args=(wconn, id + [i]))
                p.start()
                p.join()

    @unittest.skip("bug filed PE-40908")
    def test_recursion(self):
        rconn, wconn = self.Pipe(duplex=False)
        self._test_recursion(wconn, [])

        time.sleep(DELTA)
        result = []
        while rconn.poll():
            result.append(rconn.recv())

        expected = [[], [0], [0, 0], [0, 1], [1], [1, 0], [1, 1]]
        self.assertEqual(result, expected)

    @classmethod
    def _test_sentinel(cls, event):
        event.wait(10.0)

    def test_sentinel(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        event = self.Event()
        p = self.Process(target=self._test_sentinel, args=(event,))
        with self.assertRaises(ValueError):
            p.sentinel
        p.start()
        self.addCleanup(p.join)
        sentinel = p.sentinel
        self.assertIsInstance(sentinel, int)
        self.assertFalse(wait_for_handle(sentinel, timeout=0.0))
        event.set()
        p.join()
        self.assertTrue(wait_for_handle(sentinel, timeout=1))

    @classmethod
    def _test_close(cls, rc=0, q=None):
        if q is not None:
            q.get()
        sys.exit(rc)

    def test_close(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        q = self.Queue()
        p = self.Process(target=self._test_close, kwargs={"q": q})
        p.daemon = True
        p.start()
        self.assertEqual(p.is_alive(), True)
        # Child is still alive, cannot close
        with self.assertRaises(ValueError):
            p.close()

        q.put(None)
        p.join()
        self.assertEqual(p.is_alive(), False)
        self.assertEqual(p.exitcode, 0)
        p.close()
        with self.assertRaises(ValueError):
            p.is_alive()
        with self.assertRaises(ValueError):
            p.join()
        with self.assertRaises(ValueError):
            p.terminate()
        p.close()

        wr = weakref.ref(p)
        del p
        gc.collect()
        self.assertIs(wr(), None)

        close_queue(q)

    def test_many_processes(self):
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        sm = multiprocessing.get_start_method()
        N = 5 if sm == "spawn" else 100

        # Try to overwhelm the forkserver loop with events
        procs = [self.Process(target=self._test_sleep, args=(0.01,)) for i in range(N)]
        for p in procs:
            p.start()
        for p in procs:
            join_process(p)
        for p in procs:
            self.assertEqual(p.exitcode, 0)

        procs = [self.Process(target=self._sleep_some) for i in range(N)]
        for p in procs:
            p.start()
        time.sleep(0.001)  # let the children start...
        for p in procs:
            p.terminate()
        for p in procs:
            join_process(p)
        if os.name != "nt":
            exitcodes = [-signal.SIGTERM]
            if sys.platform == "darwin":
                # bpo-31510: On macOS, killing a freshly started process with
                # SIGTERM sometimes kills the process with SIGKILL.
                exitcodes.append(-signal.SIGKILL)
            for p in procs:
                self.assertIn(p.exitcode, exitcodes)

    @unittest.skip("bug filed PE-40908")
    def test_lose_target_ref(self):
        c = DummyCallable()
        wr = weakref.ref(c)
        q = self.Queue()
        p = self.Process(target=c, args=(q, c))
        del c
        p.start()
        p.join()
        gc.collect()  # For PyPy or other GCs.
        self.assertIs(wr(), None)
        self.assertEqual(q.get(), 5)
        close_queue(q)

    @classmethod
    def _test_child_fd_inflation(self, evt, q):
        q.put(test.support.fd_count())
        evt.wait()

    @unittest.skip("DRAGON: Semlock not implemented")
    def test_child_fd_inflation(self):
        # Number of fds in child processes should not grow with the
        # number of running children.
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        sm = multiprocessing.get_start_method()
        if sm == "fork":
            # The fork method by design inherits all fds from the parent,
            # trying to go against it is a lost battle
            self.skipTest("test not appropriate for {}".format(sm))

        N = 5
        evt = self.Event()
        q = self.Queue()

        procs = [self.Process(target=self._test_child_fd_inflation, args=(evt, q)) for i in range(N)]
        for p in procs:
            p.start()

        try:
            fd_counts = [q.get() for i in range(N)]
            self.assertEqual(len(set(fd_counts)), 1, fd_counts)

        finally:
            evt.set()
            for p in procs:
                p.join()
            close_queue(q)

    @classmethod
    def _test_wait_for_threads(self, evt):
        def func1():
            time.sleep(0.5)
            evt.set()

        def func2():
            time.sleep(20)
            evt.clear()

        threading.Thread(target=func1).start()
        threading.Thread(target=func2, daemon=True).start()

    def test_wait_for_threads(self):
        # A child process should wait for non-daemonic threads to end
        # before exiting
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        evt = self.Event()
        proc = self.Process(target=self._test_wait_for_threads, args=(evt,))
        proc.start()
        proc.join()
        self.assertTrue(evt.is_set())

    @classmethod
    def _test_error_on_stdio_flush(self, evt, break_std_streams={}):
        for stream_name, action in break_std_streams.items():
            if action == "close":
                stream = io.StringIO()
                stream.close()
            else:
                assert action == "remove"
                stream = None
            setattr(sys, stream_name, None)
        evt.set()

    @unittest.skip("bug filed PE-40908")
    def test_error_on_stdio_flush_1(self):
        # Check that Process works with broken standard streams
        streams = [io.StringIO(), None]
        streams[0].close()
        for stream_name in ("stdout", "stderr"):
            for stream in streams:
                old_stream = getattr(sys, stream_name)
                setattr(sys, stream_name, stream)
                try:
                    evt = self.Event()
                    proc = self.Process(target=self._test_error_on_stdio_flush, args=(evt,))
                    proc.start()
                    proc.join()
                    self.assertTrue(evt.is_set())
                    self.assertEqual(proc.exitcode, 0)
                finally:
                    setattr(sys, stream_name, old_stream)

    @unittest.skip("bug filed PE-40908")
    def test_error_on_stdio_flush_2(self):
        # Same as test_error_on_stdio_flush_1(), but standard streams are
        # broken by the child process
        for stream_name in ("stdout", "stderr"):
            for action in ("close", "remove"):
                old_stream = getattr(sys, stream_name)
                try:
                    evt = self.Event()
                    proc = self.Process(
                        target=self._test_error_on_stdio_flush, args=(evt, {stream_name: action})
                    )
                    proc.start()
                    proc.join()
                    self.assertTrue(evt.is_set())
                    self.assertEqual(proc.exitcode, 0)
                finally:
                    setattr(sys, stream_name, old_stream)

    @classmethod
    def _sleep_and_set_event(self, evt, delay=0.0):
        time.sleep(delay)
        evt.set()

    def check_forkserver_death(self, signum):
        # bpo-31308: if the forkserver process has died, we should still
        # be able to create and run new Process instances (the forkserver
        # is implicitly restarted).
        if self.TYPE == "threads":
            self.skipTest("test not appropriate for {}".format(self.TYPE))
        sm = multiprocessing.get_start_method()
        if sm != "forkserver":
            # The fork method by design inherits all fds from the parent,
            # trying to go against it is a lost battle
            self.skipTest("test not appropriate for {}".format(sm))

        from multiprocessing.forkserver import _forkserver

        _forkserver.ensure_running()

        # First process sleeps 500 ms
        delay = 0.5

        evt = self.Event()
        proc = self.Process(target=self._sleep_and_set_event, args=(evt, delay))
        proc.start()

        pid = _forkserver._forkserver_pid
        os.kill(pid, signum)
        # give time to the fork server to die and time to proc to complete
        time.sleep(delay * 2.0)

        evt2 = self.Event()
        proc2 = self.Process(target=self._sleep_and_set_event, args=(evt2,))
        proc2.start()
        proc2.join()
        self.assertTrue(evt2.is_set())
        self.assertEqual(proc2.exitcode, 0)

        proc.join()
        self.assertTrue(evt.is_set())
        self.assertIn(proc.exitcode, (0, 255))

    def test_forkserver_sigint(self):
        # Catchable signal
        self.check_forkserver_death(signal.SIGINT)

    def test_forkserver_sigkill(self):
        # Uncatchable signal
        if os.name != "nt":
            self.check_forkserver_death(signal.SIGKILL)


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_process.py
DragonHPC/dragon
import unittest
import os
import time
import queue as pyqueue


import test.support
from test.support import hashlib_helper
from test import support

import dragon  # DRAGON import before multiprocessing

import multiprocessing

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    TimingWrapper,
    setUpModule,
    tearDownModule,
    close_queue,
    DELTA,
    TIMEOUT1,
    TIMEOUT2,
    TIMEOUT3,
)

#
#
#


def queue_empty(q):
    if hasattr(q, "empty"):
        return q.empty()
    else:
        return q.qsize() == 0


def queue_full(q, maxsize):
    if hasattr(q, "full"):
        return q.full()
    else:
        return q.qsize() == maxsize


class WithProcessesTestQueue(BaseTestCase, ProcessesMixin, unittest.TestCase):
    @classmethod
    def _test_put(cls, queue, child_can_start, parent_can_continue):
        child_can_start.wait()
        for i in range(6):
            queue.get()
        parent_can_continue.set()

    def test_put(self):
        MAXSIZE = 6
        queue = self.Queue(maxsize=MAXSIZE)
        child_can_start = self.Event()
        parent_can_continue = self.Event()

        proc = self.Process(target=self._test_put, args=(queue, child_can_start, parent_can_continue))
        proc.daemon = True
        proc.start()

        self.assertEqual(queue_empty(queue), True)
        self.assertEqual(queue_full(queue, MAXSIZE), False)

        queue.put(1)
        queue.put(2, True)
        queue.put(3, True, None)
        queue.put(4, False)
        queue.put(5, False, None)
        queue.put_nowait(6)

        # the values may be in buffer but not yet in pipe so sleep a bit
        time.sleep(DELTA)

        self.assertEqual(queue_empty(queue), False)
        self.assertEqual(queue_full(queue, MAXSIZE), True)

        put = TimingWrapper(queue.put)
        put_nowait = TimingWrapper(queue.put_nowait)

        self.assertRaises(pyqueue.Full, put, 7, False)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, False, None)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put_nowait, 7)
        self.assertTimingAlmostEqual(put_nowait.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, True, TIMEOUT1)
        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT1)

        self.assertRaises(pyqueue.Full, put, 7, False, TIMEOUT2)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, True, timeout=TIMEOUT3)
        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT3)

        child_can_start.set()
        parent_can_continue.wait()

        self.assertEqual(queue_empty(queue), True)
        self.assertEqual(queue_full(queue, MAXSIZE), False)

        proc.join()
        close_queue(queue)

    @classmethod
    def _test_get(cls, queue, child_can_start, parent_can_continue):
        child_can_start.wait()
        queue.put(1)
        queue.put(2)
        queue.put(3)
        queue.put(4)
        queue.put(5)
        parent_can_continue.set()

    def test_get(self):
        queue = self.Queue()
        child_can_start = self.Event()
        parent_can_continue = self.Event()

        proc = self.Process(target=self._test_get, args=(queue, child_can_start, parent_can_continue))
        proc.daemon = True
        proc.start()

        self.assertEqual(queue_empty(queue), True)

        child_can_start.set()
        parent_can_continue.wait()

        time.sleep(DELTA)
        self.assertEqual(queue_empty(queue), False)

        self.assertEqual(queue.get(), 1)
        self.assertEqual(queue.get(True, None), 2)
        self.assertEqual(queue.get(True), 3)
        self.assertEqual(queue.get(timeout=1), 4)
        self.assertEqual(queue.get_nowait(), 5)

        self.assertEqual(queue_empty(queue), True)

        get = TimingWrapper(queue.get)
        get_nowait = TimingWrapper(queue.get_nowait)

        self.assertRaises(pyqueue.Empty, get, False)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, False, None)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get_nowait)
        self.assertTimingAlmostEqual(get_nowait.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, True, TIMEOUT1)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT1)

        self.assertRaises(pyqueue.Empty, get, False, TIMEOUT2)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, timeout=TIMEOUT3)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT3)

        proc.join()
        close_queue(queue)

    @classmethod
    def _test_fork(cls, queue):
        for i in range(10, 20):
            queue.put(i)
        # note that at this point the items may only be buffered, so the
        # process cannot shutdown until the feeder thread has finished
        # pushing items onto the pipe.

    def test_fork(self):
        # Old versions of Queue would fail to create a new feeder
        # thread for a forked process if the original process had its
        # own feeder thread.  This test checks that this no longer
        # happens.

        queue = self.Queue()

        # put items on queue so that main process starts a feeder thread
        for i in range(10):
            queue.put(i)

        # wait to make sure thread starts before we fork a new process
        time.sleep(DELTA)

        # fork process
        p = self.Process(target=self._test_fork, args=(queue,))
        p.daemon = True
        p.start()

        # check that all expected items are in the queue
        for i in range(20):
            self.assertEqual(queue.get(), i)
        self.assertRaises(pyqueue.Empty, queue.get, False)

        p.join()
        close_queue(queue)

    def test_qsize(self):
        q = self.Queue()
        try:
            self.assertEqual(q.qsize(), 0)
        except NotImplementedError:
            self.skipTest("qsize method not implemented")
        q.put(1)
        self.assertEqual(q.qsize(), 1)
        q.put(5)
        self.assertEqual(q.qsize(), 2)
        q.get()
        self.assertEqual(q.qsize(), 1)
        q.get()
        self.assertEqual(q.qsize(), 0)
        close_queue(q)

    @classmethod
    def _test_task_done(cls, q):
        for obj in iter(q.get, None):
            time.sleep(DELTA)
            q.task_done()

    def test_task_done(self):
        queue = self.JoinableQueue()

        workers = [self.Process(target=self._test_task_done, args=(queue,)) for i in range(4)]

        for p in workers:
            p.daemon = True
            p.start()

        for i in range(10):
            queue.put(i)

        queue.join()

        for p in workers:
            queue.put(None)

        for p in workers:
            p.join()

        close_queue(queue)

    def test_no_import_lock_contention(self):
        with test.support.temp_cwd():
            module_name = "imported_by_an_imported_module"
            with open(module_name + ".py", "w") as f:
                f.write(
                    """if 1:
                    import multiprocessing

                    q = multiprocessing.Queue()
                    q.put('knock knock')
                    q.get(timeout=3)
                    q.close()
                    del q
                """
                )

            with test.support.DirsOnSysPath(os.getcwd()):
                try:
                    __import__(module_name)
                except pyqueue.Empty:
                    self.fail("Probable regression on import lock contention;" " see Issue #22853")

    def test_timeout(self):
        q = multiprocessing.Queue()
        start = time.monotonic()
        self.assertRaises(pyqueue.Empty, q.get, True, 0.200)
        delta = time.monotonic() - start
        # bpo-30317: Tolerate a delta of 100 ms because of the bad clock
        # resolution on Windows (usually 15.6 ms). x86 Windows7 3.x once
        # failed because the delta was only 135.8 ms.
        self.assertGreaterEqual(delta, 0.100)
        close_queue(q)

    def test_queue_feeder_donot_stop_onexc(self):
        # bpo-30414: verify feeder handles exceptions correctly
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class NotSerializable(object):
            def __reduce__(self):
                raise AttributeError

        with test.support.captured_stderr():
            q = self.Queue()
            q.put(NotSerializable())
            q.put(True)
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))
            close_queue(q)

        with test.support.captured_stderr():
            # bpo-33078: verify that the queue size is correctly handled
            # on errors.
            q = self.Queue(maxsize=1)
            q.put(NotSerializable())
            q.put(True)
            try:
                self.assertEqual(q.qsize(), 1)
            except NotImplementedError:
                # qsize is not available on all platform as it
                # relies on sem_getvalue
                pass
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))

            # Check that the size of the queue is correct
            self.assertTrue(q.empty())
            close_queue(q)

    def test_queue_feeder_on_queue_feeder_error(self):
        # bpo-30006: verify feeder handles exceptions using the
        # _on_queue_feeder_error hook.
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class NotSerializable(object):
            """Mock unserializable object"""

            def __init__(self):
                self.reduce_was_called = False
                self.on_queue_feeder_error_was_called = False

            def __reduce__(self):
                self.reduce_was_called = True
                raise AttributeError

        class SafeQueue(multiprocessing.queues.Queue):
            """Queue with overloaded _on_queue_feeder_error hook"""

            @staticmethod
            def _on_queue_feeder_error(e, obj):
                if isinstance(e, AttributeError) and isinstance(obj, NotSerializable):
                    obj.on_queue_feeder_error_was_called = True

        not_serializable_obj = NotSerializable()
        # The captured_stderr reduces the noise in the test report
        with test.support.captured_stderr():
            q = SafeQueue(ctx=multiprocessing.get_context())
            q.put(not_serializable_obj)

            # Verify that q is still functioning correctly
            q.put(True)
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))

        # Assert that the serialization and the hook have been called correctly
        self.assertTrue(not_serializable_obj.reduce_was_called)
        self.assertTrue(not_serializable_obj.on_queue_feeder_error_was_called)

    def test_closed_queue_put_get_exceptions(self):
        for q in multiprocessing.Queue(), multiprocessing.JoinableQueue():
            q.close()
            with self.assertRaisesRegex(ValueError, "is closed"):
                q.put("foo")
            with self.assertRaisesRegex(ValueError, "is closed"):
                q.get()


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestQueue(BaseTestCase, ManagerMixin, unittest.TestCase):
    @classmethod
    def _test_put(cls, queue, child_can_start, parent_can_continue):
        child_can_start.wait()
        for i in range(6):
            queue.get()
        parent_can_continue.set()

    def test_put(self):
        MAXSIZE = 6
        queue = self.Queue(maxsize=MAXSIZE)
        child_can_start = self.Event()
        parent_can_continue = self.Event()

        proc = self.Process(target=self._test_put, args=(queue, child_can_start, parent_can_continue))
        proc.daemon = True
        proc.start()

        self.assertEqual(queue_empty(queue), True)
        self.assertEqual(queue_full(queue, MAXSIZE), False)

        queue.put(1)
        queue.put(2, True)
        queue.put(3, True, None)
        queue.put(4, False)
        queue.put(5, False, None)
        queue.put_nowait(6)

        # the values may be in buffer but not yet in pipe so sleep a bit
        time.sleep(DELTA)

        self.assertEqual(queue_empty(queue), False)
        self.assertEqual(queue_full(queue, MAXSIZE), True)

        put = TimingWrapper(queue.put)
        put_nowait = TimingWrapper(queue.put_nowait)

        self.assertRaises(pyqueue.Full, put, 7, False)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, False, None)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put_nowait, 7)
        self.assertTimingAlmostEqual(put_nowait.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, True, TIMEOUT1)
        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT1)

        self.assertRaises(pyqueue.Full, put, 7, False, TIMEOUT2)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, True, timeout=TIMEOUT3)
        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT3)

        child_can_start.set()
        parent_can_continue.wait()

        self.assertEqual(queue_empty(queue), True)
        self.assertEqual(queue_full(queue, MAXSIZE), False)

        proc.join()
        close_queue(queue)

    @classmethod
    def _test_get(cls, queue, child_can_start, parent_can_continue):
        child_can_start.wait()
        # queue.put(1)
        queue.put(2)
        queue.put(3)
        queue.put(4)
        queue.put(5)
        parent_can_continue.set()

    def test_get(self):
        queue = self.Queue()
        child_can_start = self.Event()
        parent_can_continue = self.Event()

        proc = self.Process(target=self._test_get, args=(queue, child_can_start, parent_can_continue))
        proc.daemon = True
        proc.start()

        self.assertEqual(queue_empty(queue), True)

        child_can_start.set()
        parent_can_continue.wait()

        time.sleep(DELTA)
        self.assertEqual(queue_empty(queue), False)

        # Hangs unexpectedly, remove for now
        # self.assertEqual(queue.get(), 1)
        self.assertEqual(queue.get(True, None), 2)
        self.assertEqual(queue.get(True), 3)
        self.assertEqual(queue.get(timeout=1), 4)
        self.assertEqual(queue.get_nowait(), 5)

        self.assertEqual(queue_empty(queue), True)

        get = TimingWrapper(queue.get)
        get_nowait = TimingWrapper(queue.get_nowait)

        self.assertRaises(pyqueue.Empty, get, False)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, False, None)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get_nowait)
        self.assertTimingAlmostEqual(get_nowait.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, True, TIMEOUT1)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT1)

        self.assertRaises(pyqueue.Empty, get, False, TIMEOUT2)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, timeout=TIMEOUT3)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT3)

        proc.join()
        close_queue(queue)

    @classmethod
    def _test_fork(cls, queue):
        for i in range(10, 20):
            queue.put(i)
        # note that at this point the items may only be buffered, so the
        # process cannot shutdown until the feeder thread has finished
        # pushing items onto the pipe.

    def test_fork(self):
        # Old versions of Queue would fail to create a new feeder
        # thread for a forked process if the original process had its
        # own feeder thread.  This test checks that this no longer
        # happens.

        queue = self.Queue()

        # put items on queue so that main process starts a feeder thread
        for i in range(10):
            queue.put(i)

        # wait to make sure thread starts before we fork a new process
        time.sleep(DELTA)

        # fork process
        p = self.Process(target=self._test_fork, args=(queue,))
        p.daemon = True
        p.start()

        # check that all expected items are in the queue
        for i in range(20):
            self.assertEqual(queue.get(), i)
        self.assertRaises(pyqueue.Empty, queue.get, False)

        p.join()
        close_queue(queue)

    def test_qsize(self):
        q = self.Queue()
        try:
            self.assertEqual(q.qsize(), 0)
        except NotImplementedError:
            self.skipTest("qsize method not implemented")
        q.put(1)
        self.assertEqual(q.qsize(), 1)
        q.put(5)
        self.assertEqual(q.qsize(), 2)
        q.get()
        self.assertEqual(q.qsize(), 1)
        q.get()
        self.assertEqual(q.qsize(), 0)
        close_queue(q)

    @classmethod
    def _test_task_done(cls, q):
        for obj in iter(q.get, None):
            time.sleep(DELTA)
            q.task_done()

    def test_task_done(self):
        queue = self.JoinableQueue()

        workers = [self.Process(target=self._test_task_done, args=(queue,)) for i in range(4)]

        for p in workers:
            p.daemon = True
            p.start()

        for i in range(10):
            queue.put(i)

        queue.join()

        for p in workers:
            queue.put(None)

        for p in workers:
            p.join()
        close_queue(queue)

    def test_no_import_lock_contention(self):
        with test.support.temp_cwd():
            module_name = "imported_by_an_imported_module"
            with open(module_name + ".py", "w") as f:
                f.write(
                    """if 1:
                    import multiprocessing

                    q = multiprocessing.Queue()
                    q.put('knock knock')
                    q.get(timeout=3)
                    q.close()
                    del q
                """
                )

            with test.support.DirsOnSysPath(os.getcwd()):
                try:
                    __import__(module_name)
                except pyqueue.Empty:
                    self.fail("Probable regression on import lock contention;" " see Issue #22853")

    def test_timeout(self):
        q = multiprocessing.Queue()
        start = time.monotonic()
        self.assertRaises(pyqueue.Empty, q.get, True, 0.200)
        delta = time.monotonic() - start
        # bpo-30317: Tolerate a delta of 100 ms because of the bad clock
        # resolution on Windows (usually 15.6 ms). x86 Windows7 3.x once
        # failed because the delta was only 135.8 ms.
        self.assertGreaterEqual(delta, 0.100)
        close_queue(q)

    def test_queue_feeder_donot_stop_onexc(self):
        # bpo-30414: verify feeder handles exceptions correctly
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class NotSerializable(object):
            def __reduce__(self):
                raise AttributeError

        with test.support.captured_stderr():
            q = self.Queue()
            q.put(NotSerializable())
            q.put(True)
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))
            close_queue(q)

        with test.support.captured_stderr():
            # bpo-33078: verify that the queue size is correctly handled
            # on errors.
            q = self.Queue(maxsize=1)
            q.put(NotSerializable())
            q.put(True)
            try:
                self.assertEqual(q.qsize(), 1)
            except NotImplementedError:
                # qsize is not available on all platform as it
                # relies on sem_getvalue
                pass
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))
            # Check that the size of the queue is correct
            self.assertTrue(q.empty())
            close_queue(q)

    def test_queue_feeder_on_queue_feeder_error(self):
        # bpo-30006: verify feeder handles exceptions using the
        # _on_queue_feeder_error hook.
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class NotSerializable(object):
            """Mock unserializable object"""

            def __init__(self):
                self.reduce_was_called = False
                self.on_queue_feeder_error_was_called = False

            def __reduce__(self):
                self.reduce_was_called = True
                raise AttributeError

        class SafeQueue(multiprocessing.queues.Queue):
            """Queue with overloaded _on_queue_feeder_error hook"""

            @staticmethod
            def _on_queue_feeder_error(e, obj):
                if isinstance(e, AttributeError) and isinstance(obj, NotSerializable):
                    obj.on_queue_feeder_error_was_called = True

        not_serializable_obj = NotSerializable()
        # The captured_stderr reduces the noise in the test report
        with test.support.captured_stderr():
            q = SafeQueue(ctx=multiprocessing.get_context())
            q.put(not_serializable_obj)

            # Verify that q is still functioning correctly
            q.put(True)
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))

        # Assert that the serialization and the hook have been called correctly
        self.assertTrue(not_serializable_obj.reduce_was_called)
        self.assertTrue(not_serializable_obj.on_queue_feeder_error_was_called)

    def test_closed_queue_put_get_exceptions(self):
        for q in multiprocessing.Queue(), multiprocessing.JoinableQueue():
            q.close()
            with self.assertRaisesRegex(ValueError, "is closed"):
                q.put("foo")
            with self.assertRaisesRegex(ValueError, "is closed"):
                q.get()


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestQueue(BaseTestCase, ThreadsMixin, unittest.TestCase):
    @classmethod
    def _test_put(cls, queue, child_can_start, parent_can_continue):
        child_can_start.wait()
        for i in range(6):
            queue.get()
        parent_can_continue.set()

    def test_put(self):
        MAXSIZE = 6
        queue = self.Queue(maxsize=MAXSIZE)
        child_can_start = self.Event()
        parent_can_continue = self.Event()

        proc = self.Process(target=self._test_put, args=(queue, child_can_start, parent_can_continue))
        proc.daemon = True
        proc.start()

        self.assertEqual(queue_empty(queue), True)
        self.assertEqual(queue_full(queue, MAXSIZE), False)

        queue.put(1)
        queue.put(2, True)
        queue.put(3, True, None)
        queue.put(4, False)
        queue.put(5, False, None)
        queue.put_nowait(6)

        # the values may be in buffer but not yet in pipe so sleep a bit
        time.sleep(DELTA)

        self.assertEqual(queue_empty(queue), False)
        self.assertEqual(queue_full(queue, MAXSIZE), True)

        put = TimingWrapper(queue.put)
        put_nowait = TimingWrapper(queue.put_nowait)

        self.assertRaises(pyqueue.Full, put, 7, False)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, False, None)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put_nowait, 7)
        self.assertTimingAlmostEqual(put_nowait.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, True, TIMEOUT1)
        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT1)

        self.assertRaises(pyqueue.Full, put, 7, False, TIMEOUT2)
        self.assertTimingAlmostEqual(put.elapsed, 0)

        self.assertRaises(pyqueue.Full, put, 7, True, timeout=TIMEOUT3)
        self.assertTimingAlmostEqual(put.elapsed, TIMEOUT3)

        child_can_start.set()
        parent_can_continue.wait()

        self.assertEqual(queue_empty(queue), True)
        self.assertEqual(queue_full(queue, MAXSIZE), False)

        proc.join()
        close_queue(queue)

    @classmethod
    def _test_get(cls, queue, child_can_start, parent_can_continue):
        child_can_start.wait()
        # queue.put(1)
        queue.put(2)
        queue.put(3)
        queue.put(4)
        queue.put(5)
        parent_can_continue.set()

    def test_get(self):
        queue = self.Queue()
        child_can_start = self.Event()
        parent_can_continue = self.Event()

        proc = self.Process(target=self._test_get, args=(queue, child_can_start, parent_can_continue))
        proc.daemon = True
        proc.start()

        self.assertEqual(queue_empty(queue), True)

        child_can_start.set()
        parent_can_continue.wait()

        time.sleep(DELTA)
        self.assertEqual(queue_empty(queue), False)

        # Hangs unexpectedly, remove for now
        # self.assertEqual(queue.get(), 1)
        self.assertEqual(queue.get(True, None), 2)
        self.assertEqual(queue.get(True), 3)
        self.assertEqual(queue.get(timeout=1), 4)
        self.assertEqual(queue.get_nowait(), 5)

        self.assertEqual(queue_empty(queue), True)

        get = TimingWrapper(queue.get)
        get_nowait = TimingWrapper(queue.get_nowait)

        self.assertRaises(pyqueue.Empty, get, False)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, False, None)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get_nowait)
        self.assertTimingAlmostEqual(get_nowait.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, True, TIMEOUT1)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT1)

        self.assertRaises(pyqueue.Empty, get, False, TIMEOUT2)
        self.assertTimingAlmostEqual(get.elapsed, 0)

        self.assertRaises(pyqueue.Empty, get, timeout=TIMEOUT3)
        self.assertTimingAlmostEqual(get.elapsed, TIMEOUT3)

        proc.join()
        close_queue(queue)

    @classmethod
    def _test_fork(cls, queue):
        for i in range(10, 20):
            queue.put(i)
        # note that at this point the items may only be buffered, so the
        # process cannot shutdown until the feeder thread has finished
        # pushing items onto the pipe.

    def test_fork(self):
        # Old versions of Queue would fail to create a new feeder
        # thread for a forked process if the original process had its
        # own feeder thread.  This test checks that this no longer
        # happens.

        queue = self.Queue()

        # put items on queue so that main process starts a feeder thread
        for i in range(10):
            queue.put(i)

        # wait to make sure thread starts before we fork a new process
        time.sleep(DELTA)

        # fork process
        p = self.Process(target=self._test_fork, args=(queue,))
        p.daemon = True
        p.start()

        # check that all expected items are in the queue
        for i in range(20):
            self.assertEqual(queue.get(), i)
        self.assertRaises(pyqueue.Empty, queue.get, False)

        p.join()
        close_queue(queue)

    def test_qsize(self):
        q = self.Queue()
        try:
            self.assertEqual(q.qsize(), 0)
        except NotImplementedError:
            self.skipTest("qsize method not implemented")
        q.put(1)
        self.assertEqual(q.qsize(), 1)
        q.put(5)
        self.assertEqual(q.qsize(), 2)
        q.get()
        self.assertEqual(q.qsize(), 1)
        q.get()
        self.assertEqual(q.qsize(), 0)
        close_queue(q)

    @classmethod
    def _test_task_done(cls, q):
        for obj in iter(q.get, None):
            time.sleep(DELTA)
            q.task_done()

    def test_task_done(self):
        queue = self.JoinableQueue()

        workers = [self.Process(target=self._test_task_done, args=(queue,)) for i in range(4)]

        for p in workers:
            p.daemon = True
            p.start()

        for i in range(10):
            queue.put(i)

        queue.join()

        for p in workers:
            queue.put(None)

        for p in workers:
            p.join()
        close_queue(queue)

    def test_no_import_lock_contention(self):
        with test.support.temp_cwd():
            module_name = "imported_by_an_imported_module"
            with open(module_name + ".py", "w") as f:
                f.write(
                    """if 1:
                    import multiprocessing

                    q = multiprocessing.Queue()
                    q.put('knock knock')
                    q.get(timeout=3)
                    q.close()
                    del q
                """
                )

            with test.support.DirsOnSysPath(os.getcwd()):
                try:
                    __import__(module_name)
                except pyqueue.Empty:
                    self.fail("Probable regression on import lock contention;" " see Issue #22853")

    def test_timeout(self):
        q = multiprocessing.Queue()
        start = time.monotonic()
        self.assertRaises(pyqueue.Empty, q.get, True, 0.200)
        delta = time.monotonic() - start
        # bpo-30317: Tolerate a delta of 100 ms because of the bad clock
        # resolution on Windows (usually 15.6 ms). x86 Windows7 3.x once
        # failed because the delta was only 135.8 ms.
        self.assertGreaterEqual(delta, 0.100)
        close_queue(q)

    def test_queue_feeder_donot_stop_onexc(self):
        # bpo-30414: verify feeder handles exceptions correctly
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class NotSerializable(object):
            def __reduce__(self):
                raise AttributeError

        with test.support.captured_stderr():
            q = self.Queue()
            q.put(NotSerializable())
            q.put(True)
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))
            close_queue(q)

        with test.support.captured_stderr():
            # bpo-33078: verify that the queue size is correctly handled
            # on errors.
            q = self.Queue(maxsize=1)
            q.put(NotSerializable())
            q.put(True)
            try:
                self.assertEqual(q.qsize(), 1)
            except NotImplementedError:
                # qsize is not available on all platform as it
                # relies on sem_getvalue
                pass
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))
            # Check that the size of the queue is correct
            self.assertTrue(q.empty())
            close_queue(q)

    def test_queue_feeder_on_queue_feeder_error(self):
        # bpo-30006: verify feeder handles exceptions using the
        # _on_queue_feeder_error hook.
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        class NotSerializable(object):
            """Mock unserializable object"""

            def __init__(self):
                self.reduce_was_called = False
                self.on_queue_feeder_error_was_called = False

            def __reduce__(self):
                self.reduce_was_called = True
                raise AttributeError

        class SafeQueue(multiprocessing.queues.Queue):
            """Queue with overloaded _on_queue_feeder_error hook"""

            @staticmethod
            def _on_queue_feeder_error(e, obj):
                if isinstance(e, AttributeError) and isinstance(obj, NotSerializable):
                    obj.on_queue_feeder_error_was_called = True

        not_serializable_obj = NotSerializable()
        # The captured_stderr reduces the noise in the test report
        with test.support.captured_stderr():
            q = SafeQueue(ctx=multiprocessing.get_context())
            q.put(not_serializable_obj)

            # Verify that q is still functioning correctly
            q.put(True)
            self.assertTrue(q.get(timeout=support.SHORT_TIMEOUT))

        # Assert that the serialization and the hook have been called correctly
        self.assertTrue(not_serializable_obj.reduce_was_called)
        self.assertTrue(not_serializable_obj.on_queue_feeder_error_was_called)

    def test_closed_queue_put_get_exceptions(self):
        for q in multiprocessing.Queue(), multiprocessing.JoinableQueue():
            q.close()
            with self.assertRaisesRegex(ValueError, "is closed"):
                q.put("foo")
            with self.assertRaisesRegex(ValueError, "is closed"):
                q.get()


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_queue.py
DragonHPC/dragon
""" Test mulitprocessing lock, event, condition, semaphore, barrier
"""
import unittest

from test.support import hashlib_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing

from common import (
    BaseTestCase,
    ProcessesMixin,
    ManagerMixin,
    ThreadsMixin,
    TimingWrapper,
    setUpModule,
    tearDownModule,
    get_value,
    TIMEOUT1,
    TIMEOUT2,
    TIMEOUT3,
)


class WithProcessesTestSemaphore(BaseTestCase, ProcessesMixin, unittest.TestCase):
    def _test_semaphore(self, sem):
        self.assertReturnsIfImplemented(2, get_value, sem)
        self.assertEqual(sem.acquire(), True)
        self.assertReturnsIfImplemented(1, get_value, sem)
        self.assertEqual(sem.acquire(), True)
        self.assertReturnsIfImplemented(0, get_value, sem)
        self.assertEqual(sem.acquire(False), False)
        self.assertReturnsIfImplemented(0, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(1, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(2, get_value, sem)

    def test_semaphore(self):
        sem = self.Semaphore(2)
        self._test_semaphore(sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(3, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(4, get_value, sem)

    def test_bounded_semaphore(self):
        sem = self.BoundedSemaphore(2)
        self._test_semaphore(sem)
        # Currently fails on OS/X
        # if HAVE_GETVALUE:
        #    self.assertRaises(ValueError, sem.release)
        #    self.assertReturnsIfImplemented(2, get_value, sem)

    def test_timeout(self):
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        sem = self.Semaphore(0)
        acquire = TimingWrapper(sem.acquire)

        self.assertEqual(acquire(False), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0.0)

        self.assertEqual(acquire(False, None), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0.0)

        self.assertEqual(acquire(False, TIMEOUT1), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0)

        self.assertEqual(acquire(True, TIMEOUT2), False)
        self.assertTimingAlmostEqual(acquire.elapsed, TIMEOUT2)

        self.assertEqual(acquire(timeout=TIMEOUT3), False)
        self.assertTimingAlmostEqual(acquire.elapsed, TIMEOUT3)


@unittest.skip("DRAGON: Manager not implemented")
@hashlib_helper.requires_hashdigest("md5")
class WithManagerTestSemaphore(BaseTestCase, ManagerMixin, unittest.TestCase):
    def _test_semaphore(self, sem):
        self.assertReturnsIfImplemented(2, get_value, sem)
        self.assertEqual(sem.acquire(), True)
        self.assertReturnsIfImplemented(1, get_value, sem)
        self.assertEqual(sem.acquire(), True)
        self.assertReturnsIfImplemented(0, get_value, sem)
        self.assertEqual(sem.acquire(False), False)
        self.assertReturnsIfImplemented(0, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(1, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(2, get_value, sem)

    def test_semaphore(self):
        sem = self.Semaphore(2)
        self._test_semaphore(sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(3, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(4, get_value, sem)

    def test_bounded_semaphore(self):
        sem = self.BoundedSemaphore(2)
        self._test_semaphore(sem)
        # Currently fails on OS/X
        # if HAVE_GETVALUE:
        #    self.assertRaises(ValueError, sem.release)
        #    self.assertReturnsIfImplemented(2, get_value, sem)

    def test_timeout(self):
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        sem = self.Semaphore(0)
        acquire = TimingWrapper(sem.acquire)

        self.assertEqual(acquire(False), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0.0)

        self.assertEqual(acquire(False, None), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0.0)

        self.assertEqual(acquire(False, TIMEOUT1), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0)

        self.assertEqual(acquire(True, TIMEOUT2), False)
        self.assertTimingAlmostEqual(acquire.elapsed, TIMEOUT2)

        self.assertEqual(acquire(timeout=TIMEOUT3), False)
        self.assertTimingAlmostEqual(acquire.elapsed, TIMEOUT3)


@unittest.skip("DRAGON: Threads not implemented")
class WithThreadsTestSemaphore(BaseTestCase, ThreadsMixin, unittest.TestCase):
    def _test_semaphore(self, sem):
        self.assertReturnsIfImplemented(2, get_value, sem)
        self.assertEqual(sem.acquire(), True)
        self.assertReturnsIfImplemented(1, get_value, sem)
        self.assertEqual(sem.acquire(), True)
        self.assertReturnsIfImplemented(0, get_value, sem)
        self.assertEqual(sem.acquire(False), False)
        self.assertReturnsIfImplemented(0, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(1, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(2, get_value, sem)

    def test_semaphore(self):
        sem = self.Semaphore(2)
        self._test_semaphore(sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(3, get_value, sem)
        self.assertEqual(sem.release(), None)
        self.assertReturnsIfImplemented(4, get_value, sem)

    def test_bounded_semaphore(self):
        sem = self.BoundedSemaphore(2)
        self._test_semaphore(sem)
        # Currently fails on OS/X
        # if HAVE_GETVALUE:
        #    self.assertRaises(ValueError, sem.release)
        #    self.assertReturnsIfImplemented(2, get_value, sem)

    def test_timeout(self):
        if self.TYPE != "processes":
            self.skipTest("test not appropriate for {}".format(self.TYPE))

        sem = self.Semaphore(0)
        acquire = TimingWrapper(sem.acquire)

        self.assertEqual(acquire(False), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0.0)

        self.assertEqual(acquire(False, None), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0.0)

        self.assertEqual(acquire(False, TIMEOUT1), False)
        self.assertTimingAlmostEqual(acquire.elapsed, 0)

        self.assertEqual(acquire(True, TIMEOUT2), False)
        self.assertTimingAlmostEqual(acquire.elapsed, TIMEOUT2)

        self.assertEqual(acquire(timeout=TIMEOUT3), False)
        self.assertTimingAlmostEqual(acquire.elapsed, TIMEOUT3)


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_semaphore.py
DragonHPC/dragon
"""Test multiprocessing Shared_Memory, Array, Value and Shared_ctypes
"""
import unittest

try:
    from ctypes import Structure, c_int, c_double, c_longlong
except ImportError:
    Structure = object
    c_int = c_double = c_longlong = None

import dragon  # DRAGON import before multiprocessing

import multiprocessing

try:
    from multiprocessing.sharedctypes import Value, copy

    HAS_SHAREDCTYPES = True
except ImportError:
    HAS_SHAREDCTYPES = False

from common import (
    BaseTestCase,
    ProcessesMixin,
    latin,
    setUpModule,
    tearDownModule,
)


class _Foo(Structure):
    _fields_ = [
        ("x", c_int),
        ("y", c_double),
        (
            "z",
            c_longlong,
        ),
    ]

@unittest.skip("DRAGON: Not Implemented")
class WithProcessesTestSharedCTypes(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    def setUp(self):
        if not HAS_SHAREDCTYPES:
            self.skipTest("requires multiprocessing.sharedctypes")

    @classmethod
    def _double(cls, x, y, z, foo, arr, string):
        x.value *= 2
        y.value *= 2
        z.value *= 2
        foo.x *= 2
        foo.y *= 2
        string.value *= 2
        for i in range(len(arr)):
            arr[i] *= 2

    @unittest.skip("bug filed PE-40919")
    def test_sharedctypes(self, lock=False):
        x = Value("i", 7, lock=lock)
        y = Value(c_double, 1.0 / 3.0, lock=lock)
        z = Value(c_longlong, 2**33, lock=lock)
        foo = Value(_Foo, 3, 2, lock=lock)
        arr = self.Array("d", list(range(10)), lock=lock)
        string = self.Array("c", 20, lock=lock)
        string.value = latin("hello")

        p = self.Process(target=self._double, args=(x, y, z, foo, arr, string))
        p.daemon = True
        p.start()
        p.join()

        self.assertEqual(x.value, 14)
        self.assertAlmostEqual(y.value, 2.0 / 3.0)
        self.assertEqual(z.value, 2**34)
        self.assertEqual(foo.x, 6)
        self.assertAlmostEqual(foo.y, 4.0)
        for i in range(10):
            self.assertAlmostEqual(arr[i], i * 2)
        self.assertEqual(string.value, latin("hellohello"))

    def test_synchronize(self):
        self.test_sharedctypes(lock=True)

    def test_copy(self):
        foo = _Foo(2, 5.0, 2**33)
        bar = copy(foo)
        foo.x = 0
        foo.y = 0
        foo.z = 0
        self.assertEqual(bar.x, 2)
        self.assertAlmostEqual(bar.y, 5.0)
        self.assertEqual(bar.z, 2**33)


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_shared_ctypes.py
DragonHPC/dragon
"""Test multiprocessing Shared_Memory and Heap
"""
import unittest
import unittest.mock
import os
import pickle
import signal
import subprocess
import sys
import time
import warnings
import random
import gc

import test.support
from test import support
import test.support.script_helper
from test.support import hashlib_helper

import dragon  # DRAGON import before multiprocessing

import multiprocessing

try:
    from multiprocessing import shared_memory

    HAS_SHMEM = True
except ImportError:
    HAS_SHMEM = False

from common import (
    BaseTestCase,
    ProcessesMixin,
    setUpModule,
    tearDownModule,
)


@unittest.skip("DRAGON: not implemented")
@unittest.skipUnless(HAS_SHMEM, "requires multiprocessing.shared_memory")
@hashlib_helper.requires_hashdigest("md5")
class WithProcessesTestSharedMemory(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    @staticmethod
    def _attach_existing_shmem_then_write(shmem_name_or_obj, binary_data):
        if isinstance(shmem_name_or_obj, str):
            local_sms = shared_memory.SharedMemory(shmem_name_or_obj)
        else:
            local_sms = shmem_name_or_obj
        local_sms.buf[: len(binary_data)] = binary_data
        local_sms.close()

    def test_shared_memory_basics(self):
        sms = shared_memory.SharedMemory("test01_tsmb", create=True, size=512)
        self.addCleanup(sms.unlink)

        # Verify attributes are readable.
        self.assertEqual(sms.name, "test01_tsmb")
        self.assertGreaterEqual(sms.size, 512)
        self.assertGreaterEqual(len(sms.buf), sms.size)

        # Modify contents of shared memory segment through memoryview.
        sms.buf[0] = 42
        self.assertEqual(sms.buf[0], 42)

        # Attach to existing shared memory segment.
        also_sms = shared_memory.SharedMemory("test01_tsmb")
        self.assertEqual(also_sms.buf[0], 42)
        also_sms.close()

        # Attach to existing shared memory segment but specify a new size.
        same_sms = shared_memory.SharedMemory("test01_tsmb", size=20 * sms.size)
        self.assertLess(same_sms.size, 20 * sms.size)  # Size was ignored.
        same_sms.close()

        # Creating Shared Memory Segment with -ve size
        with self.assertRaises(ValueError):
            shared_memory.SharedMemory(create=True, size=-2)

        # Attaching Shared Memory Segment without a name
        with self.assertRaises(ValueError):
            shared_memory.SharedMemory(create=False)

        # Test if shared memory segment is created properly,
        # when _make_filename returns an existing shared memory segment name
        with unittest.mock.patch("multiprocessing.shared_memory._make_filename") as mock_make_filename:

            NAME_PREFIX = shared_memory._SHM_NAME_PREFIX
            names = ["test01_fn", "test02_fn"]
            # Prepend NAME_PREFIX which can be '/psm_' or 'wnsm_', necessary
            # because some POSIX compliant systems require name to start with /
            names = [NAME_PREFIX + name for name in names]

            mock_make_filename.side_effect = names
            shm1 = shared_memory.SharedMemory(create=True, size=1)
            self.addCleanup(shm1.unlink)
            self.assertEqual(shm1._name, names[0])

            mock_make_filename.side_effect = names
            shm2 = shared_memory.SharedMemory(create=True, size=1)
            self.addCleanup(shm2.unlink)
            self.assertEqual(shm2._name, names[1])

        if shared_memory._USE_POSIX:
            # Posix Shared Memory can only be unlinked once.  Here we
            # test an implementation detail that is not observed across
            # all supported platforms (since WindowsNamedSharedMemory
            # manages unlinking on its own and unlink() does nothing).
            # True release of shared memory segment does not necessarily
            # happen until process exits, depending on the OS platform.
            with self.assertRaises(FileNotFoundError):
                sms_uno = shared_memory.SharedMemory("test01_dblunlink", create=True, size=5000)

                try:
                    self.assertGreaterEqual(sms_uno.size, 5000)

                    sms_duo = shared_memory.SharedMemory("test01_dblunlink")
                    sms_duo.unlink()  # First shm_unlink() call.
                    sms_duo.close()
                    sms_uno.close()

                finally:
                    sms_uno.unlink()  # A second shm_unlink() call is bad.

        with self.assertRaises(FileExistsError):
            # Attempting to create a new shared memory segment with a
            # name that is already in use triggers an exception.
            there_can_only_be_one_sms = shared_memory.SharedMemory("test01_tsmb", create=True, size=512)

        if shared_memory._USE_POSIX:
            # Requesting creation of a shared memory segment with the option
            # to attach to an existing segment, if that name is currently in
            # use, should not trigger an exception.
            # Note:  Using a smaller size could possibly cause truncation of
            # the existing segment but is OS platform dependent.  In the
            # case of MacOS/darwin, requesting a smaller size is disallowed.
            class OptionalAttachSharedMemory(shared_memory.SharedMemory):
                _flags = os.O_CREAT | os.O_RDWR

            ok_if_exists_sms = OptionalAttachSharedMemory("test01_tsmb")
            self.assertEqual(ok_if_exists_sms.size, sms.size)
            ok_if_exists_sms.close()

        # Attempting to attach to an existing shared memory segment when
        # no segment exists with the supplied name triggers an exception.
        with self.assertRaises(FileNotFoundError):
            nonexisting_sms = shared_memory.SharedMemory("test01_notthere")
            nonexisting_sms.unlink()  # Error should occur on prior line.

        sms.close()

        # Test creating a shared memory segment with negative size
        with self.assertRaises(ValueError):
            sms_invalid = shared_memory.SharedMemory(create=True, size=-1)

        # Test creating a shared memory segment with size 0
        with self.assertRaises(ValueError):
            sms_invalid = shared_memory.SharedMemory(create=True, size=0)

        # Test creating a shared memory segment without size argument
        with self.assertRaises(ValueError):
            sms_invalid = shared_memory.SharedMemory(create=True)

    def test_shared_memory_across_processes(self):
        # bpo-40135: don't define shared memory block's name in case of
        # the failure when we run multiprocessing tests in parallel.
        sms = shared_memory.SharedMemory(create=True, size=512)
        self.addCleanup(sms.unlink)

        # Verify remote attachment to existing block by name is working.
        p = self.Process(target=self._attach_existing_shmem_then_write, args=(sms.name, b"howdy"))
        p.daemon = True
        p.start()
        p.join()
        self.assertEqual(bytes(sms.buf[:5]), b"howdy")

        # Verify pickling of SharedMemory instance also works.
        p = self.Process(target=self._attach_existing_shmem_then_write, args=(sms, b"HELLO"))
        p.daemon = True
        p.start()
        p.join()
        self.assertEqual(bytes(sms.buf[:5]), b"HELLO")

        sms.close()

    @unittest.skipIf(os.name != "posix", "not feasible in non-posix platforms")
    def test_shared_memory_SharedMemoryServer_ignores_sigint(self):
        # bpo-36368: protect SharedMemoryManager server process from
        # KeyboardInterrupt signals.
        smm = multiprocessing.managers.SharedMemoryManager()
        smm.start()

        # make sure the manager works properly at the beginning
        sl = smm.ShareableList(range(10))

        # the manager's server should ignore KeyboardInterrupt signals, and
        # maintain its connection with the current process, and success when
        # asked to deliver memory segments.
        os.kill(smm._process.pid, signal.SIGINT)

        sl2 = smm.ShareableList(range(10))

        # test that the custom signal handler registered in the Manager does
        # not affect signal handling in the parent process.
        with self.assertRaises(KeyboardInterrupt):
            os.kill(os.getpid(), signal.SIGINT)

        smm.shutdown()

    @unittest.skipIf(os.name != "posix", "resource_tracker is posix only")
    def test_shared_memory_SharedMemoryManager_reuses_resource_tracker(self):
        # bpo-36867: test that a SharedMemoryManager uses the
        # same resource_tracker process as its parent.
        cmd = """if 1:
            from multiprocessing.managers import SharedMemoryManager


            smm = SharedMemoryManager()
            smm.start()
            sl = smm.ShareableList(range(10))
            smm.shutdown()
        """
        rc, out, err = test.support.script_helper.assert_python_ok("-c", cmd)

        # Before bpo-36867 was fixed, a SharedMemoryManager not using the same
        # resource_tracker process as its parent would make the parent's
        # tracker complain about sl being leaked even though smm.shutdown()
        # properly released sl.
        self.assertFalse(err)

    def test_shared_memory_SharedMemoryManager_basics(self):
        smm1 = multiprocessing.managers.SharedMemoryManager()
        with self.assertRaises(ValueError):
            smm1.SharedMemory(size=9)  # Fails if SharedMemoryServer not started
        smm1.start()
        lol = [smm1.ShareableList(range(i)) for i in range(5, 10)]
        lom = [smm1.SharedMemory(size=j) for j in range(32, 128, 16)]
        doppleganger_list0 = shared_memory.ShareableList(name=lol[0].shm.name)
        self.assertEqual(len(doppleganger_list0), 5)
        doppleganger_shm0 = shared_memory.SharedMemory(name=lom[0].name)
        self.assertGreaterEqual(len(doppleganger_shm0.buf), 32)
        held_name = lom[0].name
        smm1.shutdown()
        if sys.platform != "win32":
            # Calls to unlink() have no effect on Windows platform; shared
            # memory will only be released once final process exits.
            with self.assertRaises(FileNotFoundError):
                # No longer there to be attached to again.
                absent_shm = shared_memory.SharedMemory(name=held_name)

        with multiprocessing.managers.SharedMemoryManager() as smm2:
            sl = smm2.ShareableList("howdy")
            shm = smm2.SharedMemory(size=128)
            held_name = sl.shm.name
        if sys.platform != "win32":
            with self.assertRaises(FileNotFoundError):
                # No longer there to be attached to again.
                absent_sl = shared_memory.ShareableList(name=held_name)

    def test_shared_memory_ShareableList_basics(self):
        sl = shared_memory.ShareableList(["howdy", b"HoWdY", -273.154, 100, None, True, 42])
        self.addCleanup(sl.shm.unlink)

        # Verify attributes are readable.
        self.assertEqual(sl.format, "8s8sdqxxxxxx?xxxxxxxx?q")

        # Exercise len().
        self.assertEqual(len(sl), 7)

        # Exercise index().
        with warnings.catch_warnings():
            # Suppress BytesWarning when comparing against b'HoWdY'.
            warnings.simplefilter("ignore")
            with self.assertRaises(ValueError):
                sl.index("100")
            self.assertEqual(sl.index(100), 3)

        # Exercise retrieving individual values.
        self.assertEqual(sl[0], "howdy")
        self.assertEqual(sl[-2], True)

        # Exercise iterability.
        self.assertEqual(tuple(sl), ("howdy", b"HoWdY", -273.154, 100, None, True, 42))

        # Exercise modifying individual values.
        sl[3] = 42
        self.assertEqual(sl[3], 42)
        sl[4] = "some"  # Change type at a given position.
        self.assertEqual(sl[4], "some")
        self.assertEqual(sl.format, "8s8sdq8sxxxxxxx?q")
        with self.assertRaisesRegex(ValueError, "exceeds available storage"):
            sl[4] = "far too many"
        self.assertEqual(sl[4], "some")
        sl[0] = "encodés"  # Exactly 8 bytes of UTF-8 data
        self.assertEqual(sl[0], "encodés")
        self.assertEqual(sl[1], b"HoWdY")  # no spillage
        with self.assertRaisesRegex(ValueError, "exceeds available storage"):
            sl[0] = "encodées"  # Exactly 9 bytes of UTF-8 data
        self.assertEqual(sl[1], b"HoWdY")
        with self.assertRaisesRegex(ValueError, "exceeds available storage"):
            sl[1] = b"123456789"
        self.assertEqual(sl[1], b"HoWdY")

        # Exercise count().
        with warnings.catch_warnings():
            # Suppress BytesWarning when comparing against b'HoWdY'.
            warnings.simplefilter("ignore")
            self.assertEqual(sl.count(42), 2)
            self.assertEqual(sl.count(b"HoWdY"), 1)
            self.assertEqual(sl.count(b"adios"), 0)

        # Exercise creating a duplicate.
        sl_copy = shared_memory.ShareableList(sl, name="test03_duplicate")
        try:
            self.assertNotEqual(sl.shm.name, sl_copy.shm.name)
            self.assertEqual("test03_duplicate", sl_copy.shm.name)
            self.assertEqual(list(sl), list(sl_copy))
            self.assertEqual(sl.format, sl_copy.format)
            sl_copy[-1] = 77
            self.assertEqual(sl_copy[-1], 77)
            self.assertNotEqual(sl[-1], 77)
            sl_copy.shm.close()
        finally:
            sl_copy.shm.unlink()

        # Obtain a second handle on the same ShareableList.
        sl_tethered = shared_memory.ShareableList(name=sl.shm.name)
        self.assertEqual(sl.shm.name, sl_tethered.shm.name)
        sl_tethered[-1] = 880
        self.assertEqual(sl[-1], 880)
        sl_tethered.shm.close()

        sl.shm.close()

        # Exercise creating an empty ShareableList.
        empty_sl = shared_memory.ShareableList()
        try:
            self.assertEqual(len(empty_sl), 0)
            self.assertEqual(empty_sl.format, "")
            self.assertEqual(empty_sl.count("any"), 0)
            with self.assertRaises(ValueError):
                empty_sl.index(None)
            empty_sl.shm.close()
        finally:
            empty_sl.shm.unlink()

    def test_shared_memory_ShareableList_pickling(self):
        sl = shared_memory.ShareableList(range(10))
        self.addCleanup(sl.shm.unlink)

        serialized_sl = pickle.dumps(sl)
        deserialized_sl = pickle.loads(serialized_sl)
        self.assertTrue(isinstance(deserialized_sl, shared_memory.ShareableList))
        self.assertTrue(deserialized_sl[-1], 9)
        self.assertFalse(sl is deserialized_sl)
        deserialized_sl[4] = "changed"
        self.assertEqual(sl[4], "changed")

        # Verify data is not being put into the pickled representation.
        name = "a" * len(sl.shm.name)
        larger_sl = shared_memory.ShareableList(range(400))
        self.addCleanup(larger_sl.shm.unlink)
        serialized_larger_sl = pickle.dumps(larger_sl)
        self.assertTrue(len(serialized_sl) == len(serialized_larger_sl))
        larger_sl.shm.close()

        deserialized_sl.shm.close()
        sl.shm.close()

    def test_shared_memory_cleaned_after_process_termination(self):
        cmd = """if 1:
            import os, time, sys
            from multiprocessing import shared_memory

            # Create a shared_memory segment, and send the segment name
            sm = shared_memory.SharedMemory(create=True, size=10)
            sys.stdout.write(sm.name + '\\n')
            sys.stdout.flush()
            time.sleep(100)
        """
        with subprocess.Popen(
            [sys.executable, "-E", "-c", cmd], stdout=subprocess.PIPE, stderr=subprocess.PIPE
        ) as p:
            name = p.stdout.readline().strip().decode()

            # killing abruptly processes holding reference to a shared memory
            # segment should not leak the given memory segment.
            p.terminate()
            p.wait()

            deadline = time.monotonic() + support.LONG_TIMEOUT
            t = 0.1
            while time.monotonic() < deadline:
                time.sleep(t)
                t = min(t * 2, 5)
                try:
                    smm = shared_memory.SharedMemory(name, create=False)
                except FileNotFoundError:
                    break
            else:
                raise AssertionError(
                    "A SharedMemory segment was leaked after" " a process was abruptly terminated."
                )

            if os.name == "posix":
                # A warning was emitted by the subprocess' own
                # resource_tracker (on Windows, shared memory segments
                # are released automatically by the OS).
                err = p.stderr.read().decode()
                self.assertIn(
                    "resource_tracker: There appear to be 1 leaked "
                    "shared_memory objects to clean up at shutdown",
                    err,
                )


@unittest.skip("bug filed PE-40902")
class WithProcessesTestHeap(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    def setUp(self):
        super().setUp()
        # Make pristine heap for these tests
        self.old_heap = multiprocessing.heap.BufferWrapper._heap
        multiprocessing.heap.BufferWrapper._heap = multiprocessing.heap.Heap()

    def tearDown(self):
        multiprocessing.heap.BufferWrapper._heap = self.old_heap
        super().tearDown()

    def test_heap(self):
        iterations = 5000
        maxblocks = 50
        blocks = []

        # get the heap object
        heap = multiprocessing.heap.BufferWrapper._heap
        heap._DISCARD_FREE_SPACE_LARGER_THAN = 0

        # create and destroy lots of blocks of different sizes
        for i in range(iterations):
            size = int(random.lognormvariate(0, 1) * 1000)
            b = multiprocessing.heap.BufferWrapper(size)
            blocks.append(b)
            if len(blocks) > maxblocks:
                i = random.randrange(maxblocks)
                del blocks[i]
            del b

        # verify the state of the heap
        with heap._lock:
            all = []
            free = 0
            occupied = 0
            for L in list(heap._len_to_seq.values()):
                # count all free blocks in arenas
                for arena, start, stop in L:
                    all.append((heap._arenas.index(arena), start, stop, stop - start, "free"))
                    free += stop - start
            for arena, arena_blocks in heap._allocated_blocks.items():
                # count all allocated blocks in arenas
                for start, stop in arena_blocks:
                    all.append((heap._arenas.index(arena), start, stop, stop - start, "occupied"))
                    occupied += stop - start

            self.assertEqual(free + occupied, sum(arena.size for arena in heap._arenas))

            all.sort()

            for i in range(len(all) - 1):
                (arena, start, stop) = all[i][:3]
                (narena, nstart, nstop) = all[i + 1][:3]
                if arena != narena:
                    # Two different arenas
                    self.assertEqual(stop, heap._arenas[arena].size)  # last block
                    self.assertEqual(nstart, 0)  # first block
                else:
                    # Same arena: two adjacent blocks
                    self.assertEqual(stop, nstart)

        # test free'ing all blocks
        random.shuffle(blocks)
        while blocks:
            blocks.pop()

        self.assertEqual(heap._n_frees, heap._n_mallocs)
        self.assertEqual(len(heap._pending_free_blocks), 0)
        self.assertEqual(len(heap._arenas), 0)
        self.assertEqual(len(heap._allocated_blocks), 0, heap._allocated_blocks)
        self.assertEqual(len(heap._len_to_seq), 0)

    def test_free_from_gc(self):
        # Check that freeing of blocks by the garbage collector doesn't deadlock
        # (issue #12352).
        # Make sure the GC is enabled, and set lower collection thresholds to
        # make collections more frequent (and increase the probability of
        # deadlock).
        if not gc.isenabled():
            gc.enable()
            self.addCleanup(gc.disable)
        thresholds = gc.get_threshold()
        self.addCleanup(gc.set_threshold, *thresholds)
        gc.set_threshold(10)

        # perform numerous block allocations, with cyclic references to make
        # sure objects are collected asynchronously by the gc
        for i in range(5000):
            a = multiprocessing.heap.BufferWrapper(1)
            b = multiprocessing.heap.BufferWrapper(1)
            # circular references
            a.buddy = b
            b.buddy = a


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()

examples/multiprocessing/unittests/test_shared_memory.py
DragonHPC/dragon
"""Test multiprocessing Shared_Memory, Array, Value and Shared_ctypes
"""
import unittest

import dragon  # DRAGON import before multiprocessing

import multiprocessing
from multiprocessing import util

try:
    from multiprocessing.sharedctypes import Value, copy

    HAS_SHAREDCTYPES = True
except ImportError:
    HAS_SHAREDCTYPES = False

from common import (
    BaseTestCase,
    ProcessesMixin,
    latin,
    setUpModule,
    tearDownModule,
)


class WithProcessesTestValue(BaseTestCase, ProcessesMixin, unittest.TestCase):

    # DRAGON ALLOWED_TYPES = ('processes',)

    codes_values = [
        ("i", 4343, 24234),
        ("d", 3.625, -4.25),
        ("h", -232, 234),
        ("q", 2**33, 2**34),
        ("c", latin("x"), latin("y")),
    ]

    def setUp(self):
        if not HAS_SHAREDCTYPES:
            self.skipTest("requires multiprocessing.sharedctypes")

    @classmethod
    def _test(cls, values):
        for sv, cv in zip(values, cls.codes_values):
            sv.value = cv[2]

    def test_value(self, raw=False):
        if raw:
            values = [self.RawValue(code, value) for code, value, _ in self.codes_values]
        else:
            values = [self.Value(code, value) for code, value, _ in self.codes_values]

        for sv, cv in zip(values, self.codes_values):
            self.assertEqual(sv.value, cv[1])

        proc = self.Process(target=self._test, args=(values,))
        proc.daemon = True
        proc.start()
        proc.join()

        for sv, cv in zip(values, self.codes_values):
            self.assertEqual(sv.value, cv[2])

    def test_rawvalue(self):
        self.test_value(raw=True)

    def test_getobj_getlock(self):
        val1 = self.Value("i", 5)
        lock1 = val1.get_lock()
        obj1 = val1.get_obj()

        val2 = self.Value("i", 5, lock=None)
        lock2 = val2.get_lock()
        obj2 = val2.get_obj()

        lock = self.Lock()
        val3 = self.Value("i", 5, lock=lock)
        lock3 = val3.get_lock()
        obj3 = val3.get_obj()
        self.assertEqual(lock, lock3)

        arr4 = self.Value("i", 5, lock=False)
        self.assertFalse(hasattr(arr4, "get_lock"))
        self.assertFalse(hasattr(arr4, "get_obj"))

        self.assertRaises(AttributeError, self.Value, "i", 5, lock="navalue")

        arr5 = self.RawValue("i", 5)
        self.assertFalse(hasattr(arr5, "get_lock"))
        self.assertFalse(hasattr(arr5, "get_obj"))


# DRAGON
if __name__ == "__main__":
    setUpModule()
    unittest.main()
    tearDownModule()
examples/multiprocessing/unittests/test_value.py
DragonHPC/dragon
#
# Python documentation build configuration file
#
# This file is execfile()d with the current directory set to its containing dir.
#
# The contents of this file are pickled, so don't put values in the namespace
# that aren't pickleable (module imports are okay, they're removed automatically).

import sys, os, time
sys.path.append(os.path.abspath('tools/extensions'))
sys.path.append(os.path.abspath('includes'))

# General configuration
# ---------------------

extensions = ['sphinx.ext.coverage', 'sphinx.ext.doctest',
              'pyspecific', 'c_annotations', 'escape4chm',
              'asdl_highlight', 'peg_highlight']


doctest_global_setup = '''
try:
    import _tkinter
except ImportError:
    _tkinter = None
'''

manpages_url = 'https://manpages.debian.org/{path}'

# General substitutions.
project = 'Python'
copyright = '2001-%s, Python Software Foundation' % time.strftime('%Y')

# We look for the Include/patchlevel.h file in the current Python source tree
# and replace the values accordingly.
import patchlevel
version, release = patchlevel.get_version_info()

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
today = ''
# Else, today_fmt is used as the format for a strftime call.
today_fmt = '%B %d, %Y'

# By default, highlight as Python 3.
highlight_language = 'python3'

# Minimum version of sphinx required
needs_sphinx = '1.8'

# Ignore any .rst files in the venv/ directory.
exclude_patterns = ['venv/*', 'README.rst']
venvdir = os.getenv('VENVDIR')
if venvdir is not None:
    exclude_patterns.append(venvdir + '/*')

# Disable Docutils smartquotes for several translations
smartquotes_excludes = {
    'languages': ['ja', 'fr', 'zh_TW', 'zh_CN'], 'builders': ['man', 'text'],
}

# Avoid a warning with Sphinx >= 2.0
master_doc = 'contents'

# Options for HTML output
# -----------------------

# Use our custom theme.
html_theme = 'python_docs_theme'
html_theme_path = ['tools']
html_theme_options = {
    'collapsiblesidebar': True,
    'issues_url': 'https://docs.python.org/3/bugs.html',
    'root_include_title': False   # We use the version switcher instead.
}

# Short title used e.g. for <title> HTML tags.
html_short_title = '%s Documentation' % release

# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
# using the given strftime format.
html_last_updated_fmt = '%b %d, %Y'

# Path to find HTML templates.
templates_path = ['tools/templates']

# Custom sidebar templates, filenames relative to this file.
html_sidebars = {
    # Defaults taken from http://www.sphinx-doc.org/en/stable/config.html#confval-html_sidebars
    # Removes the quick search block
    '**': ['localtoc.html', 'relations.html', 'customsourcelink.html'],
    'index': ['indexsidebar.html'],
}

# Additional templates that should be rendered to pages.
html_additional_pages = {
    'download': 'download.html',
    'index': 'indexcontent.html',
}

# Output an OpenSearch description file.
html_use_opensearch = 'https://docs.python.org/' + version

# Additional static files.
html_static_path = ['tools/static']

# Output file base name for HTML help builder.
htmlhelp_basename = 'python' + release.replace('.', '')

# Split the index
html_split_index = True


# Options for LaTeX output
# ------------------------

latex_engine = 'xelatex'

# Get LaTeX to handle Unicode correctly
latex_elements = {
}

# Additional stuff for the LaTeX preamble.
latex_elements['preamble'] = r'''
\authoraddress{
  \sphinxstrong{Python Software Foundation}\\
  Email: \sphinxemail{docs@python.org}
}
\let\Verbatim=\OriginalVerbatim
\let\endVerbatim=\endOriginalVerbatim
\setcounter{tocdepth}{2}
'''

# The paper size ('letter' or 'a4').
latex_elements['papersize'] = 'a4'

# The font size ('10pt', '11pt' or '12pt').
latex_elements['pointsize'] = '10pt'

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title, author, document class [howto/manual]).
_stdauthor = r'Guido van Rossum\\and the Python development team'
latex_documents = [
    ('c-api/index', 'c-api.tex',
     'The Python/C API', _stdauthor, 'manual'),
    ('distributing/index', 'distributing.tex',
     'Distributing Python Modules', _stdauthor, 'manual'),
    ('extending/index', 'extending.tex',
     'Extending and Embedding Python', _stdauthor, 'manual'),
    ('installing/index', 'installing.tex',
     'Installing Python Modules', _stdauthor, 'manual'),
    ('library/index', 'library.tex',
     'The Python Library Reference', _stdauthor, 'manual'),
    ('reference/index', 'reference.tex',
     'The Python Language Reference', _stdauthor, 'manual'),
    ('tutorial/index', 'tutorial.tex',
     'Python Tutorial', _stdauthor, 'manual'),
    ('using/index', 'using.tex',
     'Python Setup and Usage', _stdauthor, 'manual'),
    ('faq/index', 'faq.tex',
     'Python Frequently Asked Questions', _stdauthor, 'manual'),
    ('whatsnew/' + version, 'whatsnew.tex',
     'What\'s New in Python', 'A. M. Kuchling', 'howto'),
]
# Collect all HOWTOs individually
latex_documents.extend(('howto/' + fn[:-4], 'howto-' + fn[:-4] + '.tex',
                        '', _stdauthor, 'howto')
                       for fn in os.listdir('howto')
                       if fn.endswith('.rst') and fn != 'index.rst')

# Documents to append as an appendix to all manuals.
latex_appendices = ['glossary', 'about', 'license', 'copyright']

# Options for Epub output
# -----------------------

epub_author = 'Python Documentation Authors'
epub_publisher = 'Python Software Foundation'

# Options for the coverage checker
# --------------------------------

# The coverage checker will ignore all modules/functions/classes whose names
# match any of the following regexes (using re.match).
coverage_ignore_modules = [
    r'[T|t][k|K]',
    r'Tix',
    r'distutils.*',
]

coverage_ignore_functions = [
    'test($|_)',
]

coverage_ignore_classes = [
]

# Glob patterns for C source files for C API coverage, relative to this directory.
coverage_c_path = [
    '../Include/*.h',
]

# Regexes to find C items in the source files.
coverage_c_regexes = {
    'cfunction': (r'^PyAPI_FUNC\(.*\)\s+([^_][\w_]+)'),
    'data': (r'^PyAPI_DATA\(.*\)\s+([^_][\w_]+)'),
    'macro': (r'^#define ([^_][\w_]+)\(.*\)[\s|\\]'),
}

# The coverage checker will ignore all C items whose names match these regexes
# (using re.match) -- the keys must be the same as in coverage_c_regexes.
coverage_ignore_c_items = {
#    'cfunction': [...]
}


# Options for the link checker
# ----------------------------

# Ignore certain URLs.
linkcheck_ignore = [r'https://bugs.python.org/(issue)?\d+',
                    # Ignore PEPs for now, they all have permanent redirects.
                    r'http://www.python.org/dev/peps/pep-\d+']


# Options for extensions
# ----------------------

# Relative filename of the reference count data file.
refcount_file = 'data/refcounts.dat'

# Sphinx 2 and Sphinx 3 compatibility
# -----------------------------------

# bpo-40204: Allow Sphinx 2 syntax in the C domain
c_allow_pre_v3 = True

# bpo-40204: Disable warnings on Sphinx 2 syntax of the C domain since the
# documentation is built with -W (warnings treated as errors).
c_warn_on_allowed_pre_v3 = False

external/cpython_master/Doc/conf.py
DragonHPC/dragon
"""A minimal subset of the locale module used at interpreter startup
(imported by the _io module), in order to reduce startup time.

Don't import directly from third-party code; use the `locale` module instead!
"""

import sys
import _locale

if sys.platform.startswith("win"):
    def getpreferredencoding(do_setlocale=True):
        if sys.flags.utf8_mode:
            return 'UTF-8'
        return _locale._getdefaultlocale()[1]
else:
    try:
        _locale.CODESET
    except AttributeError:
        if hasattr(sys, 'getandroidapilevel'):
            # On Android langinfo.h and CODESET are missing, and UTF-8 is
            # always used in mbstowcs() and wcstombs().
            def getpreferredencoding(do_setlocale=True):
                return 'UTF-8'
        else:
            def getpreferredencoding(do_setlocale=True):
                if sys.flags.utf8_mode:
                    return 'UTF-8'
                # This path for legacy systems needs the more complex
                # getdefaultlocale() function, import the full locale module.
                import locale
                return locale.getpreferredencoding(do_setlocale)
    else:
        def getpreferredencoding(do_setlocale=True):
            assert not do_setlocale
            if sys.flags.utf8_mode:
                return 'UTF-8'
            result = _locale.nl_langinfo(_locale.CODESET)
            if not result and sys.platform == 'darwin':
                # nl_langinfo can return an empty string
                # when the setting has an invalid value.
                # Default to UTF-8 in that case because
                # UTF-8 is the default charset on OSX and
                # returning nothing will crash the
                # interpreter.
                result = 'UTF-8'
            return result

external/cpython_master/Lib/_bootlocale.py
DragonHPC/dragon
import enum

# After the connection is lost, log warnings after this many write()s.
LOG_THRESHOLD_FOR_CONNLOST_WRITES = 5

# Seconds to wait before retrying accept().
ACCEPT_RETRY_DELAY = 1

# Number of stack entries to capture in debug mode.
# The larger the number, the slower the operation in debug mode
# (see extract_stack() in format_helpers.py).
DEBUG_STACK_DEPTH = 10

# Number of seconds to wait for SSL handshake to complete
# The default timeout matches that of Nginx.
SSL_HANDSHAKE_TIMEOUT = 60.0

# Used in sendfile fallback code.  We use fallback for platforms
# that don't support sendfile, or for TLS connections.
SENDFILE_FALLBACK_READBUFFER_SIZE = 1024 * 256

# The enum should be here to break circular dependencies between
# base_events and sslproto
class _SendfileMode(enum.Enum):
    UNSUPPORTED = enum.auto()
    TRY_NATIVE = enum.auto()
    FALLBACK = enum.auto()

external/cpython_master/Lib/asyncio/constants.py
DragonHPC/dragon
import functools
import inspect
import reprlib
import sys
import traceback

from . import constants


def _get_function_source(func):
    func = inspect.unwrap(func)
    if inspect.isfunction(func):
        code = func.__code__
        return (code.co_filename, code.co_firstlineno)
    if isinstance(func, functools.partial):
        return _get_function_source(func.func)
    if isinstance(func, functools.partialmethod):
        return _get_function_source(func.func)
    return None


def _format_callback_source(func, args):
    func_repr = _format_callback(func, args, None)
    source = _get_function_source(func)
    if source:
        func_repr += f' at {source[0]}:{source[1]}'
    return func_repr


def _format_args_and_kwargs(args, kwargs):
    """Format function arguments and keyword arguments.

    Special case for a single parameter: ('hello',) is formatted as ('hello').
    """
    # use reprlib to limit the length of the output
    items = []
    if args:
        items.extend(reprlib.repr(arg) for arg in args)
    if kwargs:
        items.extend(f'{k}={reprlib.repr(v)}' for k, v in kwargs.items())
    return '({})'.format(', '.join(items))


def _format_callback(func, args, kwargs, suffix=''):
    if isinstance(func, functools.partial):
        suffix = _format_args_and_kwargs(args, kwargs) + suffix
        return _format_callback(func.func, func.args, func.keywords, suffix)

    if hasattr(func, '__qualname__') and func.__qualname__:
        func_repr = func.__qualname__
    elif hasattr(func, '__name__') and func.__name__:
        func_repr = func.__name__
    else:
        func_repr = repr(func)

    func_repr += _format_args_and_kwargs(args, kwargs)
    if suffix:
        func_repr += suffix
    return func_repr


def extract_stack(f=None, limit=None):
    """Replacement for traceback.extract_stack() that only does the
    necessary work for asyncio debug mode.
    """
    if f is None:
        f = sys._getframe().f_back
    if limit is None:
        # Limit the amount of work to a reasonable amount, as extract_stack()
        # can be called for each coroutine and future in debug mode.
        limit = constants.DEBUG_STACK_DEPTH
    stack = traceback.StackSummary.extract(traceback.walk_stack(f),
                                           limit=limit,
                                           lookup_lines=False)
    stack.reverse()
    return stack

external/cpython_master/Lib/asyncio/format_helpers.py
DragonHPC/dragon
import unittest
from ctypes import *

try:
    WINFUNCTYPE
except NameError:
    # fake to enable this test on Linux
    WINFUNCTYPE = CFUNCTYPE

import _ctypes_test
lib = CDLL(_ctypes_test.__file__)

class CFuncPtrTestCase(unittest.TestCase):
    def test_basic(self):
        X = WINFUNCTYPE(c_int, c_int, c_int)

        def func(*args):
            return len(args)

        x = X(func)
        self.assertEqual(x.restype, c_int)
        self.assertEqual(x.argtypes, (c_int, c_int))
        self.assertEqual(sizeof(x), sizeof(c_voidp))
        self.assertEqual(sizeof(X), sizeof(c_voidp))

    def test_first(self):
        StdCallback = WINFUNCTYPE(c_int, c_int, c_int)
        CdeclCallback = CFUNCTYPE(c_int, c_int, c_int)

        def func(a, b):
            return a + b

        s = StdCallback(func)
        c = CdeclCallback(func)

        self.assertEqual(s(1, 2), 3)
        self.assertEqual(c(1, 2), 3)
        # The following no longer raises a TypeError - it is now
        # possible, as in C, to call cdecl functions with more parameters.
        #self.assertRaises(TypeError, c, 1, 2, 3)
        self.assertEqual(c(1, 2, 3, 4, 5, 6), 3)
        if not WINFUNCTYPE is CFUNCTYPE:
            self.assertRaises(TypeError, s, 1, 2, 3)

    def test_structures(self):
        WNDPROC = WINFUNCTYPE(c_long, c_int, c_int, c_int, c_int)

        def wndproc(hwnd, msg, wParam, lParam):
            return hwnd + msg + wParam + lParam

        HINSTANCE = c_int
        HICON = c_int
        HCURSOR = c_int
        LPCTSTR = c_char_p

        class WNDCLASS(Structure):
            _fields_ = [("style", c_uint),
                        ("lpfnWndProc", WNDPROC),
                        ("cbClsExtra", c_int),
                        ("cbWndExtra", c_int),
                        ("hInstance", HINSTANCE),
                        ("hIcon", HICON),
                        ("hCursor", HCURSOR),
                        ("lpszMenuName", LPCTSTR),
                        ("lpszClassName", LPCTSTR)]

        wndclass = WNDCLASS()
        wndclass.lpfnWndProc = WNDPROC(wndproc)

        WNDPROC_2 = WINFUNCTYPE(c_long, c_int, c_int, c_int, c_int)

        # This is no longer true, now that WINFUNCTYPE caches created types internally.
        ## # CFuncPtr subclasses are compared by identity, so this raises a TypeError:
        ## self.assertRaises(TypeError, setattr, wndclass,
        ##                  "lpfnWndProc", WNDPROC_2(wndproc))
        # instead:

        self.assertIs(WNDPROC, WNDPROC_2)
        # 'wndclass.lpfnWndProc' leaks 94 references.  Why?
        self.assertEqual(wndclass.lpfnWndProc(1, 2, 3, 4), 10)


        f = wndclass.lpfnWndProc

        del wndclass
        del wndproc

        self.assertEqual(f(10, 11, 12, 13), 46)

    def test_dllfunctions(self):

        def NoNullHandle(value):
            if not value:
                raise WinError()
            return value

        strchr = lib.my_strchr
        strchr.restype = c_char_p
        strchr.argtypes = (c_char_p, c_char)
        self.assertEqual(strchr(b"abcdefghi", b"b"), b"bcdefghi")
        self.assertEqual(strchr(b"abcdefghi", b"x"), None)


        strtok = lib.my_strtok
        strtok.restype = c_char_p
        # Neither of this does work: strtok changes the buffer it is passed
##        strtok.argtypes = (c_char_p, c_char_p)
##        strtok.argtypes = (c_string, c_char_p)

        def c_string(init):
            size = len(init) + 1
            return (c_char*size)(*init)

        s = b"a\nb\nc"
        b = c_string(s)

##        b = (c_char * (len(s)+1))()
##        b.value = s

##        b = c_string(s)
        self.assertEqual(strtok(b, b"\n"), b"a")
        self.assertEqual(strtok(None, b"\n"), b"b")
        self.assertEqual(strtok(None, b"\n"), b"c")
        self.assertEqual(strtok(None, b"\n"), None)

    def test_abstract(self):
        from ctypes import _CFuncPtr

        self.assertRaises(TypeError, _CFuncPtr, 13, "name", 42, "iid")

if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/ctypes/test/test_funcptr.py
DragonHPC/dragon
"""
Here is probably the place to write the docs, since the test-cases
show how the type behave.

Later...
"""

from ctypes import *
from ctypes.test import need_symbol
import sys, unittest

try:
    WINFUNCTYPE
except NameError:
    # fake to enable this test on Linux
    WINFUNCTYPE = CFUNCTYPE

import _ctypes_test
dll = CDLL(_ctypes_test.__file__)
if sys.platform == "win32":
    windll = WinDLL(_ctypes_test.__file__)

class POINT(Structure):
    _fields_ = [("x", c_int), ("y", c_int)]
class RECT(Structure):
    _fields_ = [("left", c_int), ("top", c_int),
                ("right", c_int), ("bottom", c_int)]
class FunctionTestCase(unittest.TestCase):

    def test_mro(self):
        # in Python 2.3, this raises TypeError: MRO conflict among bases classes,
        # in Python 2.2 it works.
        #
        # But in early versions of _ctypes.c, the result of tp_new
        # wasn't checked, and it even crashed Python.
        # Found by Greg Chapman.

        try:
            class X(object, Array):
                _length_ = 5
                _type_ = "i"
        except TypeError:
            pass


        from _ctypes import _Pointer
        try:
            class X(object, _Pointer):
                pass
        except TypeError:
            pass

        from _ctypes import _SimpleCData
        try:
            class X(object, _SimpleCData):
                _type_ = "i"
        except TypeError:
            pass

        try:
            class X(object, Structure):
                _fields_ = []
        except TypeError:
            pass


    @need_symbol('c_wchar')
    def test_wchar_parm(self):
        f = dll._testfunc_i_bhilfd
        f.argtypes = [c_byte, c_wchar, c_int, c_long, c_float, c_double]
        result = f(1, "x", 3, 4, 5.0, 6.0)
        self.assertEqual(result, 139)
        self.assertEqual(type(result), int)

    @need_symbol('c_wchar')
    def test_wchar_result(self):
        f = dll._testfunc_i_bhilfd
        f.argtypes = [c_byte, c_short, c_int, c_long, c_float, c_double]
        f.restype = c_wchar
        result = f(0, 0, 0, 0, 0, 0)
        self.assertEqual(result, '\x00')

    def test_voidresult(self):
        f = dll._testfunc_v
        f.restype = None
        f.argtypes = [c_int, c_int, POINTER(c_int)]
        result = c_int()
        self.assertEqual(None, f(1, 2, byref(result)))
        self.assertEqual(result.value, 3)

    def test_intresult(self):
        f = dll._testfunc_i_bhilfd
        f.argtypes = [c_byte, c_short, c_int, c_long, c_float, c_double]
        f.restype = c_int
        result = f(1, 2, 3, 4, 5.0, 6.0)
        self.assertEqual(result, 21)
        self.assertEqual(type(result), int)

        result = f(-1, -2, -3, -4, -5.0, -6.0)
        self.assertEqual(result, -21)
        self.assertEqual(type(result), int)

        # If we declare the function to return a short,
        # is the high part split off?
        f.restype = c_short
        result = f(1, 2, 3, 4, 5.0, 6.0)
        self.assertEqual(result, 21)
        self.assertEqual(type(result), int)

        result = f(1, 2, 3, 0x10004, 5.0, 6.0)
        self.assertEqual(result, 21)
        self.assertEqual(type(result), int)

        # You cannot assign character format codes as restype any longer
        self.assertRaises(TypeError, setattr, f, "restype", "i")

    def test_floatresult(self):
        f = dll._testfunc_f_bhilfd
        f.argtypes = [c_byte, c_short, c_int, c_long, c_float, c_double]
        f.restype = c_float
        result = f(1, 2, 3, 4, 5.0, 6.0)
        self.assertEqual(result, 21)
        self.assertEqual(type(result), float)

        result = f(-1, -2, -3, -4, -5.0, -6.0)
        self.assertEqual(result, -21)
        self.assertEqual(type(result), float)

    def test_doubleresult(self):
        f = dll._testfunc_d_bhilfd
        f.argtypes = [c_byte, c_short, c_int, c_long, c_float, c_double]
        f.restype = c_double
        result = f(1, 2, 3, 4, 5.0, 6.0)
        self.assertEqual(result, 21)
        self.assertEqual(type(result), float)

        result = f(-1, -2, -3, -4, -5.0, -6.0)
        self.assertEqual(result, -21)
        self.assertEqual(type(result), float)

    def test_longdoubleresult(self):
        f = dll._testfunc_D_bhilfD
        f.argtypes = [c_byte, c_short, c_int, c_long, c_float, c_longdouble]
        f.restype = c_longdouble
        result = f(1, 2, 3, 4, 5.0, 6.0)
        self.assertEqual(result, 21)
        self.assertEqual(type(result), float)

        result = f(-1, -2, -3, -4, -5.0, -6.0)
        self.assertEqual(result, -21)
        self.assertEqual(type(result), float)

    @need_symbol('c_longlong')
    def test_longlongresult(self):
        f = dll._testfunc_q_bhilfd
        f.restype = c_longlong
        f.argtypes = [c_byte, c_short, c_int, c_long, c_float, c_double]
        result = f(1, 2, 3, 4, 5.0, 6.0)
        self.assertEqual(result, 21)

        f = dll._testfunc_q_bhilfdq
        f.restype = c_longlong
        f.argtypes = [c_byte, c_short, c_int, c_long, c_float, c_double, c_longlong]
        result = f(1, 2, 3, 4, 5.0, 6.0, 21)
        self.assertEqual(result, 42)

    def test_stringresult(self):
        f = dll._testfunc_p_p
        f.argtypes = None
        f.restype = c_char_p
        result = f(b"123")
        self.assertEqual(result, b"123")

        result = f(None)
        self.assertEqual(result, None)

    def test_pointers(self):
        f = dll._testfunc_p_p
        f.restype = POINTER(c_int)
        f.argtypes = [POINTER(c_int)]

        # This only works if the value c_int(42) passed to the
        # function is still alive while the pointer (the result) is
        # used.

        v = c_int(42)

        self.assertEqual(pointer(v).contents.value, 42)
        result = f(pointer(v))
        self.assertEqual(type(result), POINTER(c_int))
        self.assertEqual(result.contents.value, 42)

        # This on works...
        result = f(pointer(v))
        self.assertEqual(result.contents.value, v.value)

        p = pointer(c_int(99))
        result = f(p)
        self.assertEqual(result.contents.value, 99)

        arg = byref(v)
        result = f(arg)
        self.assertNotEqual(result.contents, v.value)

        self.assertRaises(ArgumentError, f, byref(c_short(22)))

        # It is dangerous, however, because you don't control the lifetime
        # of the pointer:
        result = f(byref(c_int(99)))
        self.assertNotEqual(result.contents, 99)

    def test_errors(self):
        f = dll._testfunc_p_p
        f.restype = c_int

        class X(Structure):
            _fields_ = [("y", c_int)]

        self.assertRaises(TypeError, f, X()) #cannot convert parameter

    ################################################################
    def test_shorts(self):
        f = dll._testfunc_callback_i_if

        args = []
        expected = [262144, 131072, 65536, 32768, 16384, 8192, 4096, 2048,
                    1024, 512, 256, 128, 64, 32, 16, 8, 4, 2, 1]

        def callback(v):
            args.append(v)
            return v

        CallBack = CFUNCTYPE(c_int, c_int)

        cb = CallBack(callback)
        f(2**18, cb)
        self.assertEqual(args, expected)

    ################################################################


    def test_callbacks(self):
        f = dll._testfunc_callback_i_if
        f.restype = c_int
        f.argtypes = None

        MyCallback = CFUNCTYPE(c_int, c_int)

        def callback(value):
            #print "called back with", value
            return value

        cb = MyCallback(callback)
        result = f(-10, cb)
        self.assertEqual(result, -18)

        # test with prototype
        f.argtypes = [c_int, MyCallback]
        cb = MyCallback(callback)
        result = f(-10, cb)
        self.assertEqual(result, -18)

        AnotherCallback = WINFUNCTYPE(c_int, c_int, c_int, c_int, c_int)

        # check that the prototype works: we call f with wrong
        # argument types
        cb = AnotherCallback(callback)
        self.assertRaises(ArgumentError, f, -10, cb)


    def test_callbacks_2(self):
        # Can also use simple datatypes as argument type specifiers
        # for the callback function.
        # In this case the call receives an instance of that type
        f = dll._testfunc_callback_i_if
        f.restype = c_int

        MyCallback = CFUNCTYPE(c_int, c_int)

        f.argtypes = [c_int, MyCallback]

        def callback(value):
            #print "called back with", value
            self.assertEqual(type(value), int)
            return value

        cb = MyCallback(callback)
        result = f(-10, cb)
        self.assertEqual(result, -18)

    @need_symbol('c_longlong')
    def test_longlong_callbacks(self):

        f = dll._testfunc_callback_q_qf
        f.restype = c_longlong

        MyCallback = CFUNCTYPE(c_longlong, c_longlong)

        f.argtypes = [c_longlong, MyCallback]

        def callback(value):
            self.assertIsInstance(value, int)
            return value & 0x7FFFFFFF

        cb = MyCallback(callback)

        self.assertEqual(13577625587, f(1000000000000, cb))

    def test_errors(self):
        self.assertRaises(AttributeError, getattr, dll, "_xxx_yyy")
        self.assertRaises(ValueError, c_int.in_dll, dll, "_xxx_yyy")

    def test_byval(self):

        # without prototype
        ptin = POINT(1, 2)
        ptout = POINT()
        # EXPORT int _testfunc_byval(point in, point *pout)
        result = dll._testfunc_byval(ptin, byref(ptout))
        got = result, ptout.x, ptout.y
        expected = 3, 1, 2
        self.assertEqual(got, expected)

        # with prototype
        ptin = POINT(101, 102)
        ptout = POINT()
        dll._testfunc_byval.argtypes = (POINT, POINTER(POINT))
        dll._testfunc_byval.restype = c_int
        result = dll._testfunc_byval(ptin, byref(ptout))
        got = result, ptout.x, ptout.y
        expected = 203, 101, 102
        self.assertEqual(got, expected)

    def test_struct_return_2H(self):
        class S2H(Structure):
            _fields_ = [("x", c_short),
                        ("y", c_short)]
        dll.ret_2h_func.restype = S2H
        dll.ret_2h_func.argtypes = [S2H]
        inp = S2H(99, 88)
        s2h = dll.ret_2h_func(inp)
        self.assertEqual((s2h.x, s2h.y), (99*2, 88*3))

    @unittest.skipUnless(sys.platform == "win32", 'Windows-specific test')
    def test_struct_return_2H_stdcall(self):
        class S2H(Structure):
            _fields_ = [("x", c_short),
                        ("y", c_short)]

        windll.s_ret_2h_func.restype = S2H
        windll.s_ret_2h_func.argtypes = [S2H]
        s2h = windll.s_ret_2h_func(S2H(99, 88))
        self.assertEqual((s2h.x, s2h.y), (99*2, 88*3))

    def test_struct_return_8H(self):
        class S8I(Structure):
            _fields_ = [("a", c_int),
                        ("b", c_int),
                        ("c", c_int),
                        ("d", c_int),
                        ("e", c_int),
                        ("f", c_int),
                        ("g", c_int),
                        ("h", c_int)]
        dll.ret_8i_func.restype = S8I
        dll.ret_8i_func.argtypes = [S8I]
        inp = S8I(9, 8, 7, 6, 5, 4, 3, 2)
        s8i = dll.ret_8i_func(inp)
        self.assertEqual((s8i.a, s8i.b, s8i.c, s8i.d, s8i.e, s8i.f, s8i.g, s8i.h),
                             (9*2, 8*3, 7*4, 6*5, 5*6, 4*7, 3*8, 2*9))

    @unittest.skipUnless(sys.platform == "win32", 'Windows-specific test')
    def test_struct_return_8H_stdcall(self):
        class S8I(Structure):
            _fields_ = [("a", c_int),
                        ("b", c_int),
                        ("c", c_int),
                        ("d", c_int),
                        ("e", c_int),
                        ("f", c_int),
                        ("g", c_int),
                        ("h", c_int)]
        windll.s_ret_8i_func.restype = S8I
        windll.s_ret_8i_func.argtypes = [S8I]
        inp = S8I(9, 8, 7, 6, 5, 4, 3, 2)
        s8i = windll.s_ret_8i_func(inp)
        self.assertEqual(
                (s8i.a, s8i.b, s8i.c, s8i.d, s8i.e, s8i.f, s8i.g, s8i.h),
                (9*2, 8*3, 7*4, 6*5, 5*6, 4*7, 3*8, 2*9))

    def test_sf1651235(self):
        # see http://www.python.org/sf/1651235

        proto = CFUNCTYPE(c_int, RECT, POINT)
        def callback(*args):
            return 0

        callback = proto(callback)
        self.assertRaises(ArgumentError, lambda: callback((1, 2, 3, 4), POINT()))

if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/ctypes/test/test_functions.py
DragonHPC/dragon
import unittest
from ctypes import *

class MyInt(c_int):
    def __eq__(self, other):
        if type(other) != MyInt:
            return NotImplementedError
        return self.value == other.value

class Test(unittest.TestCase):

    def test_compare(self):
        self.assertEqual(MyInt(3), MyInt(3))
        self.assertNotEqual(MyInt(42), MyInt(43))

    def test_ignore_retval(self):
        # Test if the return value of a callback is ignored
        # if restype is None
        proto = CFUNCTYPE(None)
        def func():
            return (1, "abc", None)

        cb = proto(func)
        self.assertEqual(None, cb())


    def test_int_callback(self):
        args = []
        def func(arg):
            args.append(arg)
            return arg

        cb = CFUNCTYPE(None, MyInt)(func)

        self.assertEqual(None, cb(42))
        self.assertEqual(type(args[-1]), MyInt)

        cb = CFUNCTYPE(c_int, c_int)(func)

        self.assertEqual(42, cb(42))
        self.assertEqual(type(args[-1]), int)

    def test_int_struct(self):
        class X(Structure):
            _fields_ = [("x", MyInt)]

        self.assertEqual(X().x, MyInt())

        s = X()
        s.x = MyInt(42)

        self.assertEqual(s.x, MyInt(42))

if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/ctypes/test/test_simplesubclasses.py
DragonHPC/dragon
"""Simple textbox editing widget with Emacs-like keybindings."""

import curses
import curses.ascii

def rectangle(win, uly, ulx, lry, lrx):
    """Draw a rectangle with corners at the provided upper-left
    and lower-right coordinates.
    """
    win.vline(uly+1, ulx, curses.ACS_VLINE, lry - uly - 1)
    win.hline(uly, ulx+1, curses.ACS_HLINE, lrx - ulx - 1)
    win.hline(lry, ulx+1, curses.ACS_HLINE, lrx - ulx - 1)
    win.vline(uly+1, lrx, curses.ACS_VLINE, lry - uly - 1)
    win.addch(uly, ulx, curses.ACS_ULCORNER)
    win.addch(uly, lrx, curses.ACS_URCORNER)
    win.addch(lry, lrx, curses.ACS_LRCORNER)
    win.addch(lry, ulx, curses.ACS_LLCORNER)

class Textbox:
    """Editing widget using the interior of a window object.
     Supports the following Emacs-like key bindings:

    Ctrl-A      Go to left edge of window.
    Ctrl-B      Cursor left, wrapping to previous line if appropriate.
    Ctrl-D      Delete character under cursor.
    Ctrl-E      Go to right edge (stripspaces off) or end of line (stripspaces on).
    Ctrl-F      Cursor right, wrapping to next line when appropriate.
    Ctrl-G      Terminate, returning the window contents.
    Ctrl-H      Delete character backward.
    Ctrl-J      Terminate if the window is 1 line, otherwise insert newline.
    Ctrl-K      If line is blank, delete it, otherwise clear to end of line.
    Ctrl-L      Refresh screen.
    Ctrl-N      Cursor down; move down one line.
    Ctrl-O      Insert a blank line at cursor location.
    Ctrl-P      Cursor up; move up one line.

    Move operations do nothing if the cursor is at an edge where the movement
    is not possible.  The following synonyms are supported where possible:

    KEY_LEFT = Ctrl-B, KEY_RIGHT = Ctrl-F, KEY_UP = Ctrl-P, KEY_DOWN = Ctrl-N
    KEY_BACKSPACE = Ctrl-h
    """
    def __init__(self, win, insert_mode=False):
        self.win = win
        self.insert_mode = insert_mode
        self._update_max_yx()
        self.stripspaces = 1
        self.lastcmd = None
        win.keypad(1)

    def _update_max_yx(self):
        maxy, maxx = self.win.getmaxyx()
        self.maxy = maxy - 1
        self.maxx = maxx - 1

    def _end_of_line(self, y):
        """Go to the location of the first blank on the given line,
        returning the index of the last non-blank character."""
        self._update_max_yx()
        last = self.maxx
        while True:
            if curses.ascii.ascii(self.win.inch(y, last)) != curses.ascii.SP:
                last = min(self.maxx, last+1)
                break
            elif last == 0:
                break
            last = last - 1
        return last

    def _insert_printable_char(self, ch):
        self._update_max_yx()
        (y, x) = self.win.getyx()
        backyx = None
        while y < self.maxy or x < self.maxx:
            if self.insert_mode:
                oldch = self.win.inch()
            # The try-catch ignores the error we trigger from some curses
            # versions by trying to write into the lowest-rightmost spot
            # in the window.
            try:
                self.win.addch(ch)
            except curses.error:
                pass
            if not self.insert_mode or not curses.ascii.isprint(oldch):
                break
            ch = oldch
            (y, x) = self.win.getyx()
            # Remember where to put the cursor back since we are in insert_mode
            if backyx is None:
                backyx = y, x

        if backyx is not None:
            self.win.move(*backyx)

    def do_command(self, ch):
        "Process a single editing command."
        self._update_max_yx()
        (y, x) = self.win.getyx()
        self.lastcmd = ch
        if curses.ascii.isprint(ch):
            if y < self.maxy or x < self.maxx:
                self._insert_printable_char(ch)
        elif ch == curses.ascii.SOH:                           # ^a
            self.win.move(y, 0)
        elif ch in (curses.ascii.STX,curses.KEY_LEFT, curses.ascii.BS,curses.KEY_BACKSPACE):
            if x > 0:
                self.win.move(y, x-1)
            elif y == 0:
                pass
            elif self.stripspaces:
                self.win.move(y-1, self._end_of_line(y-1))
            else:
                self.win.move(y-1, self.maxx)
            if ch in (curses.ascii.BS, curses.KEY_BACKSPACE):
                self.win.delch()
        elif ch == curses.ascii.EOT:                           # ^d
            self.win.delch()
        elif ch == curses.ascii.ENQ:                           # ^e
            if self.stripspaces:
                self.win.move(y, self._end_of_line(y))
            else:
                self.win.move(y, self.maxx)
        elif ch in (curses.ascii.ACK, curses.KEY_RIGHT):       # ^f
            if x < self.maxx:
                self.win.move(y, x+1)
            elif y == self.maxy:
                pass
            else:
                self.win.move(y+1, 0)
        elif ch == curses.ascii.BEL:                           # ^g
            return 0
        elif ch == curses.ascii.NL:                            # ^j
            if self.maxy == 0:
                return 0
            elif y < self.maxy:
                self.win.move(y+1, 0)
        elif ch == curses.ascii.VT:                            # ^k
            if x == 0 and self._end_of_line(y) == 0:
                self.win.deleteln()
            else:
                # first undo the effect of self._end_of_line
                self.win.move(y, x)
                self.win.clrtoeol()
        elif ch == curses.ascii.FF:                            # ^l
            self.win.refresh()
        elif ch in (curses.ascii.SO, curses.KEY_DOWN):         # ^n
            if y < self.maxy:
                self.win.move(y+1, x)
                if x > self._end_of_line(y+1):
                    self.win.move(y+1, self._end_of_line(y+1))
        elif ch == curses.ascii.SI:                            # ^o
            self.win.insertln()
        elif ch in (curses.ascii.DLE, curses.KEY_UP):          # ^p
            if y > 0:
                self.win.move(y-1, x)
                if x > self._end_of_line(y-1):
                    self.win.move(y-1, self._end_of_line(y-1))
        return 1

    def gather(self):
        "Collect and return the contents of the window."
        result = ""
        self._update_max_yx()
        for y in range(self.maxy+1):
            self.win.move(y, 0)
            stop = self._end_of_line(y)
            if stop == 0 and self.stripspaces:
                continue
            for x in range(self.maxx+1):
                if self.stripspaces and x > stop:
                    break
                result = result + chr(curses.ascii.ascii(self.win.inch(y, x)))
            if self.maxy > 0:
                result = result + "\n"
        return result

    def edit(self, validate=None):
        "Edit in the widget window and collect the results."
        while 1:
            ch = self.win.getch()
            if validate:
                ch = validate(ch)
            if not ch:
                continue
            if not self.do_command(ch):
                break
            self.win.refresh()
        return self.gather()

if __name__ == '__main__':
    def test_editbox(stdscr):
        ncols, nlines = 9, 4
        uly, ulx = 15, 20
        stdscr.addstr(uly-2, ulx, "Use Ctrl-G to end editing.")
        win = curses.newwin(nlines, ncols, uly, ulx)
        rectangle(stdscr, uly-1, ulx-1, uly + nlines, ulx + ncols)
        stdscr.refresh()
        return Textbox(win).edit()

    str = curses.wrapper(test_editbox)
    print('Contents of text box:', repr(str))

external/cpython_master/Lib/curses/textpad.py
DragonHPC/dragon
# Copyright (C) 2001-2006 Python Software Foundation
# Author: Barry Warsaw
# Contact: email-sig@python.org

"""Class representing image/* type MIME documents."""

__all__ = ['MIMEImage']

import imghdr

from email import encoders
from email.mime.nonmultipart import MIMENonMultipart



class MIMEImage(MIMENonMultipart):
    """Class for generating image/* type MIME documents."""

    def __init__(self, _imagedata, _subtype=None,
                 _encoder=encoders.encode_base64, *, policy=None, **_params):
        """Create an image/* type MIME document.

        _imagedata is a string containing the raw image data.  If this data
        can be decoded by the standard Python `imghdr' module, then the
        subtype will be automatically included in the Content-Type header.
        Otherwise, you can specify the specific image subtype via the _subtype
        parameter.

        _encoder is a function which will perform the actual encoding for
        transport of the image data.  It takes one argument, which is this
        Image instance.  It should use get_payload() and set_payload() to
        change the payload to the encoded form.  It should also add any
        Content-Transfer-Encoding or other headers to the message as
        necessary.  The default encoding is Base64.

        Any additional keyword arguments are passed to the base class
        constructor, which turns them into parameters on the Content-Type
        header.
        """
        if _subtype is None:
            _subtype = imghdr.what(None, _imagedata)
        if _subtype is None:
            raise TypeError('Could not guess image MIME subtype')
        MIMENonMultipart.__init__(self, 'image', _subtype, policy=policy,
                                  **_params)
        self.set_payload(_imagedata)
        _encoder(self)

external/cpython_master/Lib/email/mime/image.py
DragonHPC/dragon
# Copyright (C) 2001-2006 Python Software Foundation
# Author: Ben Gertzfield
# Contact: email-sig@python.org

"""Quoted-printable content transfer encoding per RFCs 2045-2047.

This module handles the content transfer encoding method defined in RFC 2045
to encode US ASCII-like 8-bit data called `quoted-printable'.  It is used to
safely encode text that is in a character set similar to the 7-bit US ASCII
character set, but that includes some 8-bit characters that are normally not
allowed in email bodies or headers.

Quoted-printable is very space-inefficient for encoding binary files; use the
email.base64mime module for that instead.

This module provides an interface to encode and decode both headers and bodies
with quoted-printable encoding.

RFC 2045 defines a method for including character set information in an
`encoded-word' in a header.  This method is commonly used for 8-bit real names
in To:/From:/Cc: etc. fields, as well as Subject: lines.

This module does not do the line wrapping or end-of-line character
conversion necessary for proper internationalized headers; it only
does dumb encoding and decoding.  To deal with the various line
wrapping issues, use the email.header module.
"""

__all__ = [
    'body_decode',
    'body_encode',
    'body_length',
    'decode',
    'decodestring',
    'header_decode',
    'header_encode',
    'header_length',
    'quote',
    'unquote',
    ]

import re

from string import ascii_letters, digits, hexdigits

CRLF = '\r\n'
NL = '\n'
EMPTYSTRING = ''

# Build a mapping of octets to the expansion of that octet.  Since we're only
# going to have 256 of these things, this isn't terribly inefficient
# space-wise.  Remember that headers and bodies have different sets of safe
# characters.  Initialize both maps with the full expansion, and then override
# the safe bytes with the more compact form.
_QUOPRI_MAP = ['=%02X' % c for c in range(256)]
_QUOPRI_HEADER_MAP = _QUOPRI_MAP[:]
_QUOPRI_BODY_MAP = _QUOPRI_MAP[:]

# Safe header bytes which need no encoding.
for c in b'-!*+/' + ascii_letters.encode('ascii') + digits.encode('ascii'):
    _QUOPRI_HEADER_MAP[c] = chr(c)
# Headers have one other special encoding; spaces become underscores.
_QUOPRI_HEADER_MAP[ord(' ')] = '_'

# Safe body bytes which need no encoding.
for c in (b' !"#$%&\'()*+,-./0123456789:;<>'
          b'?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`'
          b'abcdefghijklmnopqrstuvwxyz{|}~\t'):
    _QUOPRI_BODY_MAP[c] = chr(c)



# Helpers
def header_check(octet):
    """Return True if the octet should be escaped with header quopri."""
    return chr(octet) != _QUOPRI_HEADER_MAP[octet]


def body_check(octet):
    """Return True if the octet should be escaped with body quopri."""
    return chr(octet) != _QUOPRI_BODY_MAP[octet]


def header_length(bytearray):
    """Return a header quoted-printable encoding length.

    Note that this does not include any RFC 2047 chrome added by
    `header_encode()`.

    :param bytearray: An array of bytes (a.k.a. octets).
    :return: The length in bytes of the byte array when it is encoded with
        quoted-printable for headers.
    """
    return sum(len(_QUOPRI_HEADER_MAP[octet]) for octet in bytearray)


def body_length(bytearray):
    """Return a body quoted-printable encoding length.

    :param bytearray: An array of bytes (a.k.a. octets).
    :return: The length in bytes of the byte array when it is encoded with
        quoted-printable for bodies.
    """
    return sum(len(_QUOPRI_BODY_MAP[octet]) for octet in bytearray)


def _max_append(L, s, maxlen, extra=''):
    if not isinstance(s, str):
        s = chr(s)
    if not L:
        L.append(s.lstrip())
    elif len(L[-1]) + len(s) <= maxlen:
        L[-1] += extra + s
    else:
        L.append(s.lstrip())


def unquote(s):
    """Turn a string in the form =AB to the ASCII character with value 0xab"""
    return chr(int(s[1:3], 16))


def quote(c):
    return _QUOPRI_MAP[ord(c)]


def header_encode(header_bytes, charset='iso-8859-1'):
    """Encode a single header line with quoted-printable (like) encoding.

    Defined in RFC 2045, this `Q' encoding is similar to quoted-printable, but
    used specifically for email header fields to allow charsets with mostly 7
    bit characters (and some 8 bit) to remain more or less readable in non-RFC
    2045 aware mail clients.

    charset names the character set to use in the RFC 2046 header.  It
    defaults to iso-8859-1.
    """
    # Return empty headers as an empty string.
    if not header_bytes:
        return ''
    # Iterate over every byte, encoding if necessary.
    encoded = header_bytes.decode('latin1').translate(_QUOPRI_HEADER_MAP)
    # Now add the RFC chrome to each encoded chunk and glue the chunks
    # together.
    return '=?%s?q?%s?=' % (charset, encoded)


_QUOPRI_BODY_ENCODE_MAP = _QUOPRI_BODY_MAP[:]
for c in b'\r\n':
    _QUOPRI_BODY_ENCODE_MAP[c] = chr(c)

def body_encode(body, maxlinelen=76, eol=NL):
    """Encode with quoted-printable, wrapping at maxlinelen characters.

    Each line of encoded text will end with eol, which defaults to "\\n".  Set
    this to "\\r\\n" if you will be using the result of this function directly
    in an email.

    Each line will be wrapped at, at most, maxlinelen characters before the
    eol string (maxlinelen defaults to 76 characters, the maximum value
    permitted by RFC 2045).  Long lines will have the 'soft line break'
    quoted-printable character "=" appended to them, so the decoded text will
    be identical to the original text.

    The minimum maxlinelen is 4 to have room for a quoted character ("=XX")
    followed by a soft line break.  Smaller values will generate a
    ValueError.

    """

    if maxlinelen < 4:
        raise ValueError("maxlinelen must be at least 4")
    if not body:
        return body

    # quote special characters
    body = body.translate(_QUOPRI_BODY_ENCODE_MAP)

    soft_break = '=' + eol
    # leave space for the '=' at the end of a line
    maxlinelen1 = maxlinelen - 1

    encoded_body = []
    append = encoded_body.append

    for line in body.splitlines():
        # break up the line into pieces no longer than maxlinelen - 1
        start = 0
        laststart = len(line) - 1 - maxlinelen
        while start <= laststart:
            stop = start + maxlinelen1
            # make sure we don't break up an escape sequence
            if line[stop - 2] == '=':
                append(line[start:stop - 1])
                start = stop - 2
            elif line[stop - 1] == '=':
                append(line[start:stop])
                start = stop - 1
            else:
                append(line[start:stop] + '=')
                start = stop

        # handle rest of line, special case if line ends in whitespace
        if line and line[-1] in ' \t':
            room = start - laststart
            if room >= 3:
                # It's a whitespace character at end-of-line, and we have room
                # for the three-character quoted encoding.
                q = quote(line[-1])
            elif room == 2:
                # There's room for the whitespace character and a soft break.
                q = line[-1] + soft_break
            else:
                # There's room only for a soft break.  The quoted whitespace
                # will be the only content on the subsequent line.
                q = soft_break + quote(line[-1])
            append(line[start:-1] + q)
        else:
            append(line[start:])

    # add back final newline if present
    if body[-1] in CRLF:
        append('')

    return eol.join(encoded_body)



# BAW: I'm not sure if the intent was for the signature of this function to be
# the same as base64MIME.decode() or not...
def decode(encoded, eol=NL):
    """Decode a quoted-printable string.

    Lines are separated with eol, which defaults to \\n.
    """
    if not encoded:
        return encoded
    # BAW: see comment in encode() above.  Again, we're building up the
    # decoded string with string concatenation, which could be done much more
    # efficiently.
    decoded = ''

    for line in encoded.splitlines():
        line = line.rstrip()
        if not line:
            decoded += eol
            continue

        i = 0
        n = len(line)
        while i < n:
            c = line[i]
            if c != '=':
                decoded += c
                i += 1
            # Otherwise, c == "=".  Are we at the end of the line?  If so, add
            # a soft line break.
            elif i+1 == n:
                i += 1
                continue
            # Decode if in form =AB
            elif i+2 < n and line[i+1] in hexdigits and line[i+2] in hexdigits:
                decoded += unquote(line[i:i+3])
                i += 3
            # Otherwise, not in form =AB, pass literally
            else:
                decoded += c
                i += 1

            if i == n:
                decoded += eol
    # Special case if original string did not end with eol
    if encoded[-1] not in '\r\n' and decoded.endswith(eol):
        decoded = decoded[:-1]
    return decoded


# For convenience and backwards compatibility w/ standard base64 module
body_decode = decode
decodestring = decode



def _unquote_match(match):
    """Turn a match in the form =AB to the ASCII character with value 0xab"""
    s = match.group(0)
    return unquote(s)


# Header decoding is done a bit differently
def header_decode(s):
    """Decode a string encoded with RFC 2045 MIME header `Q' encoding.

    This function does not parse a full MIME header value encoded with
    quoted-printable (like =?iso-8859-1?q?Hello_World?=) -- please use
    the high level email.header class for that functionality.
    """
    s = s.replace('_', ' ')
    return re.sub(r'=[a-fA-F0-9]{2}', _unquote_match, s, flags=re.ASCII)

external/cpython_master/Lib/email/quoprimime.py
DragonHPC/dragon
"""The idlelib package implements the Idle application.

Idle includes an interactive shell and editor.
Starting with Python 3.6, IDLE requires tcl/tk 8.5 or later.
Use the files named idle.* to start Idle.

The other files are private implementations.  Their details are subject to
change.  See PEP 434 for more.  Import them at your own risk.
"""
testing = False  # Set True by test.test_idle.

external/cpython_master/Lib/idlelib/__init__.py
DragonHPC/dragon
'''Complete the current word before the cursor with words in the editor.

Each menu selection or shortcut key selection replaces the word with a
different word with the same prefix. The search for matches begins
before the target and moves toward the top of the editor. It then starts
after the cursor and moves down. It then returns to the original word and
the cycle starts again.

Changing the current text line or leaving the cursor in a different
place before requesting the next selection causes AutoExpand to reset
its state.

There is only one instance of Autoexpand.
'''
import re
import string


class AutoExpand:
    wordchars = string.ascii_letters + string.digits + "_"

    def __init__(self, editwin):
        self.text = editwin.text
        self.bell = self.text.bell
        self.state = None

    def expand_word_event(self, event):
        "Replace the current word with the next expansion."
        curinsert = self.text.index("insert")
        curline = self.text.get("insert linestart", "insert lineend")
        if not self.state:
            words = self.getwords()
            index = 0
        else:
            words, index, insert, line = self.state
            if insert != curinsert or line != curline:
                words = self.getwords()
                index = 0
        if not words:
            self.bell()
            return "break"
        word = self.getprevword()
        self.text.delete("insert - %d chars" % len(word), "insert")
        newword = words[index]
        index = (index + 1) % len(words)
        if index == 0:
            self.bell()            # Warn we cycled around
        self.text.insert("insert", newword)
        curinsert = self.text.index("insert")
        curline = self.text.get("insert linestart", "insert lineend")
        self.state = words, index, curinsert, curline
        return "break"

    def getwords(self):
        "Return a list of words that match the prefix before the cursor."
        word = self.getprevword()
        if not word:
            return []
        before = self.text.get("1.0", "insert wordstart")
        wbefore = re.findall(r"\b" + word + r"\w+\b", before)
        del before
        after = self.text.get("insert wordend", "end")
        wafter = re.findall(r"\b" + word + r"\w+\b", after)
        del after
        if not wbefore and not wafter:
            return []
        words = []
        dict = {}
        # search backwards through words before
        wbefore.reverse()
        for w in wbefore:
            if dict.get(w):
                continue
            words.append(w)
            dict[w] = w
        # search onwards through words after
        for w in wafter:
            if dict.get(w):
                continue
            words.append(w)
            dict[w] = w
        words.append(word)
        return words

    def getprevword(self):
        "Return the word prefix before the cursor."
        line = self.text.get("insert linestart", "insert")
        i = len(line)
        while i > 0 and line[i-1] in self.wordchars:
            i = i-1
        return line[i:]


if __name__ == '__main__':
    from unittest import main
    main('idlelib.idle_test.test_autoexpand', verbosity=2)

external/cpython_master/Lib/idlelib/autoexpand.py
DragonHPC/dragon
"""About Dialog for IDLE

"""
import os
import sys
from platform import python_version, architecture

from tkinter import Toplevel, Frame, Label, Button, PhotoImage
from tkinter import SUNKEN, TOP, BOTTOM, LEFT, X, BOTH, W, EW, NSEW, E

from idlelib import textview


def build_bits():
    "Return bits for platform."
    if sys.platform == 'darwin':
        return '64' if sys.maxsize > 2**32 else '32'
    else:
        return architecture()[0][:2]


class AboutDialog(Toplevel):
    """Modal about dialog for idle

    """
    def __init__(self, parent, title=None, *, _htest=False, _utest=False):
        """Create popup, do not return until tk widget destroyed.

        parent - parent of this dialog
        title - string which is title of popup dialog
        _htest - bool, change box location when running htest
        _utest - bool, don't wait_window when running unittest
        """
        Toplevel.__init__(self, parent)
        self.configure(borderwidth=5)
        # place dialog below parent if running htest
        self.geometry("+%d+%d" % (
                        parent.winfo_rootx()+30,
                        parent.winfo_rooty()+(30 if not _htest else 100)))
        self.bg = "#bbbbbb"
        self.fg = "#000000"
        self.create_widgets()
        self.resizable(height=False, width=False)
        self.title(title or
                   f'About IDLE {python_version()} ({build_bits()} bit)')
        self.transient(parent)
        self.grab_set()
        self.protocol("WM_DELETE_WINDOW", self.ok)
        self.parent = parent
        self.button_ok.focus_set()
        self.bind('<Return>', self.ok)  # dismiss dialog
        self.bind('<Escape>', self.ok)  # dismiss dialog
        self._current_textview = None
        self._utest = _utest

        if not _utest:
            self.deiconify()
            self.wait_window()

    def create_widgets(self):
        frame = Frame(self, borderwidth=2, relief=SUNKEN)
        frame_buttons = Frame(self)
        frame_buttons.pack(side=BOTTOM, fill=X)
        frame.pack(side=TOP, expand=True, fill=BOTH)
        self.button_ok = Button(frame_buttons, text='Close',
                                command=self.ok)
        self.button_ok.pack(padx=5, pady=5)

        frame_background = Frame(frame, bg=self.bg)
        frame_background.pack(expand=True, fill=BOTH)

        header = Label(frame_background, text='IDLE', fg=self.fg,
                       bg=self.bg, font=('courier', 24, 'bold'))
        header.grid(row=0, column=0, sticky=E, padx=10, pady=10)

        tk_patchlevel = self.tk.call('info', 'patchlevel')
        ext = '.png' if tk_patchlevel >= '8.6' else '.gif'
        icon = os.path.join(os.path.abspath(os.path.dirname(__file__)),
                            'Icons', f'idle_48{ext}')
        self.icon_image = PhotoImage(master=self._root(), file=icon)
        logo = Label(frame_background, image=self.icon_image, bg=self.bg)
        logo.grid(row=0, column=0, sticky=W, rowspan=2, padx=10, pady=10)

        byline_text = "Python's Integrated Development\nand Learning Environment" + 5*'\n'
        byline = Label(frame_background, text=byline_text, justify=LEFT,
                       fg=self.fg, bg=self.bg)
        byline.grid(row=2, column=0, sticky=W, columnspan=3, padx=10, pady=5)
        email = Label(frame_background, text='email:  idle-dev@python.org',
                      justify=LEFT, fg=self.fg, bg=self.bg)
        email.grid(row=6, column=0, columnspan=2, sticky=W, padx=10, pady=0)
        docs = Label(frame_background, text='https://docs.python.org/' +
                     python_version()[:3] + '/library/idle.html',
                     justify=LEFT, fg=self.fg, bg=self.bg)
        docs.grid(row=7, column=0, columnspan=2, sticky=W, padx=10, pady=0)

        Frame(frame_background, borderwidth=1, relief=SUNKEN,
              height=2, bg=self.bg).grid(row=8, column=0, sticky=EW,
                                         columnspan=3, padx=5, pady=5)

        pyver = Label(frame_background,
                      text='Python version:  ' + python_version(),
                      fg=self.fg, bg=self.bg)
        pyver.grid(row=9, column=0, sticky=W, padx=10, pady=0)
        tkver = Label(frame_background, text='Tk version:  ' + tk_patchlevel,
                      fg=self.fg, bg=self.bg)
        tkver.grid(row=9, column=1, sticky=W, padx=2, pady=0)
        py_buttons = Frame(frame_background, bg=self.bg)
        py_buttons.grid(row=10, column=0, columnspan=2, sticky=NSEW)
        self.py_license = Button(py_buttons, text='License', width=8,
                                 highlightbackground=self.bg,
                                 command=self.show_py_license)
        self.py_license.pack(side=LEFT, padx=10, pady=10)
        self.py_copyright = Button(py_buttons, text='Copyright', width=8,
                                   highlightbackground=self.bg,
                                   command=self.show_py_copyright)
        self.py_copyright.pack(side=LEFT, padx=10, pady=10)
        self.py_credits = Button(py_buttons, text='Credits', width=8,
                                 highlightbackground=self.bg,
                                 command=self.show_py_credits)
        self.py_credits.pack(side=LEFT, padx=10, pady=10)

        Frame(frame_background, borderwidth=1, relief=SUNKEN,
              height=2, bg=self.bg).grid(row=11, column=0, sticky=EW,
                                         columnspan=3, padx=5, pady=5)

        idlever = Label(frame_background,
                        text='IDLE version:   ' + python_version(),
                        fg=self.fg, bg=self.bg)
        idlever.grid(row=12, column=0, sticky=W, padx=10, pady=0)
        idle_buttons = Frame(frame_background, bg=self.bg)
        idle_buttons.grid(row=13, column=0, columnspan=3, sticky=NSEW)
        self.readme = Button(idle_buttons, text='README', width=8,
                             highlightbackground=self.bg,
                             command=self.show_readme)
        self.readme.pack(side=LEFT, padx=10, pady=10)
        self.idle_news = Button(idle_buttons, text='NEWS', width=8,
                                highlightbackground=self.bg,
                                command=self.show_idle_news)
        self.idle_news.pack(side=LEFT, padx=10, pady=10)
        self.idle_credits = Button(idle_buttons, text='Credits', width=8,
                                   highlightbackground=self.bg,
                                   command=self.show_idle_credits)
        self.idle_credits.pack(side=LEFT, padx=10, pady=10)

    # License, copyright, and credits are of type _sitebuiltins._Printer
    def show_py_license(self):
        "Handle License button event."
        self.display_printer_text('About - License', license)

    def show_py_copyright(self):
        "Handle Copyright button event."
        self.display_printer_text('About - Copyright', copyright)

    def show_py_credits(self):
        "Handle Python Credits button event."
        self.display_printer_text('About - Python Credits', credits)

    # Encode CREDITS.txt to utf-8 for proper version of Loewis.
    # Specify others as ascii until need utf-8, so catch errors.
    def show_idle_credits(self):
        "Handle Idle Credits button event."
        self.display_file_text('About - Credits', 'CREDITS.txt', 'utf-8')

    def show_readme(self):
        "Handle Readme button event."
        self.display_file_text('About - Readme', 'README.txt', 'ascii')

    def show_idle_news(self):
        "Handle News button event."
        self.display_file_text('About - NEWS', 'NEWS.txt', 'utf-8')

    def display_printer_text(self, title, printer):
        """Create textview for built-in constants.

        Built-in constants have type _sitebuiltins._Printer.  The
        text is extracted from the built-in and then sent to a text
        viewer with self as the parent and title as the title of
        the popup.
        """
        printer._Printer__setup()
        text = '\n'.join(printer._Printer__lines)
        self._current_textview = textview.view_text(
            self, title, text, _utest=self._utest)

    def display_file_text(self, title, filename, encoding=None):
        """Create textview for filename.

        The filename needs to be in the current directory.  The path
        is sent to a text viewer with self as the parent, title as
        the title of the popup, and the file encoding.
        """
        fn = os.path.join(os.path.abspath(os.path.dirname(__file__)), filename)
        self._current_textview = textview.view_file(
            self, title, fn, encoding, _utest=self._utest)

    def ok(self, event=None):
        "Dismiss help_about dialog."
        self.grab_release()
        self.destroy()


if __name__ == '__main__':
    from unittest import main
    main('idlelib.idle_test.test_help_about', verbosity=2, exit=False)

    from idlelib.idle_test.htest import run
    run(AboutDialog)

external/cpython_master/Lib/idlelib/help_about.py
DragonHPC/dragon
"Test autocomplete_w, coverage 11%."

import unittest
from test.support import requires
from tkinter import Tk, Text

import idlelib.autocomplete_w as acw


class AutoCompleteWindowTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        requires('gui')
        cls.root = Tk()
        cls.root.withdraw()
        cls.text = Text(cls.root)
        cls.acw = acw.AutoCompleteWindow(cls.text)

    @classmethod
    def tearDownClass(cls):
        del cls.text, cls.acw
        cls.root.update_idletasks()
        cls.root.destroy()
        del cls.root

    def test_init(self):
        self.assertEqual(self.acw.widget, self.text)


if __name__ == '__main__':
    unittest.main(verbosity=2)

external/cpython_master/Lib/idlelib/idle_test/test_autocomplete_w.py
DragonHPC/dragon
"Test debugger, coverage 19%"

from idlelib import debugger
import unittest
from test.support import requires
requires('gui')
from tkinter import Tk


class NameSpaceTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.root = Tk()
        cls.root.withdraw()

    @classmethod
    def tearDownClass(cls):
        cls.root.destroy()
        del cls.root

    def test_init(self):
        debugger.NamespaceViewer(self.root, 'Test')


# Other classes are Idb, Debugger, and StackViewer.

if __name__ == '__main__':
    unittest.main(verbosity=2)

external/cpython_master/Lib/idlelib/idle_test/test_debugger.py
DragonHPC/dragon
'''Test (selected) IDLE Edit menu items.

Edit modules have their own test files
'''
from test.support import requires
requires('gui')
import tkinter as tk
from tkinter import ttk
import unittest
from idlelib import pyshell

class PasteTest(unittest.TestCase):
    '''Test pasting into widgets that allow pasting.

    On X11, replacing selections requires tk fix.
    '''
    @classmethod
    def setUpClass(cls):
        cls.root = root = tk.Tk()
        cls.root.withdraw()
        pyshell.fix_x11_paste(root)
        cls.text = tk.Text(root)
        cls.entry = tk.Entry(root)
        cls.tentry = ttk.Entry(root)
        cls.spin = tk.Spinbox(root)
        root.clipboard_clear()
        root.clipboard_append('two')

    @classmethod
    def tearDownClass(cls):
        del cls.text, cls.entry, cls.tentry
        cls.root.clipboard_clear()
        cls.root.update_idletasks()
        cls.root.destroy()
        del cls.root

    def test_paste_text(self):
        "Test pasting into text with and without a selection."
        text = self.text
        for tag, ans in ('', 'onetwo\n'), ('sel', 'two\n'):
            with self.subTest(tag=tag, ans=ans):
                text.delete('1.0', 'end')
                text.insert('1.0', 'one', tag)
                text.event_generate('<<Paste>>')
                self.assertEqual(text.get('1.0', 'end'), ans)

    def test_paste_entry(self):
        "Test pasting into an entry with and without a selection."
        # Generated <<Paste>> fails for tk entry without empty select
        # range for 'no selection'.  Live widget works fine.
        for entry in self.entry, self.tentry:
            for end, ans in (0, 'onetwo'), ('end', 'two'):
                with self.subTest(entry=entry, end=end, ans=ans):
                    entry.delete(0, 'end')
                    entry.insert(0, 'one')
                    entry.select_range(0, end)
                    entry.event_generate('<<Paste>>')
                    self.assertEqual(entry.get(), ans)

    def test_paste_spin(self):
        "Test pasting into a spinbox with and without a selection."
        # See note above for entry.
        spin = self.spin
        for end, ans in (0, 'onetwo'), ('end', 'two'):
            with self.subTest(end=end, ans=ans):
                spin.delete(0, 'end')
                spin.insert(0, 'one')
                spin.selection('range', 0, end)  # see note
                spin.event_generate('<<Paste>>')
                self.assertEqual(spin.get(), ans)


if __name__ == '__main__':
    unittest.main(verbosity=2)

external/cpython_master/Lib/idlelib/idle_test/test_editmenu.py
DragonHPC/dragon
"Test filelist, coverage 19%."

from idlelib import filelist
import unittest
from test.support import requires
from tkinter import Tk

class FileListTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        requires('gui')
        cls.root = Tk()
        cls.root.withdraw()

    @classmethod
    def tearDownClass(cls):
        cls.root.update_idletasks()
        for id in cls.root.tk.call('after', 'info'):
            cls.root.after_cancel(id)
        cls.root.destroy()
        del cls.root

    def test_new_empty(self):
        flist = filelist.FileList(self.root)
        self.assertEqual(flist.root, self.root)
        e = flist.new()
        self.assertEqual(type(e), flist.EditorWindow)
        e._close()


if __name__ == '__main__':
    unittest.main(verbosity=2)

external/cpython_master/Lib/idlelib/idle_test/test_filelist.py
DragonHPC/dragon
"Test undo, coverage 77%."
# Only test UndoDelegator so far.

from idlelib.undo import UndoDelegator
import unittest
from test.support import requires
requires('gui')

from unittest.mock import Mock
from tkinter import Text, Tk
from idlelib.percolator import Percolator


class UndoDelegatorTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.root = Tk()
        cls.text = Text(cls.root)
        cls.percolator = Percolator(cls.text)

    @classmethod
    def tearDownClass(cls):
        cls.percolator.redir.close()
        del cls.percolator, cls.text
        cls.root.destroy()
        del cls.root

    def setUp(self):
        self.delegator = UndoDelegator()
        self.delegator.bell = Mock()
        self.percolator.insertfilter(self.delegator)

    def tearDown(self):
        self.percolator.removefilter(self.delegator)
        self.text.delete('1.0', 'end')
        self.delegator.resetcache()

    def test_undo_event(self):
        text = self.text

        text.insert('insert', 'foobar')
        text.insert('insert', 'h')
        text.event_generate('<<undo>>')
        self.assertEqual(text.get('1.0', 'end'), '\n')

        text.insert('insert', 'foo')
        text.insert('insert', 'bar')
        text.delete('1.2', '1.4')
        text.insert('insert', 'hello')
        text.event_generate('<<undo>>')
        self.assertEqual(text.get('1.0', '1.4'), 'foar')
        text.event_generate('<<undo>>')
        self.assertEqual(text.get('1.0', '1.6'), 'foobar')
        text.event_generate('<<undo>>')
        self.assertEqual(text.get('1.0', '1.3'), 'foo')
        text.event_generate('<<undo>>')
        self.delegator.undo_event('event')
        self.assertTrue(self.delegator.bell.called)

    def test_redo_event(self):
        text = self.text

        text.insert('insert', 'foo')
        text.insert('insert', 'bar')
        text.delete('1.0', '1.3')
        text.event_generate('<<undo>>')
        text.event_generate('<<redo>>')
        self.assertEqual(text.get('1.0', '1.3'), 'bar')
        text.event_generate('<<redo>>')
        self.assertTrue(self.delegator.bell.called)

    def test_dump_event(self):
        """
        Dump_event cannot be tested directly without changing
        environment variables. So, test statements in dump_event
        indirectly
        """
        text = self.text
        d = self.delegator

        text.insert('insert', 'foo')
        text.insert('insert', 'bar')
        text.delete('1.2', '1.4')
        self.assertTupleEqual((d.pointer, d.can_merge), (3, True))
        text.event_generate('<<undo>>')
        self.assertTupleEqual((d.pointer, d.can_merge), (2, False))

    def test_get_set_saved(self):
        # test the getter method get_saved
        # test the setter method set_saved
        # indirectly test check_saved
        d = self.delegator

        self.assertTrue(d.get_saved())
        self.text.insert('insert', 'a')
        self.assertFalse(d.get_saved())
        d.saved_change_hook = Mock()

        d.set_saved(True)
        self.assertEqual(d.pointer, d.saved)
        self.assertTrue(d.saved_change_hook.called)

        d.set_saved(False)
        self.assertEqual(d.saved, -1)
        self.assertTrue(d.saved_change_hook.called)

    def test_undo_start_stop(self):
        # test the undo_block_start and undo_block_stop methods
        text = self.text

        text.insert('insert', 'foo')
        self.delegator.undo_block_start()
        text.insert('insert', 'bar')
        text.insert('insert', 'bar')
        self.delegator.undo_block_stop()
        self.assertEqual(text.get('1.0', '1.3'), 'foo')

        # test another code path
        self.delegator.undo_block_start()
        text.insert('insert', 'bar')
        self.delegator.undo_block_stop()
        self.assertEqual(text.get('1.0', '1.3'), 'foo')

    def test_addcmd(self):
        text = self.text
        # when number of undo operations exceeds max_undo
        self.delegator.max_undo = max_undo = 10
        for i in range(max_undo + 10):
            text.insert('insert', 'foo')
            self.assertLessEqual(len(self.delegator.undolist), max_undo)


if __name__ == '__main__':
    unittest.main(verbosity=2, exit=False)

external/cpython_master/Lib/idlelib/idle_test/test_undo.py
DragonHPC/dragon
"Test window, coverage 47%."

from idlelib import window
import unittest
from test.support import requires
from tkinter import Tk


class WindowListTest(unittest.TestCase):

    def test_init(self):
        wl = window.WindowList()
        self.assertEqual(wl.dict, {})
        self.assertEqual(wl.callbacks, [])

    # Further tests need mock Window.


class ListedToplevelTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        window.registry = set()
        requires('gui')
        cls.root = Tk()
        cls.root.withdraw()

    @classmethod
    def tearDownClass(cls):
        window.registry = window.WindowList()
        cls.root.update_idletasks()
##        for id in cls.root.tk.call('after', 'info'):
##            cls.root.after_cancel(id)  # Need for EditorWindow.
        cls.root.destroy()
        del cls.root

    def test_init(self):

        win = window.ListedToplevel(self.root)
        self.assertIn(win, window.registry)
        self.assertEqual(win.focused_widget, win)


if __name__ == '__main__':
    unittest.main(verbosity=2)

external/cpython_master/Lib/idlelib/idle_test/test_window.py
DragonHPC/dragon
"Test zoomheight, coverage 66%."
# Some code is system dependent.

from idlelib import zoomheight
import unittest
from test.support import requires
from tkinter import Tk
from idlelib.editor import EditorWindow


class Test(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        requires('gui')
        cls.root = Tk()
        cls.root.withdraw()
        cls.editwin = EditorWindow(root=cls.root)

    @classmethod
    def tearDownClass(cls):
        cls.editwin._close()
        cls.root.update_idletasks()
        for id in cls.root.tk.call('after', 'info'):
            cls.root.after_cancel(id)  # Need for EditorWindow.
        cls.root.destroy()
        del cls.root

    def test_init(self):
        zoom = zoomheight.ZoomHeight(self.editwin)
        self.assertIs(zoom.editwin, self.editwin)

    def test_zoom_height_event(self):
        zoom = zoomheight.ZoomHeight(self.editwin)
        zoom.zoom_height_event()


if __name__ == '__main__':
    unittest.main(verbosity=2)

external/cpython_master/Lib/idlelib/idle_test/test_zoomheight.py
DragonHPC/dragon
"""
MultiCall - a class which inherits its methods from a Tkinter widget (Text, for
example), but enables multiple calls of functions per virtual event - all
matching events will be called, not only the most specific one. This is done
by wrapping the event functions - event_add, event_delete and event_info.
MultiCall recognizes only a subset of legal event sequences. Sequences which
are not recognized are treated by the original Tk handling mechanism. A
more-specific event will be called before a less-specific event.

The recognized sequences are complete one-event sequences (no emacs-style
Ctrl-X Ctrl-C, no shortcuts like <3>), for all types of events.
Key/Button Press/Release events can have modifiers.
The recognized modifiers are Shift, Control, Option and Command for Mac, and
Control, Alt, Shift, Meta/M for other platforms.

For all events which were handled by MultiCall, a new member is added to the
event instance passed to the binded functions - mc_type. This is one of the
event type constants defined in this module (such as MC_KEYPRESS).
For Key/Button events (which are handled by MultiCall and may receive
modifiers), another member is added - mc_state. This member gives the state
of the recognized modifiers, as a combination of the modifier constants
also defined in this module (for example, MC_SHIFT).
Using these members is absolutely portable.

The order by which events are called is defined by these rules:
1. A more-specific event will be called before a less-specific event.
2. A recently-binded event will be called before a previously-binded event,
   unless this conflicts with the first rule.
Each function will be called at most once for each event.
"""
import re
import sys

import tkinter

# the event type constants, which define the meaning of mc_type
MC_KEYPRESS=0; MC_KEYRELEASE=1; MC_BUTTONPRESS=2; MC_BUTTONRELEASE=3;
MC_ACTIVATE=4; MC_CIRCULATE=5; MC_COLORMAP=6; MC_CONFIGURE=7;
MC_DEACTIVATE=8; MC_DESTROY=9; MC_ENTER=10; MC_EXPOSE=11; MC_FOCUSIN=12;
MC_FOCUSOUT=13; MC_GRAVITY=14; MC_LEAVE=15; MC_MAP=16; MC_MOTION=17;
MC_MOUSEWHEEL=18; MC_PROPERTY=19; MC_REPARENT=20; MC_UNMAP=21; MC_VISIBILITY=22;
# the modifier state constants, which define the meaning of mc_state
MC_SHIFT = 1<<0; MC_CONTROL = 1<<2; MC_ALT = 1<<3; MC_META = 1<<5
MC_OPTION = 1<<6; MC_COMMAND = 1<<7

# define the list of modifiers, to be used in complex event types.
if sys.platform == "darwin":
    _modifiers = (("Shift",), ("Control",), ("Option",), ("Command",))
    _modifier_masks = (MC_SHIFT, MC_CONTROL, MC_OPTION, MC_COMMAND)
else:
    _modifiers = (("Control",), ("Alt",), ("Shift",), ("Meta", "M"))
    _modifier_masks = (MC_CONTROL, MC_ALT, MC_SHIFT, MC_META)

# a dictionary to map a modifier name into its number
_modifier_names = dict([(name, number)
                         for number in range(len(_modifiers))
                         for name in _modifiers[number]])

# In 3.4, if no shell window is ever open, the underlying Tk widget is
# destroyed before .__del__ methods here are called.  The following
# is used to selectively ignore shutdown exceptions to avoid
# 'Exception ignored' messages.  See http://bugs.python.org/issue20167
APPLICATION_GONE = "application has been destroyed"

# A binder is a class which binds functions to one type of event. It has two
# methods: bind and unbind, which get a function and a parsed sequence, as
# returned by _parse_sequence(). There are two types of binders:
# _SimpleBinder handles event types with no modifiers and no detail.
# No Python functions are called when no events are binded.
# _ComplexBinder handles event types with modifiers and a detail.
# A Python function is called each time an event is generated.

class _SimpleBinder:
    def __init__(self, type, widget, widgetinst):
        self.type = type
        self.sequence = '<'+_types[type][0]+'>'
        self.widget = widget
        self.widgetinst = widgetinst
        self.bindedfuncs = []
        self.handlerid = None

    def bind(self, triplet, func):
        if not self.handlerid:
            def handler(event, l = self.bindedfuncs, mc_type = self.type):
                event.mc_type = mc_type
                wascalled = {}
                for i in range(len(l)-1, -1, -1):
                    func = l[i]
                    if func not in wascalled:
                        wascalled[func] = True
                        r = func(event)
                        if r:
                            return r
            self.handlerid = self.widget.bind(self.widgetinst,
                                              self.sequence, handler)
        self.bindedfuncs.append(func)

    def unbind(self, triplet, func):
        self.bindedfuncs.remove(func)
        if not self.bindedfuncs:
            self.widget.unbind(self.widgetinst, self.sequence, self.handlerid)
            self.handlerid = None

    def __del__(self):
        if self.handlerid:
            try:
                self.widget.unbind(self.widgetinst, self.sequence,
                        self.handlerid)
            except tkinter.TclError as e:
                if not APPLICATION_GONE in e.args[0]:
                    raise

# An int in range(1 << len(_modifiers)) represents a combination of modifiers
# (if the least significant bit is on, _modifiers[0] is on, and so on).
# _state_subsets gives for each combination of modifiers, or *state*,
# a list of the states which are a subset of it. This list is ordered by the
# number of modifiers is the state - the most specific state comes first.
_states = range(1 << len(_modifiers))
_state_names = [''.join(m[0]+'-'
                        for i, m in enumerate(_modifiers)
                        if (1 << i) & s)
                for s in _states]

def expand_substates(states):
    '''For each item of states return a list containing all combinations of
    that item with individual bits reset, sorted by the number of set bits.
    '''
    def nbits(n):
        "number of bits set in n base 2"
        nb = 0
        while n:
            n, rem = divmod(n, 2)
            nb += rem
        return nb
    statelist = []
    for state in states:
        substates = list(set(state & x for x in states))
        substates.sort(key=nbits, reverse=True)
        statelist.append(substates)
    return statelist

_state_subsets = expand_substates(_states)

# _state_codes gives for each state, the portable code to be passed as mc_state
_state_codes = []
for s in _states:
    r = 0
    for i in range(len(_modifiers)):
        if (1 << i) & s:
            r |= _modifier_masks[i]
    _state_codes.append(r)

class _ComplexBinder:
    # This class binds many functions, and only unbinds them when it is deleted.
    # self.handlerids is the list of seqs and ids of binded handler functions.
    # The binded functions sit in a dictionary of lists of lists, which maps
    # a detail (or None) and a state into a list of functions.
    # When a new detail is discovered, handlers for all the possible states
    # are binded.

    def __create_handler(self, lists, mc_type, mc_state):
        def handler(event, lists = lists,
                    mc_type = mc_type, mc_state = mc_state,
                    ishandlerrunning = self.ishandlerrunning,
                    doafterhandler = self.doafterhandler):
            ishandlerrunning[:] = [True]
            event.mc_type = mc_type
            event.mc_state = mc_state
            wascalled = {}
            r = None
            for l in lists:
                for i in range(len(l)-1, -1, -1):
                    func = l[i]
                    if func not in wascalled:
                        wascalled[func] = True
                        r = l[i](event)
                        if r:
                            break
                if r:
                    break
            ishandlerrunning[:] = []
            # Call all functions in doafterhandler and remove them from list
            for f in doafterhandler:
                f()
            doafterhandler[:] = []
            if r:
                return r
        return handler

    def __init__(self, type, widget, widgetinst):
        self.type = type
        self.typename = _types[type][0]
        self.widget = widget
        self.widgetinst = widgetinst
        self.bindedfuncs = {None: [[] for s in _states]}
        self.handlerids = []
        # we don't want to change the lists of functions while a handler is
        # running - it will mess up the loop and anyway, we usually want the
        # change to happen from the next event. So we have a list of functions
        # for the handler to run after it finishes calling the binded functions.
        # It calls them only once.
        # ishandlerrunning is a list. An empty one means no, otherwise - yes.
        # this is done so that it would be mutable.
        self.ishandlerrunning = []
        self.doafterhandler = []
        for s in _states:
            lists = [self.bindedfuncs[None][i] for i in _state_subsets[s]]
            handler = self.__create_handler(lists, type, _state_codes[s])
            seq = '<'+_state_names[s]+self.typename+'>'
            self.handlerids.append((seq, self.widget.bind(self.widgetinst,
                                                          seq, handler)))

    def bind(self, triplet, func):
        if triplet[2] not in self.bindedfuncs:
            self.bindedfuncs[triplet[2]] = [[] for s in _states]
            for s in _states:
                lists = [ self.bindedfuncs[detail][i]
                          for detail in (triplet[2], None)
                          for i in _state_subsets[s]       ]
                handler = self.__create_handler(lists, self.type,
                                                _state_codes[s])
                seq = "<%s%s-%s>"% (_state_names[s], self.typename, triplet[2])
                self.handlerids.append((seq, self.widget.bind(self.widgetinst,
                                                              seq, handler)))
        doit = lambda: self.bindedfuncs[triplet[2]][triplet[0]].append(func)
        if not self.ishandlerrunning:
            doit()
        else:
            self.doafterhandler.append(doit)

    def unbind(self, triplet, func):
        doit = lambda: self.bindedfuncs[triplet[2]][triplet[0]].remove(func)
        if not self.ishandlerrunning:
            doit()
        else:
            self.doafterhandler.append(doit)

    def __del__(self):
        for seq, id in self.handlerids:
            try:
                self.widget.unbind(self.widgetinst, seq, id)
            except tkinter.TclError as e:
                if not APPLICATION_GONE in e.args[0]:
                    raise

# define the list of event types to be handled by MultiEvent. the order is
# compatible with the definition of event type constants.
_types = (
    ("KeyPress", "Key"), ("KeyRelease",), ("ButtonPress", "Button"),
    ("ButtonRelease",), ("Activate",), ("Circulate",), ("Colormap",),
    ("Configure",), ("Deactivate",), ("Destroy",), ("Enter",), ("Expose",),
    ("FocusIn",), ("FocusOut",), ("Gravity",), ("Leave",), ("Map",),
    ("Motion",), ("MouseWheel",), ("Property",), ("Reparent",), ("Unmap",),
    ("Visibility",),
)

# which binder should be used for every event type?
_binder_classes = (_ComplexBinder,) * 4 + (_SimpleBinder,) * (len(_types)-4)

# A dictionary to map a type name into its number
_type_names = dict([(name, number)
                     for number in range(len(_types))
                     for name in _types[number]])

_keysym_re = re.compile(r"^\w+$")
_button_re = re.compile(r"^[1-5]$")
def _parse_sequence(sequence):
    """Get a string which should describe an event sequence. If it is
    successfully parsed as one, return a tuple containing the state (as an int),
    the event type (as an index of _types), and the detail - None if none, or a
    string if there is one. If the parsing is unsuccessful, return None.
    """
    if not sequence or sequence[0] != '<' or sequence[-1] != '>':
        return None
    words = sequence[1:-1].split('-')
    modifiers = 0
    while words and words[0] in _modifier_names:
        modifiers |= 1 << _modifier_names[words[0]]
        del words[0]
    if words and words[0] in _type_names:
        type = _type_names[words[0]]
        del words[0]
    else:
        return None
    if _binder_classes[type] is _SimpleBinder:
        if modifiers or words:
            return None
        else:
            detail = None
    else:
        # _ComplexBinder
        if type in [_type_names[s] for s in ("KeyPress", "KeyRelease")]:
            type_re = _keysym_re
        else:
            type_re = _button_re

        if not words:
            detail = None
        elif len(words) == 1 and type_re.match(words[0]):
            detail = words[0]
        else:
            return None

    return modifiers, type, detail

def _triplet_to_sequence(triplet):
    if triplet[2]:
        return '<'+_state_names[triplet[0]]+_types[triplet[1]][0]+'-'+ \
               triplet[2]+'>'
    else:
        return '<'+_state_names[triplet[0]]+_types[triplet[1]][0]+'>'

_multicall_dict = {}
def MultiCallCreator(widget):
    """Return a MultiCall class which inherits its methods from the
    given widget class (for example, Tkinter.Text). This is used
    instead of a templating mechanism.
    """
    if widget in _multicall_dict:
        return _multicall_dict[widget]

    class MultiCall (widget):
        assert issubclass(widget, tkinter.Misc)

        def __init__(self, *args, **kwargs):
            widget.__init__(self, *args, **kwargs)
            # a dictionary which maps a virtual event to a tuple with:
            #  0. the function binded
            #  1. a list of triplets - the sequences it is binded to
            self.__eventinfo = {}
            self.__binders = [_binder_classes[i](i, widget, self)
                              for i in range(len(_types))]

        def bind(self, sequence=None, func=None, add=None):
            #print("bind(%s, %s, %s)" % (sequence, func, add),
            #      file=sys.__stderr__)
            if type(sequence) is str and len(sequence) > 2 and \
               sequence[:2] == "<<" and sequence[-2:] == ">>":
                if sequence in self.__eventinfo:
                    ei = self.__eventinfo[sequence]
                    if ei[0] is not None:
                        for triplet in ei[1]:
                            self.__binders[triplet[1]].unbind(triplet, ei[0])
                    ei[0] = func
                    if ei[0] is not None:
                        for triplet in ei[1]:
                            self.__binders[triplet[1]].bind(triplet, func)
                else:
                    self.__eventinfo[sequence] = [func, []]
            return widget.bind(self, sequence, func, add)

        def unbind(self, sequence, funcid=None):
            if type(sequence) is str and len(sequence) > 2 and \
               sequence[:2] == "<<" and sequence[-2:] == ">>" and \
               sequence in self.__eventinfo:
                func, triplets = self.__eventinfo[sequence]
                if func is not None:
                    for triplet in triplets:
                        self.__binders[triplet[1]].unbind(triplet, func)
                    self.__eventinfo[sequence][0] = None
            return widget.unbind(self, sequence, funcid)

        def event_add(self, virtual, *sequences):
            #print("event_add(%s, %s)" % (repr(virtual), repr(sequences)),
            #      file=sys.__stderr__)
            if virtual not in self.__eventinfo:
                self.__eventinfo[virtual] = [None, []]

            func, triplets = self.__eventinfo[virtual]
            for seq in sequences:
                triplet = _parse_sequence(seq)
                if triplet is None:
                    #print("Tkinter event_add(%s)" % seq, file=sys.__stderr__)
                    widget.event_add(self, virtual, seq)
                else:
                    if func is not None:
                        self.__binders[triplet[1]].bind(triplet, func)
                    triplets.append(triplet)

        def event_delete(self, virtual, *sequences):
            if virtual not in self.__eventinfo:
                return
            func, triplets = self.__eventinfo[virtual]
            for seq in sequences:
                triplet = _parse_sequence(seq)
                if triplet is None:
                    #print("Tkinter event_delete: %s" % seq, file=sys.__stderr__)
                    widget.event_delete(self, virtual, seq)
                else:
                    if func is not None:
                        self.__binders[triplet[1]].unbind(triplet, func)
                    triplets.remove(triplet)

        def event_info(self, virtual=None):
            if virtual is None or virtual not in self.__eventinfo:
                return widget.event_info(self, virtual)
            else:
                return tuple(map(_triplet_to_sequence,
                                 self.__eventinfo[virtual][1])) + \
                       widget.event_info(self, virtual)

        def __del__(self):
            for virtual in self.__eventinfo:
                func, triplets = self.__eventinfo[virtual]
                if func:
                    for triplet in triplets:
                        try:
                            self.__binders[triplet[1]].unbind(triplet, func)
                        except tkinter.TclError as e:
                            if not APPLICATION_GONE in e.args[0]:
                                raise

    _multicall_dict[widget] = MultiCall
    return MultiCall


def _multi_call(parent):  # htest #
    top = tkinter.Toplevel(parent)
    top.title("Test MultiCall")
    x, y = map(int, parent.geometry().split('+')[1:])
    top.geometry("+%d+%d" % (x, y + 175))
    text = MultiCallCreator(tkinter.Text)(top)
    text.pack()
    def bindseq(seq, n=[0]):
        def handler(event):
            print(seq)
        text.bind("<<handler%d>>"%n[0], handler)
        text.event_add("<<handler%d>>"%n[0], seq)
        n[0] += 1
    bindseq("<Key>")
    bindseq("<Control-Key>")
    bindseq("<Alt-Key-a>")
    bindseq("<Control-Key-a>")
    bindseq("<Alt-Control-Key-a>")
    bindseq("<Key-b>")
    bindseq("<Control-Button-1>")
    bindseq("<Button-2>")
    bindseq("<Alt-Button-1>")
    bindseq("<FocusOut>")
    bindseq("<Enter>")
    bindseq("<Leave>")

if __name__ == "__main__":
    from unittest import main
    main('idlelib.idle_test.test_mainmenu', verbosity=2, exit=False)

    from idlelib.idle_test.htest import run
    run(_multi_call)

external/cpython_master/Lib/idlelib/multicall.py
DragonHPC/dragon
import linecache
import os
import sys

import tkinter as tk

from idlelib.debugobj import ObjectTreeItem, make_objecttreeitem
from idlelib.tree import TreeNode, TreeItem, ScrolledCanvas

def StackBrowser(root, flist=None, tb=None, top=None):
    global sc, item, node  # For testing.
    if top is None:
        top = tk.Toplevel(root)
    sc = ScrolledCanvas(top, bg="white", highlightthickness=0)
    sc.frame.pack(expand=1, fill="both")
    item = StackTreeItem(flist, tb)
    node = TreeNode(sc.canvas, None, item)
    node.expand()


class StackTreeItem(TreeItem):

    def __init__(self, flist=None, tb=None):
        self.flist = flist
        self.stack = self.get_stack(tb)
        self.text = self.get_exception()

    def get_stack(self, tb):
        if tb is None:
            tb = sys.last_traceback
        stack = []
        if tb and tb.tb_frame is None:
            tb = tb.tb_next
        while tb is not None:
            stack.append((tb.tb_frame, tb.tb_lineno))
            tb = tb.tb_next
        return stack

    def get_exception(self):
        type = sys.last_type
        value = sys.last_value
        if hasattr(type, "__name__"):
            type = type.__name__
        s = str(type)
        if value is not None:
            s = s + ": " + str(value)
        return s

    def GetText(self):
        return self.text

    def GetSubList(self):
        sublist = []
        for info in self.stack:
            item = FrameTreeItem(info, self.flist)
            sublist.append(item)
        return sublist


class FrameTreeItem(TreeItem):

    def __init__(self, info, flist):
        self.info = info
        self.flist = flist

    def GetText(self):
        frame, lineno = self.info
        try:
            modname = frame.f_globals["__name__"]
        except:
            modname = "?"
        code = frame.f_code
        filename = code.co_filename
        funcname = code.co_name
        sourceline = linecache.getline(filename, lineno)
        sourceline = sourceline.strip()
        if funcname in ("?", "", None):
            item = "%s, line %d: %s" % (modname, lineno, sourceline)
        else:
            item = "%s.%s(...), line %d: %s" % (modname, funcname,
                                             lineno, sourceline)
        return item

    def GetSubList(self):
        frame, lineno = self.info
        sublist = []
        if frame.f_globals is not frame.f_locals:
            item = VariablesTreeItem("<locals>", frame.f_locals, self.flist)
            sublist.append(item)
        item = VariablesTreeItem("<globals>", frame.f_globals, self.flist)
        sublist.append(item)
        return sublist

    def OnDoubleClick(self):
        if self.flist:
            frame, lineno = self.info
            filename = frame.f_code.co_filename
            if os.path.isfile(filename):
                self.flist.gotofileline(filename, lineno)


class VariablesTreeItem(ObjectTreeItem):

    def GetText(self):
        return self.labeltext

    def GetLabelText(self):
        return None

    def IsExpandable(self):
        return len(self.object) > 0

    def GetSubList(self):
        sublist = []
        for key in self.object.keys():
            try:
                value = self.object[key]
            except KeyError:
                continue
            def setfunction(value, key=key, object=self.object):
                object[key] = value
            item = make_objecttreeitem(key + " =", value, setfunction)
            sublist.append(item)
        return sublist


def _stack_viewer(parent):  # htest #
    from idlelib.pyshell import PyShellFileList
    top = tk.Toplevel(parent)
    top.title("Test StackViewer")
    x, y = map(int, parent.geometry().split('+')[1:])
    top.geometry("+%d+%d" % (x + 50, y + 175))
    flist = PyShellFileList(top)
    try: # to obtain a traceback object
        intentional_name_error
    except NameError:
        exc_type, exc_value, exc_tb = sys.exc_info()
    # inject stack trace to sys
    sys.last_type = exc_type
    sys.last_value = exc_value
    sys.last_traceback = exc_tb

    StackBrowser(top, flist=flist, top=top, tb=exc_tb)

    # restore sys to original state
    del sys.last_type
    del sys.last_value
    del sys.last_traceback

if __name__ == '__main__':
    from unittest import main
    main('idlelib.idle_test.test_stackviewer', verbosity=2, exit=False)

    from idlelib.idle_test.htest import run
    run(_stack_viewer)

external/cpython_master/Lib/idlelib/stackviewer.py
DragonHPC/dragon
"""This module provides the components needed to build your own __import__
function.  Undocumented functions are obsolete.

In most cases it is preferred you consider using the importlib module's
functionality over this module.

"""
# (Probably) need to stay in _imp
from _imp import (lock_held, acquire_lock, release_lock,
                  get_frozen_object, is_frozen_package,
                  init_frozen, is_builtin, is_frozen,
                  _fix_co_filename)
try:
    from _imp import create_dynamic
except ImportError:
    # Platform doesn't support dynamic loading.
    create_dynamic = None

from importlib._bootstrap import _ERR_MSG, _exec, _load, _builtin_from_name
from importlib._bootstrap_external import SourcelessFileLoader

from importlib import machinery
from importlib import util
import importlib
import os
import sys
import tokenize
import types
import warnings

warnings.warn("the imp module is deprecated in favour of importlib; "
              "see the module's documentation for alternative uses",
              DeprecationWarning, stacklevel=2)

# DEPRECATED
SEARCH_ERROR = 0
PY_SOURCE = 1
PY_COMPILED = 2
C_EXTENSION = 3
PY_RESOURCE = 4
PKG_DIRECTORY = 5
C_BUILTIN = 6
PY_FROZEN = 7
PY_CODERESOURCE = 8
IMP_HOOK = 9


def new_module(name):
    """**DEPRECATED**

    Create a new module.

    The module is not entered into sys.modules.

    """
    return types.ModuleType(name)


def get_magic():
    """**DEPRECATED**

    Return the magic number for .pyc files.
    """
    return util.MAGIC_NUMBER


def get_tag():
    """Return the magic tag for .pyc files."""
    return sys.implementation.cache_tag


def cache_from_source(path, debug_override=None):
    """**DEPRECATED**

    Given the path to a .py file, return the path to its .pyc file.

    The .py file does not need to exist; this simply returns the path to the
    .pyc file calculated as if the .py file were imported.

    If debug_override is not None, then it must be a boolean and is used in
    place of sys.flags.optimize.

    If sys.implementation.cache_tag is None then NotImplementedError is raised.

    """
    with warnings.catch_warnings():
        warnings.simplefilter('ignore')
        return util.cache_from_source(path, debug_override)


def source_from_cache(path):
    """**DEPRECATED**

    Given the path to a .pyc. file, return the path to its .py file.

    The .pyc file does not need to exist; this simply returns the path to
    the .py file calculated to correspond to the .pyc file.  If path does
    not conform to PEP 3147 format, ValueError will be raised. If
    sys.implementation.cache_tag is None then NotImplementedError is raised.

    """
    return util.source_from_cache(path)


def get_suffixes():
    """**DEPRECATED**"""
    extensions = [(s, 'rb', C_EXTENSION) for s in machinery.EXTENSION_SUFFIXES]
    source = [(s, 'r', PY_SOURCE) for s in machinery.SOURCE_SUFFIXES]
    bytecode = [(s, 'rb', PY_COMPILED) for s in machinery.BYTECODE_SUFFIXES]

    return extensions + source + bytecode


class NullImporter:

    """**DEPRECATED**

    Null import object.

    """

    def __init__(self, path):
        if path == '':
            raise ImportError('empty pathname', path='')
        elif os.path.isdir(path):
            raise ImportError('existing directory', path=path)

    def find_module(self, fullname):
        """Always returns None."""
        return None


class _HackedGetData:

    """Compatibility support for 'file' arguments of various load_*()
    functions."""

    def __init__(self, fullname, path, file=None):
        super().__init__(fullname, path)
        self.file = file

    def get_data(self, path):
        """Gross hack to contort loader to deal w/ load_*()'s bad API."""
        if self.file and path == self.path:
            # The contract of get_data() requires us to return bytes. Reopen the
            # file in binary mode if needed.
            if not self.file.closed:
                file = self.file
                if 'b' not in file.mode:
                    file.close()
            if self.file.closed:
                self.file = file = open(self.path, 'rb')

            with file:
                return file.read()
        else:
            return super().get_data(path)


class _LoadSourceCompatibility(_HackedGetData, machinery.SourceFileLoader):

    """Compatibility support for implementing load_source()."""


def load_source(name, pathname, file=None):
    loader = _LoadSourceCompatibility(name, pathname, file)
    spec = util.spec_from_file_location(name, pathname, loader=loader)
    if name in sys.modules:
        module = _exec(spec, sys.modules[name])
    else:
        module = _load(spec)
    # To allow reloading to potentially work, use a non-hacked loader which
    # won't rely on a now-closed file object.
    module.__loader__ = machinery.SourceFileLoader(name, pathname)
    module.__spec__.loader = module.__loader__
    return module


class _LoadCompiledCompatibility(_HackedGetData, SourcelessFileLoader):

    """Compatibility support for implementing load_compiled()."""


def load_compiled(name, pathname, file=None):
    """**DEPRECATED**"""
    loader = _LoadCompiledCompatibility(name, pathname, file)
    spec = util.spec_from_file_location(name, pathname, loader=loader)
    if name in sys.modules:
        module = _exec(spec, sys.modules[name])
    else:
        module = _load(spec)
    # To allow reloading to potentially work, use a non-hacked loader which
    # won't rely on a now-closed file object.
    module.__loader__ = SourcelessFileLoader(name, pathname)
    module.__spec__.loader = module.__loader__
    return module


def load_package(name, path):
    """**DEPRECATED**"""
    if os.path.isdir(path):
        extensions = (machinery.SOURCE_SUFFIXES[:] +
                      machinery.BYTECODE_SUFFIXES[:])
        for extension in extensions:
            init_path = os.path.join(path, '__init__' + extension)
            if os.path.exists(init_path):
                path = init_path
                break
        else:
            raise ValueError('{!r} is not a package'.format(path))
    spec = util.spec_from_file_location(name, path,
                                        submodule_search_locations=[])
    if name in sys.modules:
        return _exec(spec, sys.modules[name])
    else:
        return _load(spec)


def load_module(name, file, filename, details):
    """**DEPRECATED**

    Load a module, given information returned by find_module().

    The module name must include the full package name, if any.

    """
    suffix, mode, type_ = details
    if mode and (not mode.startswith(('r', 'U')) or '+' in mode):
        raise ValueError('invalid file open mode {!r}'.format(mode))
    elif file is None and type_ in {PY_SOURCE, PY_COMPILED}:
        msg = 'file object required for import (type code {})'.format(type_)
        raise ValueError(msg)
    elif type_ == PY_SOURCE:
        return load_source(name, filename, file)
    elif type_ == PY_COMPILED:
        return load_compiled(name, filename, file)
    elif type_ == C_EXTENSION and load_dynamic is not None:
        if file is None:
            with open(filename, 'rb') as opened_file:
                return load_dynamic(name, filename, opened_file)
        else:
            return load_dynamic(name, filename, file)
    elif type_ == PKG_DIRECTORY:
        return load_package(name, filename)
    elif type_ == C_BUILTIN:
        return init_builtin(name)
    elif type_ == PY_FROZEN:
        return init_frozen(name)
    else:
        msg =  "Don't know how to import {} (type code {})".format(name, type_)
        raise ImportError(msg, name=name)


def find_module(name, path=None):
    """**DEPRECATED**

    Search for a module.

    If path is omitted or None, search for a built-in, frozen or special
    module and continue search in sys.path. The module name cannot
    contain '.'; to search for a submodule of a package, pass the
    submodule name and the package's __path__.

    """
    if not isinstance(name, str):
        raise TypeError("'name' must be a str, not {}".format(type(name)))
    elif not isinstance(path, (type(None), list)):
        # Backwards-compatibility
        raise RuntimeError("'path' must be None or a list, "
                           "not {}".format(type(path)))

    if path is None:
        if is_builtin(name):
            return None, None, ('', '', C_BUILTIN)
        elif is_frozen(name):
            return None, None, ('', '', PY_FROZEN)
        else:
            path = sys.path

    for entry in path:
        package_directory = os.path.join(entry, name)
        for suffix in ['.py', machinery.BYTECODE_SUFFIXES[0]]:
            package_file_name = '__init__' + suffix
            file_path = os.path.join(package_directory, package_file_name)
            if os.path.isfile(file_path):
                return None, package_directory, ('', '', PKG_DIRECTORY)
        for suffix, mode, type_ in get_suffixes():
            file_name = name + suffix
            file_path = os.path.join(entry, file_name)
            if os.path.isfile(file_path):
                break
        else:
            continue
        break  # Break out of outer loop when breaking out of inner loop.
    else:
        raise ImportError(_ERR_MSG.format(name), name=name)

    encoding = None
    if 'b' not in mode:
        with open(file_path, 'rb') as file:
            encoding = tokenize.detect_encoding(file.readline)[0]
    file = open(file_path, mode, encoding=encoding)
    return file, file_path, (suffix, mode, type_)


def reload(module):
    """**DEPRECATED**

    Reload the module and return it.

    The module must have been successfully imported before.

    """
    return importlib.reload(module)


def init_builtin(name):
    """**DEPRECATED**

    Load and return a built-in module by name, or None is such module doesn't
    exist
    """
    try:
        return _builtin_from_name(name)
    except ImportError:
        return None


if create_dynamic:
    def load_dynamic(name, path, file=None):
        """**DEPRECATED**

        Load an extension module.
        """
        import importlib.machinery
        loader = importlib.machinery.ExtensionFileLoader(name, path)

        # Issue #24748: Skip the sys.modules check in _load_module_shim;
        # always load new extension
        spec = importlib.machinery.ModuleSpec(
            name=name, loader=loader, origin=path)
        return _load(spec)

else:
    load_dynamic = None

external/cpython_master/Lib/imp.py
DragonHPC/dragon
"""Implementation of JSONDecoder
"""
import re

from json import scanner
try:
    from _json import scanstring as c_scanstring
except ImportError:
    c_scanstring = None

__all__ = ['JSONDecoder', 'JSONDecodeError']

FLAGS = re.VERBOSE | re.MULTILINE | re.DOTALL

NaN = float('nan')
PosInf = float('inf')
NegInf = float('-inf')


class JSONDecodeError(ValueError):
    """Subclass of ValueError with the following additional properties:

    msg: The unformatted error message
    doc: The JSON document being parsed
    pos: The start index of doc where parsing failed
    lineno: The line corresponding to pos
    colno: The column corresponding to pos

    """
    # Note that this exception is used from _json
    def __init__(self, msg, doc, pos):
        lineno = doc.count('\n', 0, pos) + 1
        colno = pos - doc.rfind('\n', 0, pos)
        errmsg = '%s: line %d column %d (char %d)' % (msg, lineno, colno, pos)
        ValueError.__init__(self, errmsg)
        self.msg = msg
        self.doc = doc
        self.pos = pos
        self.lineno = lineno
        self.colno = colno

    def __reduce__(self):
        return self.__class__, (self.msg, self.doc, self.pos)


_CONSTANTS = {
    '-Infinity': NegInf,
    'Infinity': PosInf,
    'NaN': NaN,
}


STRINGCHUNK = re.compile(r'(.*?)(["\\\x00-\x1f])', FLAGS)
BACKSLASH = {
    '"': '"', '\\': '\\', '/': '/',
    'b': '\b', 'f': '\f', 'n': '\n', 'r': '\r', 't': '\t',
}

def _decode_uXXXX(s, pos):
    esc = s[pos + 1:pos + 5]
    if len(esc) == 4 and esc[1] not in 'xX':
        try:
            return int(esc, 16)
        except ValueError:
            pass
    msg = "Invalid \\uXXXX escape"
    raise JSONDecodeError(msg, s, pos)

def py_scanstring(s, end, strict=True,
        _b=BACKSLASH, _m=STRINGCHUNK.match):
    """Scan the string s for a JSON string. End is the index of the
    character in s after the quote that started the JSON string.
    Unescapes all valid JSON string escape sequences and raises ValueError
    on attempt to decode an invalid string. If strict is False then literal
    control characters are allowed in the string.

    Returns a tuple of the decoded string and the index of the character in s
    after the end quote."""
    chunks = []
    _append = chunks.append
    begin = end - 1
    while 1:
        chunk = _m(s, end)
        if chunk is None:
            raise JSONDecodeError("Unterminated string starting at", s, begin)
        end = chunk.end()
        content, terminator = chunk.groups()
        # Content is contains zero or more unescaped string characters
        if content:
            _append(content)
        # Terminator is the end of string, a literal control character,
        # or a backslash denoting that an escape sequence follows
        if terminator == '"':
            break
        elif terminator != '\\':
            if strict:
                #msg = "Invalid control character %r at" % (terminator,)
                msg = "Invalid control character {0!r} at".format(terminator)
                raise JSONDecodeError(msg, s, end)
            else:
                _append(terminator)
                continue
        try:
            esc = s[end]
        except IndexError:
            raise JSONDecodeError("Unterminated string starting at",
                                  s, begin) from None
        # If not a unicode escape sequence, must be in the lookup table
        if esc != 'u':
            try:
                char = _b[esc]
            except KeyError:
                msg = "Invalid \\escape: {0!r}".format(esc)
                raise JSONDecodeError(msg, s, end)
            end += 1
        else:
            uni = _decode_uXXXX(s, end)
            end += 5
            if 0xd800 <= uni <= 0xdbff and s[end:end + 2] == '\\u':
                uni2 = _decode_uXXXX(s, end + 1)
                if 0xdc00 <= uni2 <= 0xdfff:
                    uni = 0x10000 + (((uni - 0xd800) << 10) | (uni2 - 0xdc00))
                    end += 6
            char = chr(uni)
        _append(char)
    return ''.join(chunks), end


# Use speedup if available
scanstring = c_scanstring or py_scanstring

WHITESPACE = re.compile(r'[ \t\n\r]*', FLAGS)
WHITESPACE_STR = ' \t\n\r'


def JSONObject(s_and_end, strict, scan_once, object_hook, object_pairs_hook,
               memo=None, _w=WHITESPACE.match, _ws=WHITESPACE_STR):
    s, end = s_and_end
    pairs = []
    pairs_append = pairs.append
    # Backwards compatibility
    if memo is None:
        memo = {}
    memo_get = memo.setdefault
    # Use a slice to prevent IndexError from being raised, the following
    # check will raise a more specific ValueError if the string is empty
    nextchar = s[end:end + 1]
    # Normally we expect nextchar == '"'
    if nextchar != '"':
        if nextchar in _ws:
            end = _w(s, end).end()
            nextchar = s[end:end + 1]
        # Trivial empty object
        if nextchar == '}':
            if object_pairs_hook is not None:
                result = object_pairs_hook(pairs)
                return result, end + 1
            pairs = {}
            if object_hook is not None:
                pairs = object_hook(pairs)
            return pairs, end + 1
        elif nextchar != '"':
            raise JSONDecodeError(
                "Expecting property name enclosed in double quotes", s, end)
    end += 1
    while True:
        key, end = scanstring(s, end, strict)
        key = memo_get(key, key)
        # To skip some function call overhead we optimize the fast paths where
        # the JSON key separator is ": " or just ":".
        if s[end:end + 1] != ':':
            end = _w(s, end).end()
            if s[end:end + 1] != ':':
                raise JSONDecodeError("Expecting ':' delimiter", s, end)
        end += 1

        try:
            if s[end] in _ws:
                end += 1
                if s[end] in _ws:
                    end = _w(s, end + 1).end()
        except IndexError:
            pass

        try:
            value, end = scan_once(s, end)
        except StopIteration as err:
            raise JSONDecodeError("Expecting value", s, err.value) from None
        pairs_append((key, value))
        try:
            nextchar = s[end]
            if nextchar in _ws:
                end = _w(s, end + 1).end()
                nextchar = s[end]
        except IndexError:
            nextchar = ''
        end += 1

        if nextchar == '}':
            break
        elif nextchar != ',':
            raise JSONDecodeError("Expecting ',' delimiter", s, end - 1)
        end = _w(s, end).end()
        nextchar = s[end:end + 1]
        end += 1
        if nextchar != '"':
            raise JSONDecodeError(
                "Expecting property name enclosed in double quotes", s, end - 1)
    if object_pairs_hook is not None:
        result = object_pairs_hook(pairs)
        return result, end
    pairs = dict(pairs)
    if object_hook is not None:
        pairs = object_hook(pairs)
    return pairs, end

def JSONArray(s_and_end, scan_once, _w=WHITESPACE.match, _ws=WHITESPACE_STR):
    s, end = s_and_end
    values = []
    nextchar = s[end:end + 1]
    if nextchar in _ws:
        end = _w(s, end + 1).end()
        nextchar = s[end:end + 1]
    # Look-ahead for trivial empty array
    if nextchar == ']':
        return values, end + 1
    _append = values.append
    while True:
        try:
            value, end = scan_once(s, end)
        except StopIteration as err:
            raise JSONDecodeError("Expecting value", s, err.value) from None
        _append(value)
        nextchar = s[end:end + 1]
        if nextchar in _ws:
            end = _w(s, end + 1).end()
            nextchar = s[end:end + 1]
        end += 1
        if nextchar == ']':
            break
        elif nextchar != ',':
            raise JSONDecodeError("Expecting ',' delimiter", s, end - 1)
        try:
            if s[end] in _ws:
                end += 1
                if s[end] in _ws:
                    end = _w(s, end + 1).end()
        except IndexError:
            pass

    return values, end


class JSONDecoder(object):
    """Simple JSON <http://json.org> decoder

    Performs the following translations in decoding by default:

    +---------------+-------------------+
    | JSON          | Python            |
    +===============+===================+
    | object        | dict              |
    +---------------+-------------------+
    | array         | list              |
    +---------------+-------------------+
    | string        | str               |
    +---------------+-------------------+
    | number (int)  | int               |
    +---------------+-------------------+
    | number (real) | float             |
    +---------------+-------------------+
    | true          | True              |
    +---------------+-------------------+
    | false         | False             |
    +---------------+-------------------+
    | null          | None              |
    +---------------+-------------------+

    It also understands ``NaN``, ``Infinity``, and ``-Infinity`` as
    their corresponding ``float`` values, which is outside the JSON spec.

    """

    def __init__(self, *, object_hook=None, parse_float=None,
            parse_int=None, parse_constant=None, strict=True,
            object_pairs_hook=None):
        """``object_hook``, if specified, will be called with the result
        of every JSON object decoded and its return value will be used in
        place of the given ``dict``.  This can be used to provide custom
        deserializations (e.g. to support JSON-RPC class hinting).

        ``object_pairs_hook``, if specified will be called with the result of
        every JSON object decoded with an ordered list of pairs.  The return
        value of ``object_pairs_hook`` will be used instead of the ``dict``.
        This feature can be used to implement custom decoders.
        If ``object_hook`` is also defined, the ``object_pairs_hook`` takes
        priority.

        ``parse_float``, if specified, will be called with the string
        of every JSON float to be decoded. By default this is equivalent to
        float(num_str). This can be used to use another datatype or parser
        for JSON floats (e.g. decimal.Decimal).

        ``parse_int``, if specified, will be called with the string
        of every JSON int to be decoded. By default this is equivalent to
        int(num_str). This can be used to use another datatype or parser
        for JSON integers (e.g. float).

        ``parse_constant``, if specified, will be called with one of the
        following strings: -Infinity, Infinity, NaN.
        This can be used to raise an exception if invalid JSON numbers
        are encountered.

        If ``strict`` is false (true is the default), then control
        characters will be allowed inside strings.  Control characters in
        this context are those with character codes in the 0-31 range,
        including ``'\\t'`` (tab), ``'\\n'``, ``'\\r'`` and ``'\\0'``.
        """
        self.object_hook = object_hook
        self.parse_float = parse_float or float
        self.parse_int = parse_int or int
        self.parse_constant = parse_constant or _CONSTANTS.__getitem__
        self.strict = strict
        self.object_pairs_hook = object_pairs_hook
        self.parse_object = JSONObject
        self.parse_array = JSONArray
        self.parse_string = scanstring
        self.memo = {}
        self.scan_once = scanner.make_scanner(self)


    def decode(self, s, _w=WHITESPACE.match):
        """Return the Python representation of ``s`` (a ``str`` instance
        containing a JSON document).

        """
        obj, end = self.raw_decode(s, idx=_w(s, 0).end())
        end = _w(s, end).end()
        if end != len(s):
            raise JSONDecodeError("Extra data", s, end)
        return obj

    def raw_decode(self, s, idx=0):
        """Decode a JSON document from ``s`` (a ``str`` beginning with
        a JSON document) and return a 2-tuple of the Python
        representation and the index in ``s`` where the document ended.

        This can be used to decode a JSON document from a string that may
        have extraneous data at the end.

        """
        try:
            obj, end = self.scan_once(s, idx)
        except StopIteration as err:
            raise JSONDecodeError("Expecting value", s, err.value) from None
        return obj, end

external/cpython_master/Lib/json/decoder.py
DragonHPC/dragon
# Copyright 2006 Google, Inc. All Rights Reserved.
# Licensed to PSF under a Contributor Agreement.

"""Fixer for exec.

This converts usages of the exec statement into calls to a built-in
exec() function.

exec code in ns1, ns2 -> exec(code, ns1, ns2)
"""

# Local imports
from .. import fixer_base
from ..fixer_util import Comma, Name, Call


class FixExec(fixer_base.BaseFix):
    BM_compatible = True

    PATTERN = """
    exec_stmt< 'exec' a=any 'in' b=any [',' c=any] >
    |
    exec_stmt< 'exec' (not atom<'(' [any] ')'>) a=any >
    """

    def transform(self, node, results):
        assert results
        syms = self.syms
        a = results["a"]
        b = results.get("b")
        c = results.get("c")
        args = [a.clone()]
        args[0].prefix = ""
        if b is not None:
            args.extend([Comma(), b.clone()])
        if c is not None:
            args.extend([Comma(), c.clone()])

        return Call(Name("exec"), args, prefix=node.prefix)

external/cpython_master/Lib/lib2to3/fixes/fix_exec.py
DragonHPC/dragon
# Copyright 2006 Google, Inc. All Rights Reserved.
# Licensed to PSF under a Contributor Agreement.

"""Fixer for has_key().

Calls to .has_key() methods are expressed in terms of the 'in'
operator:

    d.has_key(k) -> k in d

CAVEATS:
1) While the primary target of this fixer is dict.has_key(), the
   fixer will change any has_key() method call, regardless of its
   class.

2) Cases like this will not be converted:

    m = d.has_key
    if m(k):
        ...

   Only *calls* to has_key() are converted. While it is possible to
   convert the above to something like

    m = d.__contains__
    if m(k):
        ...

   this is currently not done.
"""

# Local imports
from .. import pytree
from .. import fixer_base
from ..fixer_util import Name, parenthesize


class FixHasKey(fixer_base.BaseFix):
    BM_compatible = True

    PATTERN = """
    anchor=power<
        before=any+
        trailer< '.' 'has_key' >
        trailer<
            '('
            ( not(arglist | argument<any '=' any>) arg=any
            | arglist<(not argument<any '=' any>) arg=any ','>
            )
            ')'
        >
        after=any*
    >
    |
    negation=not_test<
        'not'
        anchor=power<
            before=any+
            trailer< '.' 'has_key' >
            trailer<
                '('
                ( not(arglist | argument<any '=' any>) arg=any
                | arglist<(not argument<any '=' any>) arg=any ','>
                )
                ')'
            >
        >
    >
    """

    def transform(self, node, results):
        assert results
        syms = self.syms
        if (node.parent.type == syms.not_test and
            self.pattern.match(node.parent)):
            # Don't transform a node matching the first alternative of the
            # pattern when its parent matches the second alternative
            return None
        negation = results.get("negation")
        anchor = results["anchor"]
        prefix = node.prefix
        before = [n.clone() for n in results["before"]]
        arg = results["arg"].clone()
        after = results.get("after")
        if after:
            after = [n.clone() for n in after]
        if arg.type in (syms.comparison, syms.not_test, syms.and_test,
                        syms.or_test, syms.test, syms.lambdef, syms.argument):
            arg = parenthesize(arg)
        if len(before) == 1:
            before = before[0]
        else:
            before = pytree.Node(syms.power, before)
        before.prefix = " "
        n_op = Name("in", prefix=" ")
        if negation:
            n_not = Name("not", prefix=" ")
            n_op = pytree.Node(syms.comp_op, (n_not, n_op))
        new = pytree.Node(syms.comparison, (arg, n_op, before))
        if after:
            new = parenthesize(new)
            new = pytree.Node(syms.power, (new,) + tuple(after))
        if node.parent.type in (syms.comparison, syms.expr, syms.xor_expr,
                                syms.and_expr, syms.shift_expr,
                                syms.arith_expr, syms.term,
                                syms.factor, syms.power):
            new = parenthesize(new)
        new.prefix = prefix
        return new

external/cpython_master/Lib/lib2to3/fixes/fix_has_key.py
DragonHPC/dragon
# Copyright 2007 Google, Inc. All Rights Reserved.
# Licensed to PSF under a Contributor Agreement.

"""Fixer that changes map(F, ...) into list(map(F, ...)) unless there
exists a 'from future_builtins import map' statement in the top-level
namespace.

As a special case, map(None, X) is changed into list(X).  (This is
necessary because the semantics are changed in this case -- the new
map(None, X) is equivalent to [(x,) for x in X].)

We avoid the transformation (except for the special case mentioned
above) if the map() call is directly contained in iter(<>), list(<>),
tuple(<>), sorted(<>), ...join(<>), or for V in <>:.

NOTE: This is still not correct if the original code was depending on
map(F, X, Y, ...) to go on until the longest argument is exhausted,
substituting None for missing values -- like zip(), it now stops as
soon as the shortest argument is exhausted.
"""

# Local imports
from ..pgen2 import token
from .. import fixer_base
from ..fixer_util import Name, ArgList, Call, ListComp, in_special_context
from ..pygram import python_symbols as syms
from ..pytree import Node


class FixMap(fixer_base.ConditionalFix):
    BM_compatible = True

    PATTERN = """
    map_none=power<
        'map'
        trailer< '(' arglist< 'None' ',' arg=any [','] > ')' >
        [extra_trailers=trailer*]
    >
    |
    map_lambda=power<
        'map'
        trailer<
            '('
            arglist<
                lambdef< 'lambda'
                         (fp=NAME | vfpdef< '(' fp=NAME ')'> ) ':' xp=any
                >
                ','
                it=any
            >
            ')'
        >
        [extra_trailers=trailer*]
    >
    |
    power<
        'map' args=trailer< '(' [any] ')' >
        [extra_trailers=trailer*]
    >
    """

    skip_on = 'future_builtins.map'

    def transform(self, node, results):
        if self.should_skip(node):
            return

        trailers = []
        if 'extra_trailers' in results:
            for t in results['extra_trailers']:
                trailers.append(t.clone())

        if node.parent.type == syms.simple_stmt:
            self.warning(node, "You should use a for loop here")
            new = node.clone()
            new.prefix = ""
            new = Call(Name("list"), [new])
        elif "map_lambda" in results:
            new = ListComp(results["xp"].clone(),
                           results["fp"].clone(),
                           results["it"].clone())
            new = Node(syms.power, [new] + trailers, prefix="")

        else:
            if "map_none" in results:
                new = results["arg"].clone()
                new.prefix = ""
            else:
                if "args" in results:
                    args = results["args"]
                    if args.type == syms.trailer and \
                       args.children[1].type == syms.arglist and \
                       args.children[1].children[0].type == token.NAME and \
                       args.children[1].children[0].value == "None":
                        self.warning(node, "cannot convert map(None, ...) "
                                     "with multiple arguments because map() "
                                     "now truncates to the shortest sequence")
                        return

                    new = Node(syms.power, [Name("map"), args.clone()])
                    new.prefix = ""

                if in_special_context(node):
                    return None

            new = Node(syms.power, [Name("list"), ArgList([new])] + trailers)
            new.prefix = ""

        new.prefix = node.prefix
        return new

external/cpython_master/Lib/lib2to3/fixes/fix_map.py
DragonHPC/dragon
r"""Fixer for unicode.

* Changes unicode to str and unichr to chr.

* If "...\u..." is not unicode literal change it into "...\\u...".

* Change u"..." into "...".

"""

from ..pgen2 import token
from .. import fixer_base

_mapping = {"unichr" : "chr", "unicode" : "str"}

class FixUnicode(fixer_base.BaseFix):
    BM_compatible = True
    PATTERN = "STRING | 'unicode' | 'unichr'"

    def start_tree(self, tree, filename):
        super(FixUnicode, self).start_tree(tree, filename)
        self.unicode_literals = 'unicode_literals' in tree.future_features

    def transform(self, node, results):
        if node.type == token.NAME:
            new = node.clone()
            new.value = _mapping[node.value]
            return new
        elif node.type == token.STRING:
            val = node.value
            if not self.unicode_literals and val[0] in '\'"' and '\\' in val:
                val = r'\\'.join([
                    v.replace('\\u', r'\\u').replace('\\U', r'\\U')
                    for v in val.split(r'\\')
                ])
            if val[0] in 'uU':
                val = val[1:]
            if val == node.value:
                return node
            new = node.clone()
            new.value = val
            return new

external/cpython_master/Lib/lib2to3/fixes/fix_unicode.py
DragonHPC/dragon
# This is empty so trying to fetch the fixer class gives an AttributeError

external/cpython_master/Lib/lib2to3/tests/data/fixers/no_fixer_cls.py
DragonHPC/dragon
# This file is used to verify that 2to3 falls back to a slower, iterative pattern matching
# scheme in the event that the faster recursive system fails due to infinite recursion.
from ctypes import *
STRING = c_char_p


OSUnknownByteOrder = 0
UIT_PROMPT = 1
P_PGID = 2
P_PID = 1
UIT_ERROR = 5
UIT_INFO = 4
UIT_NONE = 0
P_ALL = 0
UIT_VERIFY = 2
OSBigEndian = 2
UIT_BOOLEAN = 3
OSLittleEndian = 1
__darwin_nl_item = c_int
__darwin_wctrans_t = c_int
__darwin_wctype_t = c_ulong
__int8_t = c_byte
__uint8_t = c_ubyte
__int16_t = c_short
__uint16_t = c_ushort
__int32_t = c_int
__uint32_t = c_uint
__int64_t = c_longlong
__uint64_t = c_ulonglong
__darwin_intptr_t = c_long
__darwin_natural_t = c_uint
__darwin_ct_rune_t = c_int
class __mbstate_t(Union):
    pass
__mbstate_t._pack_ = 4
__mbstate_t._fields_ = [
    ('__mbstate8', c_char * 128),
    ('_mbstateL', c_longlong),
]
assert sizeof(__mbstate_t) == 128, sizeof(__mbstate_t)
assert alignment(__mbstate_t) == 4, alignment(__mbstate_t)
__darwin_mbstate_t = __mbstate_t
__darwin_ptrdiff_t = c_int
__darwin_size_t = c_ulong
__darwin_va_list = STRING
__darwin_wchar_t = c_int
__darwin_rune_t = __darwin_wchar_t
__darwin_wint_t = c_int
__darwin_clock_t = c_ulong
__darwin_socklen_t = __uint32_t
__darwin_ssize_t = c_long
__darwin_time_t = c_long
sig_atomic_t = c_int
class sigcontext(Structure):
    pass
sigcontext._fields_ = [
    ('sc_onstack', c_int),
    ('sc_mask', c_int),
    ('sc_eax', c_uint),
    ('sc_ebx', c_uint),
    ('sc_ecx', c_uint),
    ('sc_edx', c_uint),
    ('sc_edi', c_uint),
    ('sc_esi', c_uint),
    ('sc_ebp', c_uint),
    ('sc_esp', c_uint),
    ('sc_ss', c_uint),
    ('sc_eflags', c_uint),
    ('sc_eip', c_uint),
    ('sc_cs', c_uint),
    ('sc_ds', c_uint),
    ('sc_es', c_uint),
    ('sc_fs', c_uint),
    ('sc_gs', c_uint),
]
assert sizeof(sigcontext) == 72, sizeof(sigcontext)
assert alignment(sigcontext) == 4, alignment(sigcontext)
u_int8_t = c_ubyte
u_int16_t = c_ushort
u_int32_t = c_uint
u_int64_t = c_ulonglong
int32_t = c_int
register_t = int32_t
user_addr_t = u_int64_t
user_size_t = u_int64_t
int64_t = c_longlong
user_ssize_t = int64_t
user_long_t = int64_t
user_ulong_t = u_int64_t
user_time_t = int64_t
syscall_arg_t = u_int64_t

# values for unnamed enumeration
class aes_key_st(Structure):
    pass
aes_key_st._fields_ = [
    ('rd_key', c_ulong * 60),
    ('rounds', c_int),
]
assert sizeof(aes_key_st) == 244, sizeof(aes_key_st)
assert alignment(aes_key_st) == 4, alignment(aes_key_st)
AES_KEY = aes_key_st
class asn1_ctx_st(Structure):
    pass
asn1_ctx_st._fields_ = [
    ('p', POINTER(c_ubyte)),
    ('eos', c_int),
    ('error', c_int),
    ('inf', c_int),
    ('tag', c_int),
    ('xclass', c_int),
    ('slen', c_long),
    ('max', POINTER(c_ubyte)),
    ('q', POINTER(c_ubyte)),
    ('pp', POINTER(POINTER(c_ubyte))),
    ('line', c_int),
]
assert sizeof(asn1_ctx_st) == 44, sizeof(asn1_ctx_st)
assert alignment(asn1_ctx_st) == 4, alignment(asn1_ctx_st)
ASN1_CTX = asn1_ctx_st
class asn1_object_st(Structure):
    pass
asn1_object_st._fields_ = [
    ('sn', STRING),
    ('ln', STRING),
    ('nid', c_int),
    ('length', c_int),
    ('data', POINTER(c_ubyte)),
    ('flags', c_int),
]
assert sizeof(asn1_object_st) == 24, sizeof(asn1_object_st)
assert alignment(asn1_object_st) == 4, alignment(asn1_object_st)
ASN1_OBJECT = asn1_object_st
class asn1_string_st(Structure):
    pass
asn1_string_st._fields_ = [
    ('length', c_int),
    ('type', c_int),
    ('data', POINTER(c_ubyte)),
    ('flags', c_long),
]
assert sizeof(asn1_string_st) == 16, sizeof(asn1_string_st)
assert alignment(asn1_string_st) == 4, alignment(asn1_string_st)
ASN1_STRING = asn1_string_st
class ASN1_ENCODING_st(Structure):
    pass
ASN1_ENCODING_st._fields_ = [
    ('enc', POINTER(c_ubyte)),
    ('len', c_long),
    ('modified', c_int),
]
assert sizeof(ASN1_ENCODING_st) == 12, sizeof(ASN1_ENCODING_st)
assert alignment(ASN1_ENCODING_st) == 4, alignment(ASN1_ENCODING_st)
ASN1_ENCODING = ASN1_ENCODING_st
class asn1_string_table_st(Structure):
    pass
asn1_string_table_st._fields_ = [
    ('nid', c_int),
    ('minsize', c_long),
    ('maxsize', c_long),
    ('mask', c_ulong),
    ('flags', c_ulong),
]
assert sizeof(asn1_string_table_st) == 20, sizeof(asn1_string_table_st)
assert alignment(asn1_string_table_st) == 4, alignment(asn1_string_table_st)
ASN1_STRING_TABLE = asn1_string_table_st
class ASN1_TEMPLATE_st(Structure):
    pass
ASN1_TEMPLATE_st._fields_ = [
]
ASN1_TEMPLATE = ASN1_TEMPLATE_st
class ASN1_ITEM_st(Structure):
    pass
ASN1_ITEM = ASN1_ITEM_st
ASN1_ITEM_st._fields_ = [
]
class ASN1_TLC_st(Structure):
    pass
ASN1_TLC = ASN1_TLC_st
ASN1_TLC_st._fields_ = [
]
class ASN1_VALUE_st(Structure):
    pass
ASN1_VALUE_st._fields_ = [
]
ASN1_VALUE = ASN1_VALUE_st
ASN1_ITEM_EXP = ASN1_ITEM
class asn1_type_st(Structure):
    pass
class N12asn1_type_st4DOLLAR_11E(Union):
    pass
ASN1_BOOLEAN = c_int
ASN1_INTEGER = asn1_string_st
ASN1_ENUMERATED = asn1_string_st
ASN1_BIT_STRING = asn1_string_st
ASN1_OCTET_STRING = asn1_string_st
ASN1_PRINTABLESTRING = asn1_string_st
ASN1_T61STRING = asn1_string_st
ASN1_IA5STRING = asn1_string_st
ASN1_GENERALSTRING = asn1_string_st
ASN1_BMPSTRING = asn1_string_st
ASN1_UNIVERSALSTRING = asn1_string_st
ASN1_UTCTIME = asn1_string_st
ASN1_GENERALIZEDTIME = asn1_string_st
ASN1_VISIBLESTRING = asn1_string_st
ASN1_UTF8STRING = asn1_string_st
N12asn1_type_st4DOLLAR_11E._fields_ = [
    ('ptr', STRING),
    ('boolean', ASN1_BOOLEAN),
    ('asn1_string', POINTER(ASN1_STRING)),
    ('object', POINTER(ASN1_OBJECT)),
    ('integer', POINTER(ASN1_INTEGER)),
    ('enumerated', POINTER(ASN1_ENUMERATED)),
    ('bit_string', POINTER(ASN1_BIT_STRING)),
    ('octet_string', POINTER(ASN1_OCTET_STRING)),
    ('printablestring', POINTER(ASN1_PRINTABLESTRING)),
    ('t61string', POINTER(ASN1_T61STRING)),
    ('ia5string', POINTER(ASN1_IA5STRING)),
    ('generalstring', POINTER(ASN1_GENERALSTRING)),
    ('bmpstring', POINTER(ASN1_BMPSTRING)),
    ('universalstring', POINTER(ASN1_UNIVERSALSTRING)),
    ('utctime', POINTER(ASN1_UTCTIME)),
    ('generalizedtime', POINTER(ASN1_GENERALIZEDTIME)),
    ('visiblestring', POINTER(ASN1_VISIBLESTRING)),
    ('utf8string', POINTER(ASN1_UTF8STRING)),
    ('set', POINTER(ASN1_STRING)),
    ('sequence', POINTER(ASN1_STRING)),
]
assert sizeof(N12asn1_type_st4DOLLAR_11E) == 4, sizeof(N12asn1_type_st4DOLLAR_11E)
assert alignment(N12asn1_type_st4DOLLAR_11E) == 4, alignment(N12asn1_type_st4DOLLAR_11E)
asn1_type_st._fields_ = [
    ('type', c_int),
    ('value', N12asn1_type_st4DOLLAR_11E),
]
assert sizeof(asn1_type_st) == 8, sizeof(asn1_type_st)
assert alignment(asn1_type_st) == 4, alignment(asn1_type_st)
ASN1_TYPE = asn1_type_st
class asn1_method_st(Structure):
    pass
asn1_method_st._fields_ = [
    ('i2d', CFUNCTYPE(c_int)),
    ('d2i', CFUNCTYPE(STRING)),
    ('create', CFUNCTYPE(STRING)),
    ('destroy', CFUNCTYPE(None)),
]
assert sizeof(asn1_method_st) == 16, sizeof(asn1_method_st)
assert alignment(asn1_method_st) == 4, alignment(asn1_method_st)
ASN1_METHOD = asn1_method_st
class asn1_header_st(Structure):
    pass
asn1_header_st._fields_ = [
    ('header', POINTER(ASN1_OCTET_STRING)),
    ('data', STRING),
    ('meth', POINTER(ASN1_METHOD)),
]
assert sizeof(asn1_header_st) == 12, sizeof(asn1_header_st)
assert alignment(asn1_header_st) == 4, alignment(asn1_header_st)
ASN1_HEADER = asn1_header_st
class BIT_STRING_BITNAME_st(Structure):
    pass
BIT_STRING_BITNAME_st._fields_ = [
    ('bitnum', c_int),
    ('lname', STRING),
    ('sname', STRING),
]
assert sizeof(BIT_STRING_BITNAME_st) == 12, sizeof(BIT_STRING_BITNAME_st)
assert alignment(BIT_STRING_BITNAME_st) == 4, alignment(BIT_STRING_BITNAME_st)
BIT_STRING_BITNAME = BIT_STRING_BITNAME_st
class bio_st(Structure):
    pass
BIO = bio_st
bio_info_cb = CFUNCTYPE(None, POINTER(bio_st), c_int, STRING, c_int, c_long, c_long)
class bio_method_st(Structure):
    pass
bio_method_st._fields_ = [
    ('type', c_int),
    ('name', STRING),
    ('bwrite', CFUNCTYPE(c_int, POINTER(BIO), STRING, c_int)),
    ('bread', CFUNCTYPE(c_int, POINTER(BIO), STRING, c_int)),
    ('bputs', CFUNCTYPE(c_int, POINTER(BIO), STRING)),
    ('bgets', CFUNCTYPE(c_int, POINTER(BIO), STRING, c_int)),
    ('ctrl', CFUNCTYPE(c_long, POINTER(BIO), c_int, c_long, c_void_p)),
    ('create', CFUNCTYPE(c_int, POINTER(BIO))),
    ('destroy', CFUNCTYPE(c_int, POINTER(BIO))),
    ('callback_ctrl', CFUNCTYPE(c_long, POINTER(BIO), c_int, POINTER(bio_info_cb))),
]
assert sizeof(bio_method_st) == 40, sizeof(bio_method_st)
assert alignment(bio_method_st) == 4, alignment(bio_method_st)
BIO_METHOD = bio_method_st
class crypto_ex_data_st(Structure):
    pass
class stack_st(Structure):
    pass
STACK = stack_st
crypto_ex_data_st._fields_ = [
    ('sk', POINTER(STACK)),
    ('dummy', c_int),
]
assert sizeof(crypto_ex_data_st) == 8, sizeof(crypto_ex_data_st)
assert alignment(crypto_ex_data_st) == 4, alignment(crypto_ex_data_st)
CRYPTO_EX_DATA = crypto_ex_data_st
bio_st._fields_ = [
    ('method', POINTER(BIO_METHOD)),
    ('callback', CFUNCTYPE(c_long, POINTER(bio_st), c_int, STRING, c_int, c_long, c_long)),
    ('cb_arg', STRING),
    ('init', c_int),
    ('shutdown', c_int),
    ('flags', c_int),
    ('retry_reason', c_int),
    ('num', c_int),
    ('ptr', c_void_p),
    ('next_bio', POINTER(bio_st)),
    ('prev_bio', POINTER(bio_st)),
    ('references', c_int),
    ('num_read', c_ulong),
    ('num_write', c_ulong),
    ('ex_data', CRYPTO_EX_DATA),
]
assert sizeof(bio_st) == 64, sizeof(bio_st)
assert alignment(bio_st) == 4, alignment(bio_st)
class bio_f_buffer_ctx_struct(Structure):
    pass
bio_f_buffer_ctx_struct._fields_ = [
    ('ibuf_size', c_int),
    ('obuf_size', c_int),
    ('ibuf', STRING),
    ('ibuf_len', c_int),
    ('ibuf_off', c_int),
    ('obuf', STRING),
    ('obuf_len', c_int),
    ('obuf_off', c_int),
]
assert sizeof(bio_f_buffer_ctx_struct) == 32, sizeof(bio_f_buffer_ctx_struct)
assert alignment(bio_f_buffer_ctx_struct) == 4, alignment(bio_f_buffer_ctx_struct)
BIO_F_BUFFER_CTX = bio_f_buffer_ctx_struct
class hostent(Structure):
    pass
hostent._fields_ = [
]
class bf_key_st(Structure):
    pass
bf_key_st._fields_ = [
    ('P', c_uint * 18),
    ('S', c_uint * 1024),
]
assert sizeof(bf_key_st) == 4168, sizeof(bf_key_st)
assert alignment(bf_key_st) == 4, alignment(bf_key_st)
BF_KEY = bf_key_st
class bignum_st(Structure):
    pass
bignum_st._fields_ = [
    ('d', POINTER(c_ulong)),
    ('top', c_int),
    ('dmax', c_int),
    ('neg', c_int),
    ('flags', c_int),
]
assert sizeof(bignum_st) == 20, sizeof(bignum_st)
assert alignment(bignum_st) == 4, alignment(bignum_st)
BIGNUM = bignum_st
class bignum_ctx(Structure):
    pass
bignum_ctx._fields_ = [
]
BN_CTX = bignum_ctx
class bn_blinding_st(Structure):
    pass
bn_blinding_st._fields_ = [
    ('init', c_int),
    ('A', POINTER(BIGNUM)),
    ('Ai', POINTER(BIGNUM)),
    ('mod', POINTER(BIGNUM)),
    ('thread_id', c_ulong),
]
assert sizeof(bn_blinding_st) == 20, sizeof(bn_blinding_st)
assert alignment(bn_blinding_st) == 4, alignment(bn_blinding_st)
BN_BLINDING = bn_blinding_st
class bn_mont_ctx_st(Structure):
    pass
bn_mont_ctx_st._fields_ = [
    ('ri', c_int),
    ('RR', BIGNUM),
    ('N', BIGNUM),
    ('Ni', BIGNUM),
    ('n0', c_ulong),
    ('flags', c_int),
]
assert sizeof(bn_mont_ctx_st) == 72, sizeof(bn_mont_ctx_st)
assert alignment(bn_mont_ctx_st) == 4, alignment(bn_mont_ctx_st)
BN_MONT_CTX = bn_mont_ctx_st
class bn_recp_ctx_st(Structure):
    pass
bn_recp_ctx_st._fields_ = [
    ('N', BIGNUM),
    ('Nr', BIGNUM),
    ('num_bits', c_int),
    ('shift', c_int),
    ('flags', c_int),
]
assert sizeof(bn_recp_ctx_st) == 52, sizeof(bn_recp_ctx_st)
assert alignment(bn_recp_ctx_st) == 4, alignment(bn_recp_ctx_st)
BN_RECP_CTX = bn_recp_ctx_st
class buf_mem_st(Structure):
    pass
buf_mem_st._fields_ = [
    ('length', c_int),
    ('data', STRING),
    ('max', c_int),
]
assert sizeof(buf_mem_st) == 12, sizeof(buf_mem_st)
assert alignment(buf_mem_st) == 4, alignment(buf_mem_st)
BUF_MEM = buf_mem_st
class cast_key_st(Structure):
    pass
cast_key_st._fields_ = [
    ('data', c_ulong * 32),
    ('short_key', c_int),
]
assert sizeof(cast_key_st) == 132, sizeof(cast_key_st)
assert alignment(cast_key_st) == 4, alignment(cast_key_st)
CAST_KEY = cast_key_st
class comp_method_st(Structure):
    pass
comp_method_st._fields_ = [
    ('type', c_int),
    ('name', STRING),
    ('init', CFUNCTYPE(c_int)),
    ('finish', CFUNCTYPE(None)),
    ('compress', CFUNCTYPE(c_int)),
    ('expand', CFUNCTYPE(c_int)),
    ('ctrl', CFUNCTYPE(c_long)),
    ('callback_ctrl', CFUNCTYPE(c_long)),
]
assert sizeof(comp_method_st) == 32, sizeof(comp_method_st)
assert alignment(comp_method_st) == 4, alignment(comp_method_st)
COMP_METHOD = comp_method_st
class comp_ctx_st(Structure):
    pass
comp_ctx_st._fields_ = [
    ('meth', POINTER(COMP_METHOD)),
    ('compress_in', c_ulong),
    ('compress_out', c_ulong),
    ('expand_in', c_ulong),
    ('expand_out', c_ulong),
    ('ex_data', CRYPTO_EX_DATA),
]
assert sizeof(comp_ctx_st) == 28, sizeof(comp_ctx_st)
assert alignment(comp_ctx_st) == 4, alignment(comp_ctx_st)
COMP_CTX = comp_ctx_st
class CRYPTO_dynlock_value(Structure):
    pass
CRYPTO_dynlock_value._fields_ = [
]
class CRYPTO_dynlock(Structure):
    pass
CRYPTO_dynlock._fields_ = [
    ('references', c_int),
    ('data', POINTER(CRYPTO_dynlock_value)),
]
assert sizeof(CRYPTO_dynlock) == 8, sizeof(CRYPTO_dynlock)
assert alignment(CRYPTO_dynlock) == 4, alignment(CRYPTO_dynlock)
BIO_dummy = bio_st
CRYPTO_EX_new = CFUNCTYPE(c_int, c_void_p, c_void_p, POINTER(CRYPTO_EX_DATA), c_int, c_long, c_void_p)
CRYPTO_EX_free = CFUNCTYPE(None, c_void_p, c_void_p, POINTER(CRYPTO_EX_DATA), c_int, c_long, c_void_p)
CRYPTO_EX_dup = CFUNCTYPE(c_int, POINTER(CRYPTO_EX_DATA), POINTER(CRYPTO_EX_DATA), c_void_p, c_int, c_long, c_void_p)
class crypto_ex_data_func_st(Structure):
    pass
crypto_ex_data_func_st._fields_ = [
    ('argl', c_long),
    ('argp', c_void_p),
    ('new_func', POINTER(CRYPTO_EX_new)),
    ('free_func', POINTER(CRYPTO_EX_free)),
    ('dup_func', POINTER(CRYPTO_EX_dup)),
]
assert sizeof(crypto_ex_data_func_st) == 20, sizeof(crypto_ex_data_func_st)
assert alignment(crypto_ex_data_func_st) == 4, alignment(crypto_ex_data_func_st)
CRYPTO_EX_DATA_FUNCS = crypto_ex_data_func_st
class st_CRYPTO_EX_DATA_IMPL(Structure):
    pass
CRYPTO_EX_DATA_IMPL = st_CRYPTO_EX_DATA_IMPL
st_CRYPTO_EX_DATA_IMPL._fields_ = [
]
CRYPTO_MEM_LEAK_CB = CFUNCTYPE(c_void_p, c_ulong, STRING, c_int, c_int, c_void_p)
DES_cblock = c_ubyte * 8
const_DES_cblock = c_ubyte * 8
class DES_ks(Structure):
    pass
class N6DES_ks3DOLLAR_9E(Union):
    pass
N6DES_ks3DOLLAR_9E._fields_ = [
    ('cblock', DES_cblock),
    ('deslong', c_ulong * 2),
]
assert sizeof(N6DES_ks3DOLLAR_9E) == 8, sizeof(N6DES_ks3DOLLAR_9E)
assert alignment(N6DES_ks3DOLLAR_9E) == 4, alignment(N6DES_ks3DOLLAR_9E)
DES_ks._fields_ = [
    ('ks', N6DES_ks3DOLLAR_9E * 16),
]
assert sizeof(DES_ks) == 128, sizeof(DES_ks)
assert alignment(DES_ks) == 4, alignment(DES_ks)
DES_key_schedule = DES_ks
_ossl_old_des_cblock = c_ubyte * 8
class _ossl_old_des_ks_struct(Structure):
    pass
class N23_ossl_old_des_ks_struct4DOLLAR_10E(Union):
    pass
N23_ossl_old_des_ks_struct4DOLLAR_10E._fields_ = [
    ('_', _ossl_old_des_cblock),
    ('pad', c_ulong * 2),
]
assert sizeof(N23_ossl_old_des_ks_struct4DOLLAR_10E) == 8, sizeof(N23_ossl_old_des_ks_struct4DOLLAR_10E)
assert alignment(N23_ossl_old_des_ks_struct4DOLLAR_10E) == 4, alignment(N23_ossl_old_des_ks_struct4DOLLAR_10E)
_ossl_old_des_ks_struct._fields_ = [
    ('ks', N23_ossl_old_des_ks_struct4DOLLAR_10E),
]
assert sizeof(_ossl_old_des_ks_struct) == 8, sizeof(_ossl_old_des_ks_struct)
assert alignment(_ossl_old_des_ks_struct) == 4, alignment(_ossl_old_des_ks_struct)
_ossl_old_des_key_schedule = _ossl_old_des_ks_struct * 16
class dh_st(Structure):
    pass
DH = dh_st
class dh_method(Structure):
    pass
dh_method._fields_ = [
    ('name', STRING),
    ('generate_key', CFUNCTYPE(c_int, POINTER(DH))),
    ('compute_key', CFUNCTYPE(c_int, POINTER(c_ubyte), POINTER(BIGNUM), POINTER(DH))),
    ('bn_mod_exp', CFUNCTYPE(c_int, POINTER(DH), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BN_CTX), POINTER(BN_MONT_CTX))),
    ('init', CFUNCTYPE(c_int, POINTER(DH))),
    ('finish', CFUNCTYPE(c_int, POINTER(DH))),
    ('flags', c_int),
    ('app_data', STRING),
]
assert sizeof(dh_method) == 32, sizeof(dh_method)
assert alignment(dh_method) == 4, alignment(dh_method)
DH_METHOD = dh_method
class engine_st(Structure):
    pass
ENGINE = engine_st
dh_st._fields_ = [
    ('pad', c_int),
    ('version', c_int),
    ('p', POINTER(BIGNUM)),
    ('g', POINTER(BIGNUM)),
    ('length', c_long),
    ('pub_key', POINTER(BIGNUM)),
    ('priv_key', POINTER(BIGNUM)),
    ('flags', c_int),
    ('method_mont_p', STRING),
    ('q', POINTER(BIGNUM)),
    ('j', POINTER(BIGNUM)),
    ('seed', POINTER(c_ubyte)),
    ('seedlen', c_int),
    ('counter', POINTER(BIGNUM)),
    ('references', c_int),
    ('ex_data', CRYPTO_EX_DATA),
    ('meth', POINTER(DH_METHOD)),
    ('engine', POINTER(ENGINE)),
]
assert sizeof(dh_st) == 76, sizeof(dh_st)
assert alignment(dh_st) == 4, alignment(dh_st)
class dsa_st(Structure):
    pass
DSA = dsa_st
class DSA_SIG_st(Structure):
    pass
DSA_SIG_st._fields_ = [
    ('r', POINTER(BIGNUM)),
    ('s', POINTER(BIGNUM)),
]
assert sizeof(DSA_SIG_st) == 8, sizeof(DSA_SIG_st)
assert alignment(DSA_SIG_st) == 4, alignment(DSA_SIG_st)
DSA_SIG = DSA_SIG_st
class dsa_method(Structure):
    pass
dsa_method._fields_ = [
    ('name', STRING),
    ('dsa_do_sign', CFUNCTYPE(POINTER(DSA_SIG), POINTER(c_ubyte), c_int, POINTER(DSA))),
    ('dsa_sign_setup', CFUNCTYPE(c_int, POINTER(DSA), POINTER(BN_CTX), POINTER(POINTER(BIGNUM)), POINTER(POINTER(BIGNUM)))),
    ('dsa_do_verify', CFUNCTYPE(c_int, POINTER(c_ubyte), c_int, POINTER(DSA_SIG), POINTER(DSA))),
    ('dsa_mod_exp', CFUNCTYPE(c_int, POINTER(DSA), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BN_CTX), POINTER(BN_MONT_CTX))),
    ('bn_mod_exp', CFUNCTYPE(c_int, POINTER(DSA), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BN_CTX), POINTER(BN_MONT_CTX))),
    ('init', CFUNCTYPE(c_int, POINTER(DSA))),
    ('finish', CFUNCTYPE(c_int, POINTER(DSA))),
    ('flags', c_int),
    ('app_data', STRING),
]
assert sizeof(dsa_method) == 40, sizeof(dsa_method)
assert alignment(dsa_method) == 4, alignment(dsa_method)
DSA_METHOD = dsa_method
dsa_st._fields_ = [
    ('pad', c_int),
    ('version', c_long),
    ('write_params', c_int),
    ('p', POINTER(BIGNUM)),
    ('q', POINTER(BIGNUM)),
    ('g', POINTER(BIGNUM)),
    ('pub_key', POINTER(BIGNUM)),
    ('priv_key', POINTER(BIGNUM)),
    ('kinv', POINTER(BIGNUM)),
    ('r', POINTER(BIGNUM)),
    ('flags', c_int),
    ('method_mont_p', STRING),
    ('references', c_int),
    ('ex_data', CRYPTO_EX_DATA),
    ('meth', POINTER(DSA_METHOD)),
    ('engine', POINTER(ENGINE)),
]
assert sizeof(dsa_st) == 68, sizeof(dsa_st)
assert alignment(dsa_st) == 4, alignment(dsa_st)
class evp_pkey_st(Structure):
    pass
class N11evp_pkey_st4DOLLAR_12E(Union):
    pass
class rsa_st(Structure):
    pass
N11evp_pkey_st4DOLLAR_12E._fields_ = [
    ('ptr', STRING),
    ('rsa', POINTER(rsa_st)),
    ('dsa', POINTER(dsa_st)),
    ('dh', POINTER(dh_st)),
]
assert sizeof(N11evp_pkey_st4DOLLAR_12E) == 4, sizeof(N11evp_pkey_st4DOLLAR_12E)
assert alignment(N11evp_pkey_st4DOLLAR_12E) == 4, alignment(N11evp_pkey_st4DOLLAR_12E)
evp_pkey_st._fields_ = [
    ('type', c_int),
    ('save_type', c_int),
    ('references', c_int),
    ('pkey', N11evp_pkey_st4DOLLAR_12E),
    ('save_parameters', c_int),
    ('attributes', POINTER(STACK)),
]
assert sizeof(evp_pkey_st) == 24, sizeof(evp_pkey_st)
assert alignment(evp_pkey_st) == 4, alignment(evp_pkey_st)
class env_md_st(Structure):
    pass
class env_md_ctx_st(Structure):
    pass
EVP_MD_CTX = env_md_ctx_st
env_md_st._fields_ = [
    ('type', c_int),
    ('pkey_type', c_int),
    ('md_size', c_int),
    ('flags', c_ulong),
    ('init', CFUNCTYPE(c_int, POINTER(EVP_MD_CTX))),
    ('update', CFUNCTYPE(c_int, POINTER(EVP_MD_CTX), c_void_p, c_ulong)),
    ('final', CFUNCTYPE(c_int, POINTER(EVP_MD_CTX), POINTER(c_ubyte))),
    ('copy', CFUNCTYPE(c_int, POINTER(EVP_MD_CTX), POINTER(EVP_MD_CTX))),
    ('cleanup', CFUNCTYPE(c_int, POINTER(EVP_MD_CTX))),
    ('sign', CFUNCTYPE(c_int)),
    ('verify', CFUNCTYPE(c_int)),
    ('required_pkey_type', c_int * 5),
    ('block_size', c_int),
    ('ctx_size', c_int),
]
assert sizeof(env_md_st) == 72, sizeof(env_md_st)
assert alignment(env_md_st) == 4, alignment(env_md_st)
EVP_MD = env_md_st
env_md_ctx_st._fields_ = [
    ('digest', POINTER(EVP_MD)),
    ('engine', POINTER(ENGINE)),
    ('flags', c_ulong),
    ('md_data', c_void_p),
]
assert sizeof(env_md_ctx_st) == 16, sizeof(env_md_ctx_st)
assert alignment(env_md_ctx_st) == 4, alignment(env_md_ctx_st)
class evp_cipher_st(Structure):
    pass
class evp_cipher_ctx_st(Structure):
    pass
EVP_CIPHER_CTX = evp_cipher_ctx_st
evp_cipher_st._fields_ = [
    ('nid', c_int),
    ('block_size', c_int),
    ('key_len', c_int),
    ('iv_len', c_int),
    ('flags', c_ulong),
    ('init', CFUNCTYPE(c_int, POINTER(EVP_CIPHER_CTX), POINTER(c_ubyte), POINTER(c_ubyte), c_int)),
    ('do_cipher', CFUNCTYPE(c_int, POINTER(EVP_CIPHER_CTX), POINTER(c_ubyte), POINTER(c_ubyte), c_uint)),
    ('cleanup', CFUNCTYPE(c_int, POINTER(EVP_CIPHER_CTX))),
    ('ctx_size', c_int),
    ('set_asn1_parameters', CFUNCTYPE(c_int, POINTER(EVP_CIPHER_CTX), POINTER(ASN1_TYPE))),
    ('get_asn1_parameters', CFUNCTYPE(c_int, POINTER(EVP_CIPHER_CTX), POINTER(ASN1_TYPE))),
    ('ctrl', CFUNCTYPE(c_int, POINTER(EVP_CIPHER_CTX), c_int, c_int, c_void_p)),
    ('app_data', c_void_p),
]
assert sizeof(evp_cipher_st) == 52, sizeof(evp_cipher_st)
assert alignment(evp_cipher_st) == 4, alignment(evp_cipher_st)
class evp_cipher_info_st(Structure):
    pass
EVP_CIPHER = evp_cipher_st
evp_cipher_info_st._fields_ = [
    ('cipher', POINTER(EVP_CIPHER)),
    ('iv', c_ubyte * 16),
]
assert sizeof(evp_cipher_info_st) == 20, sizeof(evp_cipher_info_st)
assert alignment(evp_cipher_info_st) == 4, alignment(evp_cipher_info_st)
EVP_CIPHER_INFO = evp_cipher_info_st
evp_cipher_ctx_st._fields_ = [
    ('cipher', POINTER(EVP_CIPHER)),
    ('engine', POINTER(ENGINE)),
    ('encrypt', c_int),
    ('buf_len', c_int),
    ('oiv', c_ubyte * 16),
    ('iv', c_ubyte * 16),
    ('buf', c_ubyte * 32),
    ('num', c_int),
    ('app_data', c_void_p),
    ('key_len', c_int),
    ('flags', c_ulong),
    ('cipher_data', c_void_p),
    ('final_used', c_int),
    ('block_mask', c_int),
    ('final', c_ubyte * 32),
]
assert sizeof(evp_cipher_ctx_st) == 140, sizeof(evp_cipher_ctx_st)
assert alignment(evp_cipher_ctx_st) == 4, alignment(evp_cipher_ctx_st)
class evp_Encode_Ctx_st(Structure):
    pass
evp_Encode_Ctx_st._fields_ = [
    ('num', c_int),
    ('length', c_int),
    ('enc_data', c_ubyte * 80),
    ('line_num', c_int),
    ('expect_nl', c_int),
]
assert sizeof(evp_Encode_Ctx_st) == 96, sizeof(evp_Encode_Ctx_st)
assert alignment(evp_Encode_Ctx_st) == 4, alignment(evp_Encode_Ctx_st)
EVP_ENCODE_CTX = evp_Encode_Ctx_st
EVP_PBE_KEYGEN = CFUNCTYPE(c_int, POINTER(EVP_CIPHER_CTX), STRING, c_int, POINTER(ASN1_TYPE), POINTER(EVP_CIPHER), POINTER(EVP_MD), c_int)
class lhash_node_st(Structure):
    pass
lhash_node_st._fields_ = [
    ('data', c_void_p),
    ('next', POINTER(lhash_node_st)),
    ('hash', c_ulong),
]
assert sizeof(lhash_node_st) == 12, sizeof(lhash_node_st)
assert alignment(lhash_node_st) == 4, alignment(lhash_node_st)
LHASH_NODE = lhash_node_st
LHASH_COMP_FN_TYPE = CFUNCTYPE(c_int, c_void_p, c_void_p)
LHASH_HASH_FN_TYPE = CFUNCTYPE(c_ulong, c_void_p)
LHASH_DOALL_FN_TYPE = CFUNCTYPE(None, c_void_p)
LHASH_DOALL_ARG_FN_TYPE = CFUNCTYPE(None, c_void_p, c_void_p)
class lhash_st(Structure):
    pass
lhash_st._fields_ = [
    ('b', POINTER(POINTER(LHASH_NODE))),
    ('comp', LHASH_COMP_FN_TYPE),
    ('hash', LHASH_HASH_FN_TYPE),
    ('num_nodes', c_uint),
    ('num_alloc_nodes', c_uint),
    ('p', c_uint),
    ('pmax', c_uint),
    ('up_load', c_ulong),
    ('down_load', c_ulong),
    ('num_items', c_ulong),
    ('num_expands', c_ulong),
    ('num_expand_reallocs', c_ulong),
    ('num_contracts', c_ulong),
    ('num_contract_reallocs', c_ulong),
    ('num_hash_calls', c_ulong),
    ('num_comp_calls', c_ulong),
    ('num_insert', c_ulong),
    ('num_replace', c_ulong),
    ('num_delete', c_ulong),
    ('num_no_delete', c_ulong),
    ('num_retrieve', c_ulong),
    ('num_retrieve_miss', c_ulong),
    ('num_hash_comps', c_ulong),
    ('error', c_int),
]
assert sizeof(lhash_st) == 96, sizeof(lhash_st)
assert alignment(lhash_st) == 4, alignment(lhash_st)
LHASH = lhash_st
class MD2state_st(Structure):
    pass
MD2state_st._fields_ = [
    ('num', c_int),
    ('data', c_ubyte * 16),
    ('cksm', c_uint * 16),
    ('state', c_uint * 16),
]
assert sizeof(MD2state_st) == 148, sizeof(MD2state_st)
assert alignment(MD2state_st) == 4, alignment(MD2state_st)
MD2_CTX = MD2state_st
class MD4state_st(Structure):
    pass
MD4state_st._fields_ = [
    ('A', c_uint),
    ('B', c_uint),
    ('C', c_uint),
    ('D', c_uint),
    ('Nl', c_uint),
    ('Nh', c_uint),
    ('data', c_uint * 16),
    ('num', c_int),
]
assert sizeof(MD4state_st) == 92, sizeof(MD4state_st)
assert alignment(MD4state_st) == 4, alignment(MD4state_st)
MD4_CTX = MD4state_st
class MD5state_st(Structure):
    pass
MD5state_st._fields_ = [
    ('A', c_uint),
    ('B', c_uint),
    ('C', c_uint),
    ('D', c_uint),
    ('Nl', c_uint),
    ('Nh', c_uint),
    ('data', c_uint * 16),
    ('num', c_int),
]
assert sizeof(MD5state_st) == 92, sizeof(MD5state_st)
assert alignment(MD5state_st) == 4, alignment(MD5state_st)
MD5_CTX = MD5state_st
class mdc2_ctx_st(Structure):
    pass
mdc2_ctx_st._fields_ = [
    ('num', c_int),
    ('data', c_ubyte * 8),
    ('h', DES_cblock),
    ('hh', DES_cblock),
    ('pad_type', c_int),
]
assert sizeof(mdc2_ctx_st) == 32, sizeof(mdc2_ctx_st)
assert alignment(mdc2_ctx_st) == 4, alignment(mdc2_ctx_st)
MDC2_CTX = mdc2_ctx_st
class obj_name_st(Structure):
    pass
obj_name_st._fields_ = [
    ('type', c_int),
    ('alias', c_int),
    ('name', STRING),
    ('data', STRING),
]
assert sizeof(obj_name_st) == 16, sizeof(obj_name_st)
assert alignment(obj_name_st) == 4, alignment(obj_name_st)
OBJ_NAME = obj_name_st
ASN1_TIME = asn1_string_st
ASN1_NULL = c_int
EVP_PKEY = evp_pkey_st
class x509_st(Structure):
    pass
X509 = x509_st
class X509_algor_st(Structure):
    pass
X509_ALGOR = X509_algor_st
class X509_crl_st(Structure):
    pass
X509_CRL = X509_crl_st
class X509_name_st(Structure):
    pass
X509_NAME = X509_name_st
class x509_store_st(Structure):
    pass
X509_STORE = x509_store_st
class x509_store_ctx_st(Structure):
    pass
X509_STORE_CTX = x509_store_ctx_st
engine_st._fields_ = [
]
class PEM_Encode_Seal_st(Structure):
    pass
PEM_Encode_Seal_st._fields_ = [
    ('encode', EVP_ENCODE_CTX),
    ('md', EVP_MD_CTX),
    ('cipher', EVP_CIPHER_CTX),
]
assert sizeof(PEM_Encode_Seal_st) == 252, sizeof(PEM_Encode_Seal_st)
assert alignment(PEM_Encode_Seal_st) == 4, alignment(PEM_Encode_Seal_st)
PEM_ENCODE_SEAL_CTX = PEM_Encode_Seal_st
class pem_recip_st(Structure):
    pass
pem_recip_st._fields_ = [
    ('name', STRING),
    ('dn', POINTER(X509_NAME)),
    ('cipher', c_int),
    ('key_enc', c_int),
]
assert sizeof(pem_recip_st) == 16, sizeof(pem_recip_st)
assert alignment(pem_recip_st) == 4, alignment(pem_recip_st)
PEM_USER = pem_recip_st
class pem_ctx_st(Structure):
    pass
class N10pem_ctx_st4DOLLAR_16E(Structure):
    pass
N10pem_ctx_st4DOLLAR_16E._fields_ = [
    ('version', c_int),
    ('mode', c_int),
]
assert sizeof(N10pem_ctx_st4DOLLAR_16E) == 8, sizeof(N10pem_ctx_st4DOLLAR_16E)
assert alignment(N10pem_ctx_st4DOLLAR_16E) == 4, alignment(N10pem_ctx_st4DOLLAR_16E)
class N10pem_ctx_st4DOLLAR_17E(Structure):
    pass
N10pem_ctx_st4DOLLAR_17E._fields_ = [
    ('cipher', c_int),
]
assert sizeof(N10pem_ctx_st4DOLLAR_17E) == 4, sizeof(N10pem_ctx_st4DOLLAR_17E)
assert alignment(N10pem_ctx_st4DOLLAR_17E) == 4, alignment(N10pem_ctx_st4DOLLAR_17E)
pem_ctx_st._fields_ = [
    ('type', c_int),
    ('proc_type', N10pem_ctx_st4DOLLAR_16E),
    ('domain', STRING),
    ('DEK_info', N10pem_ctx_st4DOLLAR_17E),
    ('originator', POINTER(PEM_USER)),
    ('num_recipient', c_int),
    ('recipient', POINTER(POINTER(PEM_USER))),
    ('x509_chain', POINTER(STACK)),
    ('md', POINTER(EVP_MD)),
    ('md_enc', c_int),
    ('md_len', c_int),
    ('md_data', STRING),
    ('dec', POINTER(EVP_CIPHER)),
    ('key_len', c_int),
    ('key', POINTER(c_ubyte)),
    ('data_enc', c_int),
    ('data_len', c_int),
    ('data', POINTER(c_ubyte)),
]
assert sizeof(pem_ctx_st) == 76, sizeof(pem_ctx_st)
assert alignment(pem_ctx_st) == 4, alignment(pem_ctx_st)
PEM_CTX = pem_ctx_st
pem_password_cb = CFUNCTYPE(c_int, STRING, c_int, c_int, c_void_p)
class pkcs7_issuer_and_serial_st(Structure):
    pass
pkcs7_issuer_and_serial_st._fields_ = [
    ('issuer', POINTER(X509_NAME)),
    ('serial', POINTER(ASN1_INTEGER)),
]
assert sizeof(pkcs7_issuer_and_serial_st) == 8, sizeof(pkcs7_issuer_and_serial_st)
assert alignment(pkcs7_issuer_and_serial_st) == 4, alignment(pkcs7_issuer_and_serial_st)
PKCS7_ISSUER_AND_SERIAL = pkcs7_issuer_and_serial_st
class pkcs7_signer_info_st(Structure):
    pass
pkcs7_signer_info_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('issuer_and_serial', POINTER(PKCS7_ISSUER_AND_SERIAL)),
    ('digest_alg', POINTER(X509_ALGOR)),
    ('auth_attr', POINTER(STACK)),
    ('digest_enc_alg', POINTER(X509_ALGOR)),
    ('enc_digest', POINTER(ASN1_OCTET_STRING)),
    ('unauth_attr', POINTER(STACK)),
    ('pkey', POINTER(EVP_PKEY)),
]
assert sizeof(pkcs7_signer_info_st) == 32, sizeof(pkcs7_signer_info_st)
assert alignment(pkcs7_signer_info_st) == 4, alignment(pkcs7_signer_info_st)
PKCS7_SIGNER_INFO = pkcs7_signer_info_st
class pkcs7_recip_info_st(Structure):
    pass
pkcs7_recip_info_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('issuer_and_serial', POINTER(PKCS7_ISSUER_AND_SERIAL)),
    ('key_enc_algor', POINTER(X509_ALGOR)),
    ('enc_key', POINTER(ASN1_OCTET_STRING)),
    ('cert', POINTER(X509)),
]
assert sizeof(pkcs7_recip_info_st) == 20, sizeof(pkcs7_recip_info_st)
assert alignment(pkcs7_recip_info_st) == 4, alignment(pkcs7_recip_info_st)
PKCS7_RECIP_INFO = pkcs7_recip_info_st
class pkcs7_signed_st(Structure):
    pass
class pkcs7_st(Structure):
    pass
pkcs7_signed_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('md_algs', POINTER(STACK)),
    ('cert', POINTER(STACK)),
    ('crl', POINTER(STACK)),
    ('signer_info', POINTER(STACK)),
    ('contents', POINTER(pkcs7_st)),
]
assert sizeof(pkcs7_signed_st) == 24, sizeof(pkcs7_signed_st)
assert alignment(pkcs7_signed_st) == 4, alignment(pkcs7_signed_st)
PKCS7_SIGNED = pkcs7_signed_st
class pkcs7_enc_content_st(Structure):
    pass
pkcs7_enc_content_st._fields_ = [
    ('content_type', POINTER(ASN1_OBJECT)),
    ('algorithm', POINTER(X509_ALGOR)),
    ('enc_data', POINTER(ASN1_OCTET_STRING)),
    ('cipher', POINTER(EVP_CIPHER)),
]
assert sizeof(pkcs7_enc_content_st) == 16, sizeof(pkcs7_enc_content_st)
assert alignment(pkcs7_enc_content_st) == 4, alignment(pkcs7_enc_content_st)
PKCS7_ENC_CONTENT = pkcs7_enc_content_st
class pkcs7_enveloped_st(Structure):
    pass
pkcs7_enveloped_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('recipientinfo', POINTER(STACK)),
    ('enc_data', POINTER(PKCS7_ENC_CONTENT)),
]
assert sizeof(pkcs7_enveloped_st) == 12, sizeof(pkcs7_enveloped_st)
assert alignment(pkcs7_enveloped_st) == 4, alignment(pkcs7_enveloped_st)
PKCS7_ENVELOPE = pkcs7_enveloped_st
class pkcs7_signedandenveloped_st(Structure):
    pass
pkcs7_signedandenveloped_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('md_algs', POINTER(STACK)),
    ('cert', POINTER(STACK)),
    ('crl', POINTER(STACK)),
    ('signer_info', POINTER(STACK)),
    ('enc_data', POINTER(PKCS7_ENC_CONTENT)),
    ('recipientinfo', POINTER(STACK)),
]
assert sizeof(pkcs7_signedandenveloped_st) == 28, sizeof(pkcs7_signedandenveloped_st)
assert alignment(pkcs7_signedandenveloped_st) == 4, alignment(pkcs7_signedandenveloped_st)
PKCS7_SIGN_ENVELOPE = pkcs7_signedandenveloped_st
class pkcs7_digest_st(Structure):
    pass
pkcs7_digest_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('md', POINTER(X509_ALGOR)),
    ('contents', POINTER(pkcs7_st)),
    ('digest', POINTER(ASN1_OCTET_STRING)),
]
assert sizeof(pkcs7_digest_st) == 16, sizeof(pkcs7_digest_st)
assert alignment(pkcs7_digest_st) == 4, alignment(pkcs7_digest_st)
PKCS7_DIGEST = pkcs7_digest_st
class pkcs7_encrypted_st(Structure):
    pass
pkcs7_encrypted_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('enc_data', POINTER(PKCS7_ENC_CONTENT)),
]
assert sizeof(pkcs7_encrypted_st) == 8, sizeof(pkcs7_encrypted_st)
assert alignment(pkcs7_encrypted_st) == 4, alignment(pkcs7_encrypted_st)
PKCS7_ENCRYPT = pkcs7_encrypted_st
class N8pkcs7_st4DOLLAR_15E(Union):
    pass
N8pkcs7_st4DOLLAR_15E._fields_ = [
    ('ptr', STRING),
    ('data', POINTER(ASN1_OCTET_STRING)),
    ('sign', POINTER(PKCS7_SIGNED)),
    ('enveloped', POINTER(PKCS7_ENVELOPE)),
    ('signed_and_enveloped', POINTER(PKCS7_SIGN_ENVELOPE)),
    ('digest', POINTER(PKCS7_DIGEST)),
    ('encrypted', POINTER(PKCS7_ENCRYPT)),
    ('other', POINTER(ASN1_TYPE)),
]
assert sizeof(N8pkcs7_st4DOLLAR_15E) == 4, sizeof(N8pkcs7_st4DOLLAR_15E)
assert alignment(N8pkcs7_st4DOLLAR_15E) == 4, alignment(N8pkcs7_st4DOLLAR_15E)
pkcs7_st._fields_ = [
    ('asn1', POINTER(c_ubyte)),
    ('length', c_long),
    ('state', c_int),
    ('detached', c_int),
    ('type', POINTER(ASN1_OBJECT)),
    ('d', N8pkcs7_st4DOLLAR_15E),
]
assert sizeof(pkcs7_st) == 24, sizeof(pkcs7_st)
assert alignment(pkcs7_st) == 4, alignment(pkcs7_st)
PKCS7 = pkcs7_st
class rc2_key_st(Structure):
    pass
rc2_key_st._fields_ = [
    ('data', c_uint * 64),
]
assert sizeof(rc2_key_st) == 256, sizeof(rc2_key_st)
assert alignment(rc2_key_st) == 4, alignment(rc2_key_st)
RC2_KEY = rc2_key_st
class rc4_key_st(Structure):
    pass
rc4_key_st._fields_ = [
    ('x', c_ubyte),
    ('y', c_ubyte),
    ('data', c_ubyte * 256),
]
assert sizeof(rc4_key_st) == 258, sizeof(rc4_key_st)
assert alignment(rc4_key_st) == 1, alignment(rc4_key_st)
RC4_KEY = rc4_key_st
class rc5_key_st(Structure):
    pass
rc5_key_st._fields_ = [
    ('rounds', c_int),
    ('data', c_ulong * 34),
]
assert sizeof(rc5_key_st) == 140, sizeof(rc5_key_st)
assert alignment(rc5_key_st) == 4, alignment(rc5_key_st)
RC5_32_KEY = rc5_key_st
class RIPEMD160state_st(Structure):
    pass
RIPEMD160state_st._fields_ = [
    ('A', c_uint),
    ('B', c_uint),
    ('C', c_uint),
    ('D', c_uint),
    ('E', c_uint),
    ('Nl', c_uint),
    ('Nh', c_uint),
    ('data', c_uint * 16),
    ('num', c_int),
]
assert sizeof(RIPEMD160state_st) == 96, sizeof(RIPEMD160state_st)
assert alignment(RIPEMD160state_st) == 4, alignment(RIPEMD160state_st)
RIPEMD160_CTX = RIPEMD160state_st
RSA = rsa_st
class rsa_meth_st(Structure):
    pass
rsa_meth_st._fields_ = [
    ('name', STRING),
    ('rsa_pub_enc', CFUNCTYPE(c_int, c_int, POINTER(c_ubyte), POINTER(c_ubyte), POINTER(RSA), c_int)),
    ('rsa_pub_dec', CFUNCTYPE(c_int, c_int, POINTER(c_ubyte), POINTER(c_ubyte), POINTER(RSA), c_int)),
    ('rsa_priv_enc', CFUNCTYPE(c_int, c_int, POINTER(c_ubyte), POINTER(c_ubyte), POINTER(RSA), c_int)),
    ('rsa_priv_dec', CFUNCTYPE(c_int, c_int, POINTER(c_ubyte), POINTER(c_ubyte), POINTER(RSA), c_int)),
    ('rsa_mod_exp', CFUNCTYPE(c_int, POINTER(BIGNUM), POINTER(BIGNUM), POINTER(RSA))),
    ('bn_mod_exp', CFUNCTYPE(c_int, POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BIGNUM), POINTER(BN_CTX), POINTER(BN_MONT_CTX))),
    ('init', CFUNCTYPE(c_int, POINTER(RSA))),
    ('finish', CFUNCTYPE(c_int, POINTER(RSA))),
    ('flags', c_int),
    ('app_data', STRING),
    ('rsa_sign', CFUNCTYPE(c_int, c_int, POINTER(c_ubyte), c_uint, POINTER(c_ubyte), POINTER(c_uint), POINTER(RSA))),
    ('rsa_verify', CFUNCTYPE(c_int, c_int, POINTER(c_ubyte), c_uint, POINTER(c_ubyte), c_uint, POINTER(RSA))),
]
assert sizeof(rsa_meth_st) == 52, sizeof(rsa_meth_st)
assert alignment(rsa_meth_st) == 4, alignment(rsa_meth_st)
RSA_METHOD = rsa_meth_st
rsa_st._fields_ = [
    ('pad', c_int),
    ('version', c_long),
    ('meth', POINTER(RSA_METHOD)),
    ('engine', POINTER(ENGINE)),
    ('n', POINTER(BIGNUM)),
    ('e', POINTER(BIGNUM)),
    ('d', POINTER(BIGNUM)),
    ('p', POINTER(BIGNUM)),
    ('q', POINTER(BIGNUM)),
    ('dmp1', POINTER(BIGNUM)),
    ('dmq1', POINTER(BIGNUM)),
    ('iqmp', POINTER(BIGNUM)),
    ('ex_data', CRYPTO_EX_DATA),
    ('references', c_int),
    ('flags', c_int),
    ('_method_mod_n', POINTER(BN_MONT_CTX)),
    ('_method_mod_p', POINTER(BN_MONT_CTX)),
    ('_method_mod_q', POINTER(BN_MONT_CTX)),
    ('bignum_data', STRING),
    ('blinding', POINTER(BN_BLINDING)),
]
assert sizeof(rsa_st) == 84, sizeof(rsa_st)
assert alignment(rsa_st) == 4, alignment(rsa_st)
openssl_fptr = CFUNCTYPE(None)
class SHAstate_st(Structure):
    pass
SHAstate_st._fields_ = [
    ('h0', c_uint),
    ('h1', c_uint),
    ('h2', c_uint),
    ('h3', c_uint),
    ('h4', c_uint),
    ('Nl', c_uint),
    ('Nh', c_uint),
    ('data', c_uint * 16),
    ('num', c_int),
]
assert sizeof(SHAstate_st) == 96, sizeof(SHAstate_st)
assert alignment(SHAstate_st) == 4, alignment(SHAstate_st)
SHA_CTX = SHAstate_st
class ssl_st(Structure):
    pass
ssl_crock_st = POINTER(ssl_st)
class ssl_cipher_st(Structure):
    pass
ssl_cipher_st._fields_ = [
    ('valid', c_int),
    ('name', STRING),
    ('id', c_ulong),
    ('algorithms', c_ulong),
    ('algo_strength', c_ulong),
    ('algorithm2', c_ulong),
    ('strength_bits', c_int),
    ('alg_bits', c_int),
    ('mask', c_ulong),
    ('mask_strength', c_ulong),
]
assert sizeof(ssl_cipher_st) == 40, sizeof(ssl_cipher_st)
assert alignment(ssl_cipher_st) == 4, alignment(ssl_cipher_st)
SSL_CIPHER = ssl_cipher_st
SSL = ssl_st
class ssl_ctx_st(Structure):
    pass
SSL_CTX = ssl_ctx_st
class ssl_method_st(Structure):
    pass
class ssl3_enc_method(Structure):
    pass
ssl_method_st._fields_ = [
    ('version', c_int),
    ('ssl_new', CFUNCTYPE(c_int, POINTER(SSL))),
    ('ssl_clear', CFUNCTYPE(None, POINTER(SSL))),
    ('ssl_free', CFUNCTYPE(None, POINTER(SSL))),
    ('ssl_accept', CFUNCTYPE(c_int, POINTER(SSL))),
    ('ssl_connect', CFUNCTYPE(c_int, POINTER(SSL))),
    ('ssl_read', CFUNCTYPE(c_int, POINTER(SSL), c_void_p, c_int)),
    ('ssl_peek', CFUNCTYPE(c_int, POINTER(SSL), c_void_p, c_int)),
    ('ssl_write', CFUNCTYPE(c_int, POINTER(SSL), c_void_p, c_int)),
    ('ssl_shutdown', CFUNCTYPE(c_int, POINTER(SSL))),
    ('ssl_renegotiate', CFUNCTYPE(c_int, POINTER(SSL))),
    ('ssl_renegotiate_check', CFUNCTYPE(c_int, POINTER(SSL))),
    ('ssl_ctrl', CFUNCTYPE(c_long, POINTER(SSL), c_int, c_long, c_void_p)),
    ('ssl_ctx_ctrl', CFUNCTYPE(c_long, POINTER(SSL_CTX), c_int, c_long, c_void_p)),
    ('get_cipher_by_char', CFUNCTYPE(POINTER(SSL_CIPHER), POINTER(c_ubyte))),
    ('put_cipher_by_char', CFUNCTYPE(c_int, POINTER(SSL_CIPHER), POINTER(c_ubyte))),
    ('ssl_pending', CFUNCTYPE(c_int, POINTER(SSL))),
    ('num_ciphers', CFUNCTYPE(c_int)),
    ('get_cipher', CFUNCTYPE(POINTER(SSL_CIPHER), c_uint)),
    ('get_ssl_method', CFUNCTYPE(POINTER(ssl_method_st), c_int)),
    ('get_timeout', CFUNCTYPE(c_long)),
    ('ssl3_enc', POINTER(ssl3_enc_method)),
    ('ssl_version', CFUNCTYPE(c_int)),
    ('ssl_callback_ctrl', CFUNCTYPE(c_long, POINTER(SSL), c_int, CFUNCTYPE(None))),
    ('ssl_ctx_callback_ctrl', CFUNCTYPE(c_long, POINTER(SSL_CTX), c_int, CFUNCTYPE(None))),
]
assert sizeof(ssl_method_st) == 100, sizeof(ssl_method_st)
assert alignment(ssl_method_st) == 4, alignment(ssl_method_st)
ssl3_enc_method._fields_ = [
]
SSL_METHOD = ssl_method_st
class ssl_session_st(Structure):
    pass
class sess_cert_st(Structure):
    pass
ssl_session_st._fields_ = [
    ('ssl_version', c_int),
    ('key_arg_length', c_uint),
    ('key_arg', c_ubyte * 8),
    ('master_key_length', c_int),
    ('master_key', c_ubyte * 48),
    ('session_id_length', c_uint),
    ('session_id', c_ubyte * 32),
    ('sid_ctx_length', c_uint),
    ('sid_ctx', c_ubyte * 32),
    ('not_resumable', c_int),
    ('sess_cert', POINTER(sess_cert_st)),
    ('peer', POINTER(X509)),
    ('verify_result', c_long),
    ('references', c_int),
    ('timeout', c_long),
    ('time', c_long),
    ('compress_meth', c_int),
    ('cipher', POINTER(SSL_CIPHER)),
    ('cipher_id', c_ulong),
    ('ciphers', POINTER(STACK)),
    ('ex_data', CRYPTO_EX_DATA),
    ('prev', POINTER(ssl_session_st)),
    ('next', POINTER(ssl_session_st)),
]
assert sizeof(ssl_session_st) == 200, sizeof(ssl_session_st)
assert alignment(ssl_session_st) == 4, alignment(ssl_session_st)
sess_cert_st._fields_ = [
]
SSL_SESSION = ssl_session_st
GEN_SESSION_CB = CFUNCTYPE(c_int, POINTER(SSL), POINTER(c_ubyte), POINTER(c_uint))
class ssl_comp_st(Structure):
    pass
ssl_comp_st._fields_ = [
    ('id', c_int),
    ('name', STRING),
    ('method', POINTER(COMP_METHOD)),
]
assert sizeof(ssl_comp_st) == 12, sizeof(ssl_comp_st)
assert alignment(ssl_comp_st) == 4, alignment(ssl_comp_st)
SSL_COMP = ssl_comp_st
class N10ssl_ctx_st4DOLLAR_18E(Structure):
    pass
N10ssl_ctx_st4DOLLAR_18E._fields_ = [
    ('sess_connect', c_int),
    ('sess_connect_renegotiate', c_int),
    ('sess_connect_good', c_int),
    ('sess_accept', c_int),
    ('sess_accept_renegotiate', c_int),
    ('sess_accept_good', c_int),
    ('sess_miss', c_int),
    ('sess_timeout', c_int),
    ('sess_cache_full', c_int),
    ('sess_hit', c_int),
    ('sess_cb_hit', c_int),
]
assert sizeof(N10ssl_ctx_st4DOLLAR_18E) == 44, sizeof(N10ssl_ctx_st4DOLLAR_18E)
assert alignment(N10ssl_ctx_st4DOLLAR_18E) == 4, alignment(N10ssl_ctx_st4DOLLAR_18E)
class cert_st(Structure):
    pass
ssl_ctx_st._fields_ = [
    ('method', POINTER(SSL_METHOD)),
    ('cipher_list', POINTER(STACK)),
    ('cipher_list_by_id', POINTER(STACK)),
    ('cert_store', POINTER(x509_store_st)),
    ('sessions', POINTER(lhash_st)),
    ('session_cache_size', c_ulong),
    ('session_cache_head', POINTER(ssl_session_st)),
    ('session_cache_tail', POINTER(ssl_session_st)),
    ('session_cache_mode', c_int),
    ('session_timeout', c_long),
    ('new_session_cb', CFUNCTYPE(c_int, POINTER(ssl_st), POINTER(SSL_SESSION))),
    ('remove_session_cb', CFUNCTYPE(None, POINTER(ssl_ctx_st), POINTER(SSL_SESSION))),
    ('get_session_cb', CFUNCTYPE(POINTER(SSL_SESSION), POINTER(ssl_st), POINTER(c_ubyte), c_int, POINTER(c_int))),
    ('stats', N10ssl_ctx_st4DOLLAR_18E),
    ('references', c_int),
    ('app_verify_callback', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), c_void_p)),
    ('app_verify_arg', c_void_p),
    ('default_passwd_callback', POINTER(pem_password_cb)),
    ('default_passwd_callback_userdata', c_void_p),
    ('client_cert_cb', CFUNCTYPE(c_int, POINTER(SSL), POINTER(POINTER(X509)), POINTER(POINTER(EVP_PKEY)))),
    ('ex_data', CRYPTO_EX_DATA),
    ('rsa_md5', POINTER(EVP_MD)),
    ('md5', POINTER(EVP_MD)),
    ('sha1', POINTER(EVP_MD)),
    ('extra_certs', POINTER(STACK)),
    ('comp_methods', POINTER(STACK)),
    ('info_callback', CFUNCTYPE(None, POINTER(SSL), c_int, c_int)),
    ('client_CA', POINTER(STACK)),
    ('options', c_ulong),
    ('mode', c_ulong),
    ('max_cert_list', c_long),
    ('cert', POINTER(cert_st)),
    ('read_ahead', c_int),
    ('msg_callback', CFUNCTYPE(None, c_int, c_int, c_int, c_void_p, c_ulong, POINTER(SSL), c_void_p)),
    ('msg_callback_arg', c_void_p),
    ('verify_mode', c_int),
    ('verify_depth', c_int),
    ('sid_ctx_length', c_uint),
    ('sid_ctx', c_ubyte * 32),
    ('default_verify_callback', CFUNCTYPE(c_int, c_int, POINTER(X509_STORE_CTX))),
    ('generate_session_id', GEN_SESSION_CB),
    ('purpose', c_int),
    ('trust', c_int),
    ('quiet_shutdown', c_int),
]
assert sizeof(ssl_ctx_st) == 248, sizeof(ssl_ctx_st)
assert alignment(ssl_ctx_st) == 4, alignment(ssl_ctx_st)
cert_st._fields_ = [
]
class ssl2_state_st(Structure):
    pass
class ssl3_state_st(Structure):
    pass
ssl_st._fields_ = [
    ('version', c_int),
    ('type', c_int),
    ('method', POINTER(SSL_METHOD)),
    ('rbio', POINTER(BIO)),
    ('wbio', POINTER(BIO)),
    ('bbio', POINTER(BIO)),
    ('rwstate', c_int),
    ('in_handshake', c_int),
    ('handshake_func', CFUNCTYPE(c_int)),
    ('server', c_int),
    ('new_session', c_int),
    ('quiet_shutdown', c_int),
    ('shutdown', c_int),
    ('state', c_int),
    ('rstate', c_int),
    ('init_buf', POINTER(BUF_MEM)),
    ('init_msg', c_void_p),
    ('init_num', c_int),
    ('init_off', c_int),
    ('packet', POINTER(c_ubyte)),
    ('packet_length', c_uint),
    ('s2', POINTER(ssl2_state_st)),
    ('s3', POINTER(ssl3_state_st)),
    ('read_ahead', c_int),
    ('msg_callback', CFUNCTYPE(None, c_int, c_int, c_int, c_void_p, c_ulong, POINTER(SSL), c_void_p)),
    ('msg_callback_arg', c_void_p),
    ('hit', c_int),
    ('purpose', c_int),
    ('trust', c_int),
    ('cipher_list', POINTER(STACK)),
    ('cipher_list_by_id', POINTER(STACK)),
    ('enc_read_ctx', POINTER(EVP_CIPHER_CTX)),
    ('read_hash', POINTER(EVP_MD)),
    ('expand', POINTER(COMP_CTX)),
    ('enc_write_ctx', POINTER(EVP_CIPHER_CTX)),
    ('write_hash', POINTER(EVP_MD)),
    ('compress', POINTER(COMP_CTX)),
    ('cert', POINTER(cert_st)),
    ('sid_ctx_length', c_uint),
    ('sid_ctx', c_ubyte * 32),
    ('session', POINTER(SSL_SESSION)),
    ('generate_session_id', GEN_SESSION_CB),
    ('verify_mode', c_int),
    ('verify_depth', c_int),
    ('verify_callback', CFUNCTYPE(c_int, c_int, POINTER(X509_STORE_CTX))),
    ('info_callback', CFUNCTYPE(None, POINTER(SSL), c_int, c_int)),
    ('error', c_int),
    ('error_code', c_int),
    ('ctx', POINTER(SSL_CTX)),
    ('debug', c_int),
    ('verify_result', c_long),
    ('ex_data', CRYPTO_EX_DATA),
    ('client_CA', POINTER(STACK)),
    ('references', c_int),
    ('options', c_ulong),
    ('mode', c_ulong),
    ('max_cert_list', c_long),
    ('first_packet', c_int),
    ('client_version', c_int),
]
assert sizeof(ssl_st) == 268, sizeof(ssl_st)
assert alignment(ssl_st) == 4, alignment(ssl_st)
class N13ssl2_state_st4DOLLAR_19E(Structure):
    pass
N13ssl2_state_st4DOLLAR_19E._fields_ = [
    ('conn_id_length', c_uint),
    ('cert_type', c_uint),
    ('cert_length', c_uint),
    ('csl', c_uint),
    ('clear', c_uint),
    ('enc', c_uint),
    ('ccl', c_ubyte * 32),
    ('cipher_spec_length', c_uint),
    ('session_id_length', c_uint),
    ('clen', c_uint),
    ('rlen', c_uint),
]
assert sizeof(N13ssl2_state_st4DOLLAR_19E) == 72, sizeof(N13ssl2_state_st4DOLLAR_19E)
assert alignment(N13ssl2_state_st4DOLLAR_19E) == 4, alignment(N13ssl2_state_st4DOLLAR_19E)
ssl2_state_st._fields_ = [
    ('three_byte_header', c_int),
    ('clear_text', c_int),
    ('escape', c_int),
    ('ssl2_rollback', c_int),
    ('wnum', c_uint),
    ('wpend_tot', c_int),
    ('wpend_buf', POINTER(c_ubyte)),
    ('wpend_off', c_int),
    ('wpend_len', c_int),
    ('wpend_ret', c_int),
    ('rbuf_left', c_int),
    ('rbuf_offs', c_int),
    ('rbuf', POINTER(c_ubyte)),
    ('wbuf', POINTER(c_ubyte)),
    ('write_ptr', POINTER(c_ubyte)),
    ('padding', c_uint),
    ('rlength', c_uint),
    ('ract_data_length', c_int),
    ('wlength', c_uint),
    ('wact_data_length', c_int),
    ('ract_data', POINTER(c_ubyte)),
    ('wact_data', POINTER(c_ubyte)),
    ('mac_data', POINTER(c_ubyte)),
    ('read_key', POINTER(c_ubyte)),
    ('write_key', POINTER(c_ubyte)),
    ('challenge_length', c_uint),
    ('challenge', c_ubyte * 32),
    ('conn_id_length', c_uint),
    ('conn_id', c_ubyte * 16),
    ('key_material_length', c_uint),
    ('key_material', c_ubyte * 48),
    ('read_sequence', c_ulong),
    ('write_sequence', c_ulong),
    ('tmp', N13ssl2_state_st4DOLLAR_19E),
]
assert sizeof(ssl2_state_st) == 288, sizeof(ssl2_state_st)
assert alignment(ssl2_state_st) == 4, alignment(ssl2_state_st)
SSL2_STATE = ssl2_state_st
class ssl3_record_st(Structure):
    pass
ssl3_record_st._fields_ = [
    ('type', c_int),
    ('length', c_uint),
    ('off', c_uint),
    ('data', POINTER(c_ubyte)),
    ('input', POINTER(c_ubyte)),
    ('comp', POINTER(c_ubyte)),
]
assert sizeof(ssl3_record_st) == 24, sizeof(ssl3_record_st)
assert alignment(ssl3_record_st) == 4, alignment(ssl3_record_st)
SSL3_RECORD = ssl3_record_st
class ssl3_buffer_st(Structure):
    pass
size_t = __darwin_size_t
ssl3_buffer_st._fields_ = [
    ('buf', POINTER(c_ubyte)),
    ('len', size_t),
    ('offset', c_int),
    ('left', c_int),
]
assert sizeof(ssl3_buffer_st) == 16, sizeof(ssl3_buffer_st)
assert alignment(ssl3_buffer_st) == 4, alignment(ssl3_buffer_st)
SSL3_BUFFER = ssl3_buffer_st
class N13ssl3_state_st4DOLLAR_20E(Structure):
    pass
N13ssl3_state_st4DOLLAR_20E._fields_ = [
    ('cert_verify_md', c_ubyte * 72),
    ('finish_md', c_ubyte * 72),
    ('finish_md_len', c_int),
    ('peer_finish_md', c_ubyte * 72),
    ('peer_finish_md_len', c_int),
    ('message_size', c_ulong),
    ('message_type', c_int),
    ('new_cipher', POINTER(SSL_CIPHER)),
    ('dh', POINTER(DH)),
    ('next_state', c_int),
    ('reuse_message', c_int),
    ('cert_req', c_int),
    ('ctype_num', c_int),
    ('ctype', c_char * 7),
    ('ca_names', POINTER(STACK)),
    ('use_rsa_tmp', c_int),
    ('key_block_length', c_int),
    ('key_block', POINTER(c_ubyte)),
    ('new_sym_enc', POINTER(EVP_CIPHER)),
    ('new_hash', POINTER(EVP_MD)),
    ('new_compression', POINTER(SSL_COMP)),
    ('cert_request', c_int),
]
assert sizeof(N13ssl3_state_st4DOLLAR_20E) == 296, sizeof(N13ssl3_state_st4DOLLAR_20E)
assert alignment(N13ssl3_state_st4DOLLAR_20E) == 4, alignment(N13ssl3_state_st4DOLLAR_20E)
ssl3_state_st._fields_ = [
    ('flags', c_long),
    ('delay_buf_pop_ret', c_int),
    ('read_sequence', c_ubyte * 8),
    ('read_mac_secret', c_ubyte * 36),
    ('write_sequence', c_ubyte * 8),
    ('write_mac_secret', c_ubyte * 36),
    ('server_random', c_ubyte * 32),
    ('client_random', c_ubyte * 32),
    ('need_empty_fragments', c_int),
    ('empty_fragment_done', c_int),
    ('rbuf', SSL3_BUFFER),
    ('wbuf', SSL3_BUFFER),
    ('rrec', SSL3_RECORD),
    ('wrec', SSL3_RECORD),
    ('alert_fragment', c_ubyte * 2),
    ('alert_fragment_len', c_uint),
    ('handshake_fragment', c_ubyte * 4),
    ('handshake_fragment_len', c_uint),
    ('wnum', c_uint),
    ('wpend_tot', c_int),
    ('wpend_type', c_int),
    ('wpend_ret', c_int),
    ('wpend_buf', POINTER(c_ubyte)),
    ('finish_dgst1', EVP_MD_CTX),
    ('finish_dgst2', EVP_MD_CTX),
    ('change_cipher_spec', c_int),
    ('warn_alert', c_int),
    ('fatal_alert', c_int),
    ('alert_dispatch', c_int),
    ('send_alert', c_ubyte * 2),
    ('renegotiate', c_int),
    ('total_renegotiations', c_int),
    ('num_renegotiations', c_int),
    ('in_read_app_data', c_int),
    ('tmp', N13ssl3_state_st4DOLLAR_20E),
]
assert sizeof(ssl3_state_st) == 648, sizeof(ssl3_state_st)
assert alignment(ssl3_state_st) == 4, alignment(ssl3_state_st)
SSL3_STATE = ssl3_state_st
stack_st._fields_ = [
    ('num', c_int),
    ('data', POINTER(STRING)),
    ('sorted', c_int),
    ('num_alloc', c_int),
    ('comp', CFUNCTYPE(c_int, POINTER(STRING), POINTER(STRING))),
]
assert sizeof(stack_st) == 20, sizeof(stack_st)
assert alignment(stack_st) == 4, alignment(stack_st)
class ui_st(Structure):
    pass
ui_st._fields_ = [
]
UI = ui_st
class ui_method_st(Structure):
    pass
ui_method_st._fields_ = [
]
UI_METHOD = ui_method_st
class ui_string_st(Structure):
    pass
ui_string_st._fields_ = [
]
UI_STRING = ui_string_st

# values for enumeration 'UI_string_types'
UI_string_types = c_int # enum
class X509_objects_st(Structure):
    pass
X509_objects_st._fields_ = [
    ('nid', c_int),
    ('a2i', CFUNCTYPE(c_int)),
    ('i2a', CFUNCTYPE(c_int)),
]
assert sizeof(X509_objects_st) == 12, sizeof(X509_objects_st)
assert alignment(X509_objects_st) == 4, alignment(X509_objects_st)
X509_OBJECTS = X509_objects_st
X509_algor_st._fields_ = [
    ('algorithm', POINTER(ASN1_OBJECT)),
    ('parameter', POINTER(ASN1_TYPE)),
]
assert sizeof(X509_algor_st) == 8, sizeof(X509_algor_st)
assert alignment(X509_algor_st) == 4, alignment(X509_algor_st)
class X509_val_st(Structure):
    pass
X509_val_st._fields_ = [
    ('notBefore', POINTER(ASN1_TIME)),
    ('notAfter', POINTER(ASN1_TIME)),
]
assert sizeof(X509_val_st) == 8, sizeof(X509_val_st)
assert alignment(X509_val_st) == 4, alignment(X509_val_st)
X509_VAL = X509_val_st
class X509_pubkey_st(Structure):
    pass
X509_pubkey_st._fields_ = [
    ('algor', POINTER(X509_ALGOR)),
    ('public_key', POINTER(ASN1_BIT_STRING)),
    ('pkey', POINTER(EVP_PKEY)),
]
assert sizeof(X509_pubkey_st) == 12, sizeof(X509_pubkey_st)
assert alignment(X509_pubkey_st) == 4, alignment(X509_pubkey_st)
X509_PUBKEY = X509_pubkey_st
class X509_sig_st(Structure):
    pass
X509_sig_st._fields_ = [
    ('algor', POINTER(X509_ALGOR)),
    ('digest', POINTER(ASN1_OCTET_STRING)),
]
assert sizeof(X509_sig_st) == 8, sizeof(X509_sig_st)
assert alignment(X509_sig_st) == 4, alignment(X509_sig_st)
X509_SIG = X509_sig_st
class X509_name_entry_st(Structure):
    pass
X509_name_entry_st._fields_ = [
    ('object', POINTER(ASN1_OBJECT)),
    ('value', POINTER(ASN1_STRING)),
    ('set', c_int),
    ('size', c_int),
]
assert sizeof(X509_name_entry_st) == 16, sizeof(X509_name_entry_st)
assert alignment(X509_name_entry_st) == 4, alignment(X509_name_entry_st)
X509_NAME_ENTRY = X509_name_entry_st
X509_name_st._fields_ = [
    ('entries', POINTER(STACK)),
    ('modified', c_int),
    ('bytes', POINTER(BUF_MEM)),
    ('hash', c_ulong),
]
assert sizeof(X509_name_st) == 16, sizeof(X509_name_st)
assert alignment(X509_name_st) == 4, alignment(X509_name_st)
class X509_extension_st(Structure):
    pass
X509_extension_st._fields_ = [
    ('object', POINTER(ASN1_OBJECT)),
    ('critical', ASN1_BOOLEAN),
    ('value', POINTER(ASN1_OCTET_STRING)),
]
assert sizeof(X509_extension_st) == 12, sizeof(X509_extension_st)
assert alignment(X509_extension_st) == 4, alignment(X509_extension_st)
X509_EXTENSION = X509_extension_st
class x509_attributes_st(Structure):
    pass
class N18x509_attributes_st4DOLLAR_13E(Union):
    pass
N18x509_attributes_st4DOLLAR_13E._fields_ = [
    ('ptr', STRING),
    ('set', POINTER(STACK)),
    ('single', POINTER(ASN1_TYPE)),
]
assert sizeof(N18x509_attributes_st4DOLLAR_13E) == 4, sizeof(N18x509_attributes_st4DOLLAR_13E)
assert alignment(N18x509_attributes_st4DOLLAR_13E) == 4, alignment(N18x509_attributes_st4DOLLAR_13E)
x509_attributes_st._fields_ = [
    ('object', POINTER(ASN1_OBJECT)),
    ('single', c_int),
    ('value', N18x509_attributes_st4DOLLAR_13E),
]
assert sizeof(x509_attributes_st) == 12, sizeof(x509_attributes_st)
assert alignment(x509_attributes_st) == 4, alignment(x509_attributes_st)
X509_ATTRIBUTE = x509_attributes_st
class X509_req_info_st(Structure):
    pass
X509_req_info_st._fields_ = [
    ('enc', ASN1_ENCODING),
    ('version', POINTER(ASN1_INTEGER)),
    ('subject', POINTER(X509_NAME)),
    ('pubkey', POINTER(X509_PUBKEY)),
    ('attributes', POINTER(STACK)),
]
assert sizeof(X509_req_info_st) == 28, sizeof(X509_req_info_st)
assert alignment(X509_req_info_st) == 4, alignment(X509_req_info_st)
X509_REQ_INFO = X509_req_info_st
class X509_req_st(Structure):
    pass
X509_req_st._fields_ = [
    ('req_info', POINTER(X509_REQ_INFO)),
    ('sig_alg', POINTER(X509_ALGOR)),
    ('signature', POINTER(ASN1_BIT_STRING)),
    ('references', c_int),
]
assert sizeof(X509_req_st) == 16, sizeof(X509_req_st)
assert alignment(X509_req_st) == 4, alignment(X509_req_st)
X509_REQ = X509_req_st
class x509_cinf_st(Structure):
    pass
x509_cinf_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('serialNumber', POINTER(ASN1_INTEGER)),
    ('signature', POINTER(X509_ALGOR)),
    ('issuer', POINTER(X509_NAME)),
    ('validity', POINTER(X509_VAL)),
    ('subject', POINTER(X509_NAME)),
    ('key', POINTER(X509_PUBKEY)),
    ('issuerUID', POINTER(ASN1_BIT_STRING)),
    ('subjectUID', POINTER(ASN1_BIT_STRING)),
    ('extensions', POINTER(STACK)),
]
assert sizeof(x509_cinf_st) == 40, sizeof(x509_cinf_st)
assert alignment(x509_cinf_st) == 4, alignment(x509_cinf_st)
X509_CINF = x509_cinf_st
class x509_cert_aux_st(Structure):
    pass
x509_cert_aux_st._fields_ = [
    ('trust', POINTER(STACK)),
    ('reject', POINTER(STACK)),
    ('alias', POINTER(ASN1_UTF8STRING)),
    ('keyid', POINTER(ASN1_OCTET_STRING)),
    ('other', POINTER(STACK)),
]
assert sizeof(x509_cert_aux_st) == 20, sizeof(x509_cert_aux_st)
assert alignment(x509_cert_aux_st) == 4, alignment(x509_cert_aux_st)
X509_CERT_AUX = x509_cert_aux_st
class AUTHORITY_KEYID_st(Structure):
    pass
x509_st._fields_ = [
    ('cert_info', POINTER(X509_CINF)),
    ('sig_alg', POINTER(X509_ALGOR)),
    ('signature', POINTER(ASN1_BIT_STRING)),
    ('valid', c_int),
    ('references', c_int),
    ('name', STRING),
    ('ex_data', CRYPTO_EX_DATA),
    ('ex_pathlen', c_long),
    ('ex_flags', c_ulong),
    ('ex_kusage', c_ulong),
    ('ex_xkusage', c_ulong),
    ('ex_nscert', c_ulong),
    ('skid', POINTER(ASN1_OCTET_STRING)),
    ('akid', POINTER(AUTHORITY_KEYID_st)),
    ('sha1_hash', c_ubyte * 20),
    ('aux', POINTER(X509_CERT_AUX)),
]
assert sizeof(x509_st) == 84, sizeof(x509_st)
assert alignment(x509_st) == 4, alignment(x509_st)
AUTHORITY_KEYID_st._fields_ = [
]
class x509_trust_st(Structure):
    pass
x509_trust_st._fields_ = [
    ('trust', c_int),
    ('flags', c_int),
    ('check_trust', CFUNCTYPE(c_int, POINTER(x509_trust_st), POINTER(X509), c_int)),
    ('name', STRING),
    ('arg1', c_int),
    ('arg2', c_void_p),
]
assert sizeof(x509_trust_st) == 24, sizeof(x509_trust_st)
assert alignment(x509_trust_st) == 4, alignment(x509_trust_st)
X509_TRUST = x509_trust_st
class X509_revoked_st(Structure):
    pass
X509_revoked_st._fields_ = [
    ('serialNumber', POINTER(ASN1_INTEGER)),
    ('revocationDate', POINTER(ASN1_TIME)),
    ('extensions', POINTER(STACK)),
    ('sequence', c_int),
]
assert sizeof(X509_revoked_st) == 16, sizeof(X509_revoked_st)
assert alignment(X509_revoked_st) == 4, alignment(X509_revoked_st)
X509_REVOKED = X509_revoked_st
class X509_crl_info_st(Structure):
    pass
X509_crl_info_st._fields_ = [
    ('version', POINTER(ASN1_INTEGER)),
    ('sig_alg', POINTER(X509_ALGOR)),
    ('issuer', POINTER(X509_NAME)),
    ('lastUpdate', POINTER(ASN1_TIME)),
    ('nextUpdate', POINTER(ASN1_TIME)),
    ('revoked', POINTER(STACK)),
    ('extensions', POINTER(STACK)),
    ('enc', ASN1_ENCODING),
]
assert sizeof(X509_crl_info_st) == 40, sizeof(X509_crl_info_st)
assert alignment(X509_crl_info_st) == 4, alignment(X509_crl_info_st)
X509_CRL_INFO = X509_crl_info_st
X509_crl_st._fields_ = [
    ('crl', POINTER(X509_CRL_INFO)),
    ('sig_alg', POINTER(X509_ALGOR)),
    ('signature', POINTER(ASN1_BIT_STRING)),
    ('references', c_int),
]
assert sizeof(X509_crl_st) == 16, sizeof(X509_crl_st)
assert alignment(X509_crl_st) == 4, alignment(X509_crl_st)
class private_key_st(Structure):
    pass
private_key_st._fields_ = [
    ('version', c_int),
    ('enc_algor', POINTER(X509_ALGOR)),
    ('enc_pkey', POINTER(ASN1_OCTET_STRING)),
    ('dec_pkey', POINTER(EVP_PKEY)),
    ('key_length', c_int),
    ('key_data', STRING),
    ('key_free', c_int),
    ('cipher', EVP_CIPHER_INFO),
    ('references', c_int),
]
assert sizeof(private_key_st) == 52, sizeof(private_key_st)
assert alignment(private_key_st) == 4, alignment(private_key_st)
X509_PKEY = private_key_st
class X509_info_st(Structure):
    pass
X509_info_st._fields_ = [
    ('x509', POINTER(X509)),
    ('crl', POINTER(X509_CRL)),
    ('x_pkey', POINTER(X509_PKEY)),
    ('enc_cipher', EVP_CIPHER_INFO),
    ('enc_len', c_int),
    ('enc_data', STRING),
    ('references', c_int),
]
assert sizeof(X509_info_st) == 44, sizeof(X509_info_st)
assert alignment(X509_info_st) == 4, alignment(X509_info_st)
X509_INFO = X509_info_st
class Netscape_spkac_st(Structure):
    pass
Netscape_spkac_st._fields_ = [
    ('pubkey', POINTER(X509_PUBKEY)),
    ('challenge', POINTER(ASN1_IA5STRING)),
]
assert sizeof(Netscape_spkac_st) == 8, sizeof(Netscape_spkac_st)
assert alignment(Netscape_spkac_st) == 4, alignment(Netscape_spkac_st)
NETSCAPE_SPKAC = Netscape_spkac_st
class Netscape_spki_st(Structure):
    pass
Netscape_spki_st._fields_ = [
    ('spkac', POINTER(NETSCAPE_SPKAC)),
    ('sig_algor', POINTER(X509_ALGOR)),
    ('signature', POINTER(ASN1_BIT_STRING)),
]
assert sizeof(Netscape_spki_st) == 12, sizeof(Netscape_spki_st)
assert alignment(Netscape_spki_st) == 4, alignment(Netscape_spki_st)
NETSCAPE_SPKI = Netscape_spki_st
class Netscape_certificate_sequence(Structure):
    pass
Netscape_certificate_sequence._fields_ = [
    ('type', POINTER(ASN1_OBJECT)),
    ('certs', POINTER(STACK)),
]
assert sizeof(Netscape_certificate_sequence) == 8, sizeof(Netscape_certificate_sequence)
assert alignment(Netscape_certificate_sequence) == 4, alignment(Netscape_certificate_sequence)
NETSCAPE_CERT_SEQUENCE = Netscape_certificate_sequence
class PBEPARAM_st(Structure):
    pass
PBEPARAM_st._fields_ = [
    ('salt', POINTER(ASN1_OCTET_STRING)),
    ('iter', POINTER(ASN1_INTEGER)),
]
assert sizeof(PBEPARAM_st) == 8, sizeof(PBEPARAM_st)
assert alignment(PBEPARAM_st) == 4, alignment(PBEPARAM_st)
PBEPARAM = PBEPARAM_st
class PBE2PARAM_st(Structure):
    pass
PBE2PARAM_st._fields_ = [
    ('keyfunc', POINTER(X509_ALGOR)),
    ('encryption', POINTER(X509_ALGOR)),
]
assert sizeof(PBE2PARAM_st) == 8, sizeof(PBE2PARAM_st)
assert alignment(PBE2PARAM_st) == 4, alignment(PBE2PARAM_st)
PBE2PARAM = PBE2PARAM_st
class PBKDF2PARAM_st(Structure):
    pass
PBKDF2PARAM_st._fields_ = [
    ('salt', POINTER(ASN1_TYPE)),
    ('iter', POINTER(ASN1_INTEGER)),
    ('keylength', POINTER(ASN1_INTEGER)),
    ('prf', POINTER(X509_ALGOR)),
]
assert sizeof(PBKDF2PARAM_st) == 16, sizeof(PBKDF2PARAM_st)
assert alignment(PBKDF2PARAM_st) == 4, alignment(PBKDF2PARAM_st)
PBKDF2PARAM = PBKDF2PARAM_st
class pkcs8_priv_key_info_st(Structure):
    pass
pkcs8_priv_key_info_st._fields_ = [
    ('broken', c_int),
    ('version', POINTER(ASN1_INTEGER)),
    ('pkeyalg', POINTER(X509_ALGOR)),
    ('pkey', POINTER(ASN1_TYPE)),
    ('attributes', POINTER(STACK)),
]
assert sizeof(pkcs8_priv_key_info_st) == 20, sizeof(pkcs8_priv_key_info_st)
assert alignment(pkcs8_priv_key_info_st) == 4, alignment(pkcs8_priv_key_info_st)
PKCS8_PRIV_KEY_INFO = pkcs8_priv_key_info_st
class x509_hash_dir_st(Structure):
    pass
x509_hash_dir_st._fields_ = [
    ('num_dirs', c_int),
    ('dirs', POINTER(STRING)),
    ('dirs_type', POINTER(c_int)),
    ('num_dirs_alloced', c_int),
]
assert sizeof(x509_hash_dir_st) == 16, sizeof(x509_hash_dir_st)
assert alignment(x509_hash_dir_st) == 4, alignment(x509_hash_dir_st)
X509_HASH_DIR_CTX = x509_hash_dir_st
class x509_file_st(Structure):
    pass
x509_file_st._fields_ = [
    ('num_paths', c_int),
    ('num_alloced', c_int),
    ('paths', POINTER(STRING)),
    ('path_type', POINTER(c_int)),
]
assert sizeof(x509_file_st) == 16, sizeof(x509_file_st)
assert alignment(x509_file_st) == 4, alignment(x509_file_st)
X509_CERT_FILE_CTX = x509_file_st
class x509_object_st(Structure):
    pass
class N14x509_object_st4DOLLAR_14E(Union):
    pass
N14x509_object_st4DOLLAR_14E._fields_ = [
    ('ptr', STRING),
    ('x509', POINTER(X509)),
    ('crl', POINTER(X509_CRL)),
    ('pkey', POINTER(EVP_PKEY)),
]
assert sizeof(N14x509_object_st4DOLLAR_14E) == 4, sizeof(N14x509_object_st4DOLLAR_14E)
assert alignment(N14x509_object_st4DOLLAR_14E) == 4, alignment(N14x509_object_st4DOLLAR_14E)
x509_object_st._fields_ = [
    ('type', c_int),
    ('data', N14x509_object_st4DOLLAR_14E),
]
assert sizeof(x509_object_st) == 8, sizeof(x509_object_st)
assert alignment(x509_object_st) == 4, alignment(x509_object_st)
X509_OBJECT = x509_object_st
class x509_lookup_st(Structure):
    pass
X509_LOOKUP = x509_lookup_st
class x509_lookup_method_st(Structure):
    pass
x509_lookup_method_st._fields_ = [
    ('name', STRING),
    ('new_item', CFUNCTYPE(c_int, POINTER(X509_LOOKUP))),
    ('free', CFUNCTYPE(None, POINTER(X509_LOOKUP))),
    ('init', CFUNCTYPE(c_int, POINTER(X509_LOOKUP))),
    ('shutdown', CFUNCTYPE(c_int, POINTER(X509_LOOKUP))),
    ('ctrl', CFUNCTYPE(c_int, POINTER(X509_LOOKUP), c_int, STRING, c_long, POINTER(STRING))),
    ('get_by_subject', CFUNCTYPE(c_int, POINTER(X509_LOOKUP), c_int, POINTER(X509_NAME), POINTER(X509_OBJECT))),
    ('get_by_issuer_serial', CFUNCTYPE(c_int, POINTER(X509_LOOKUP), c_int, POINTER(X509_NAME), POINTER(ASN1_INTEGER), POINTER(X509_OBJECT))),
    ('get_by_fingerprint', CFUNCTYPE(c_int, POINTER(X509_LOOKUP), c_int, POINTER(c_ubyte), c_int, POINTER(X509_OBJECT))),
    ('get_by_alias', CFUNCTYPE(c_int, POINTER(X509_LOOKUP), c_int, STRING, c_int, POINTER(X509_OBJECT))),
]
assert sizeof(x509_lookup_method_st) == 40, sizeof(x509_lookup_method_st)
assert alignment(x509_lookup_method_st) == 4, alignment(x509_lookup_method_st)
X509_LOOKUP_METHOD = x509_lookup_method_st
x509_store_st._fields_ = [
    ('cache', c_int),
    ('objs', POINTER(STACK)),
    ('get_cert_methods', POINTER(STACK)),
    ('flags', c_ulong),
    ('purpose', c_int),
    ('trust', c_int),
    ('verify', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX))),
    ('verify_cb', CFUNCTYPE(c_int, c_int, POINTER(X509_STORE_CTX))),
    ('get_issuer', CFUNCTYPE(c_int, POINTER(POINTER(X509)), POINTER(X509_STORE_CTX), POINTER(X509))),
    ('check_issued', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(X509), POINTER(X509))),
    ('check_revocation', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX))),
    ('get_crl', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(POINTER(X509_CRL)), POINTER(X509))),
    ('check_crl', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(X509_CRL))),
    ('cert_crl', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(X509_CRL), POINTER(X509))),
    ('cleanup', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX))),
    ('ex_data', CRYPTO_EX_DATA),
    ('references', c_int),
    ('depth', c_int),
]
assert sizeof(x509_store_st) == 76, sizeof(x509_store_st)
assert alignment(x509_store_st) == 4, alignment(x509_store_st)
x509_lookup_st._fields_ = [
    ('init', c_int),
    ('skip', c_int),
    ('method', POINTER(X509_LOOKUP_METHOD)),
    ('method_data', STRING),
    ('store_ctx', POINTER(X509_STORE)),
]
assert sizeof(x509_lookup_st) == 20, sizeof(x509_lookup_st)
assert alignment(x509_lookup_st) == 4, alignment(x509_lookup_st)
time_t = __darwin_time_t
x509_store_ctx_st._fields_ = [
    ('ctx', POINTER(X509_STORE)),
    ('current_method', c_int),
    ('cert', POINTER(X509)),
    ('untrusted', POINTER(STACK)),
    ('purpose', c_int),
    ('trust', c_int),
    ('check_time', time_t),
    ('flags', c_ulong),
    ('other_ctx', c_void_p),
    ('verify', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX))),
    ('verify_cb', CFUNCTYPE(c_int, c_int, POINTER(X509_STORE_CTX))),
    ('get_issuer', CFUNCTYPE(c_int, POINTER(POINTER(X509)), POINTER(X509_STORE_CTX), POINTER(X509))),
    ('check_issued', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(X509), POINTER(X509))),
    ('check_revocation', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX))),
    ('get_crl', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(POINTER(X509_CRL)), POINTER(X509))),
    ('check_crl', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(X509_CRL))),
    ('cert_crl', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX), POINTER(X509_CRL), POINTER(X509))),
    ('cleanup', CFUNCTYPE(c_int, POINTER(X509_STORE_CTX))),
    ('depth', c_int),
    ('valid', c_int),
    ('last_untrusted', c_int),
    ('chain', POINTER(STACK)),
    ('error_depth', c_int),
    ('error', c_int),
    ('current_cert', POINTER(X509)),
    ('current_issuer', POINTER(X509)),
    ('current_crl', POINTER(X509_CRL)),
    ('ex_data', CRYPTO_EX_DATA),
]
assert sizeof(x509_store_ctx_st) == 116, sizeof(x509_store_ctx_st)
assert alignment(x509_store_ctx_st) == 4, alignment(x509_store_ctx_st)
va_list = __darwin_va_list
__darwin_off_t = __int64_t
fpos_t = __darwin_off_t
class __sbuf(Structure):
    pass
__sbuf._fields_ = [
    ('_base', POINTER(c_ubyte)),
    ('_size', c_int),
]
assert sizeof(__sbuf) == 8, sizeof(__sbuf)
assert alignment(__sbuf) == 4, alignment(__sbuf)
class __sFILEX(Structure):
    pass
__sFILEX._fields_ = [
]
class __sFILE(Structure):
    pass
__sFILE._pack_ = 4
__sFILE._fields_ = [
    ('_p', POINTER(c_ubyte)),
    ('_r', c_int),
    ('_w', c_int),
    ('_flags', c_short),
    ('_file', c_short),
    ('_bf', __sbuf),
    ('_lbfsize', c_int),
    ('_cookie', c_void_p),
    ('_close', CFUNCTYPE(c_int, c_void_p)),
    ('_read', CFUNCTYPE(c_int, c_void_p, STRING, c_int)),
    ('_seek', CFUNCTYPE(fpos_t, c_void_p, c_longlong, c_int)),
    ('_write', CFUNCTYPE(c_int, c_void_p, STRING, c_int)),
    ('_ub', __sbuf),
    ('_extra', POINTER(__sFILEX)),
    ('_ur', c_int),
    ('_ubuf', c_ubyte * 3),
    ('_nbuf', c_ubyte * 1),
    ('_lb', __sbuf),
    ('_blksize', c_int),
    ('_offset', fpos_t),
]
assert sizeof(__sFILE) == 88, sizeof(__sFILE)
assert alignment(__sFILE) == 4, alignment(__sFILE)
FILE = __sFILE
ct_rune_t = __darwin_ct_rune_t
rune_t = __darwin_rune_t
class div_t(Structure):
    pass
div_t._fields_ = [
    ('quot', c_int),
    ('rem', c_int),
]
assert sizeof(div_t) == 8, sizeof(div_t)
assert alignment(div_t) == 4, alignment(div_t)
class ldiv_t(Structure):
    pass
ldiv_t._fields_ = [
    ('quot', c_long),
    ('rem', c_long),
]
assert sizeof(ldiv_t) == 8, sizeof(ldiv_t)
assert alignment(ldiv_t) == 4, alignment(ldiv_t)
class lldiv_t(Structure):
    pass
lldiv_t._pack_ = 4
lldiv_t._fields_ = [
    ('quot', c_longlong),
    ('rem', c_longlong),
]
assert sizeof(lldiv_t) == 16, sizeof(lldiv_t)
assert alignment(lldiv_t) == 4, alignment(lldiv_t)
__darwin_dev_t = __int32_t
dev_t = __darwin_dev_t
__darwin_mode_t = __uint16_t
mode_t = __darwin_mode_t
class mcontext(Structure):
    pass
mcontext._fields_ = [
]
class mcontext64(Structure):
    pass
mcontext64._fields_ = [
]
class __darwin_pthread_handler_rec(Structure):
    pass
__darwin_pthread_handler_rec._fields_ = [
    ('__routine', CFUNCTYPE(None, c_void_p)),
    ('__arg', c_void_p),
    ('__next', POINTER(__darwin_pthread_handler_rec)),
]
assert sizeof(__darwin_pthread_handler_rec) == 12, sizeof(__darwin_pthread_handler_rec)
assert alignment(__darwin_pthread_handler_rec) == 4, alignment(__darwin_pthread_handler_rec)
class _opaque_pthread_attr_t(Structure):
    pass
_opaque_pthread_attr_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 36),
]
assert sizeof(_opaque_pthread_attr_t) == 40, sizeof(_opaque_pthread_attr_t)
assert alignment(_opaque_pthread_attr_t) == 4, alignment(_opaque_pthread_attr_t)
class _opaque_pthread_cond_t(Structure):
    pass
_opaque_pthread_cond_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 24),
]
assert sizeof(_opaque_pthread_cond_t) == 28, sizeof(_opaque_pthread_cond_t)
assert alignment(_opaque_pthread_cond_t) == 4, alignment(_opaque_pthread_cond_t)
class _opaque_pthread_condattr_t(Structure):
    pass
_opaque_pthread_condattr_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 4),
]
assert sizeof(_opaque_pthread_condattr_t) == 8, sizeof(_opaque_pthread_condattr_t)
assert alignment(_opaque_pthread_condattr_t) == 4, alignment(_opaque_pthread_condattr_t)
class _opaque_pthread_mutex_t(Structure):
    pass
_opaque_pthread_mutex_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 40),
]
assert sizeof(_opaque_pthread_mutex_t) == 44, sizeof(_opaque_pthread_mutex_t)
assert alignment(_opaque_pthread_mutex_t) == 4, alignment(_opaque_pthread_mutex_t)
class _opaque_pthread_mutexattr_t(Structure):
    pass
_opaque_pthread_mutexattr_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 8),
]
assert sizeof(_opaque_pthread_mutexattr_t) == 12, sizeof(_opaque_pthread_mutexattr_t)
assert alignment(_opaque_pthread_mutexattr_t) == 4, alignment(_opaque_pthread_mutexattr_t)
class _opaque_pthread_once_t(Structure):
    pass
_opaque_pthread_once_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 4),
]
assert sizeof(_opaque_pthread_once_t) == 8, sizeof(_opaque_pthread_once_t)
assert alignment(_opaque_pthread_once_t) == 4, alignment(_opaque_pthread_once_t)
class _opaque_pthread_rwlock_t(Structure):
    pass
_opaque_pthread_rwlock_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 124),
]
assert sizeof(_opaque_pthread_rwlock_t) == 128, sizeof(_opaque_pthread_rwlock_t)
assert alignment(_opaque_pthread_rwlock_t) == 4, alignment(_opaque_pthread_rwlock_t)
class _opaque_pthread_rwlockattr_t(Structure):
    pass
_opaque_pthread_rwlockattr_t._fields_ = [
    ('__sig', c_long),
    ('__opaque', c_char * 12),
]
assert sizeof(_opaque_pthread_rwlockattr_t) == 16, sizeof(_opaque_pthread_rwlockattr_t)
assert alignment(_opaque_pthread_rwlockattr_t) == 4, alignment(_opaque_pthread_rwlockattr_t)
class _opaque_pthread_t(Structure):
    pass
_opaque_pthread_t._fields_ = [
    ('__sig', c_long),
    ('__cleanup_stack', POINTER(__darwin_pthread_handler_rec)),
    ('__opaque', c_char * 596),
]
assert sizeof(_opaque_pthread_t) == 604, sizeof(_opaque_pthread_t)
assert alignment(_opaque_pthread_t) == 4, alignment(_opaque_pthread_t)
__darwin_blkcnt_t = __int64_t
__darwin_blksize_t = __int32_t
__darwin_fsblkcnt_t = c_uint
__darwin_fsfilcnt_t = c_uint
__darwin_gid_t = __uint32_t
__darwin_id_t = __uint32_t
__darwin_ino_t = __uint32_t
__darwin_mach_port_name_t = __darwin_natural_t
__darwin_mach_port_t = __darwin_mach_port_name_t
__darwin_mcontext_t = POINTER(mcontext)
__darwin_mcontext64_t = POINTER(mcontext64)
__darwin_pid_t = __int32_t
__darwin_pthread_attr_t = _opaque_pthread_attr_t
__darwin_pthread_cond_t = _opaque_pthread_cond_t
__darwin_pthread_condattr_t = _opaque_pthread_condattr_t
__darwin_pthread_key_t = c_ulong
__darwin_pthread_mutex_t = _opaque_pthread_mutex_t
__darwin_pthread_mutexattr_t = _opaque_pthread_mutexattr_t
__darwin_pthread_once_t = _opaque_pthread_once_t
__darwin_pthread_rwlock_t = _opaque_pthread_rwlock_t
__darwin_pthread_rwlockattr_t = _opaque_pthread_rwlockattr_t
__darwin_pthread_t = POINTER(_opaque_pthread_t)
__darwin_sigset_t = __uint32_t
__darwin_suseconds_t = __int32_t
__darwin_uid_t = __uint32_t
__darwin_useconds_t = __uint32_t
__darwin_uuid_t = c_ubyte * 16
class sigaltstack(Structure):
    pass
sigaltstack._fields_ = [
    ('ss_sp', c_void_p),
    ('ss_size', __darwin_size_t),
    ('ss_flags', c_int),
]
assert sizeof(sigaltstack) == 12, sizeof(sigaltstack)
assert alignment(sigaltstack) == 4, alignment(sigaltstack)
__darwin_stack_t = sigaltstack
class ucontext(Structure):
    pass
ucontext._fields_ = [
    ('uc_onstack', c_int),
    ('uc_sigmask', __darwin_sigset_t),
    ('uc_stack', __darwin_stack_t),
    ('uc_link', POINTER(ucontext)),
    ('uc_mcsize', __darwin_size_t),
    ('uc_mcontext', __darwin_mcontext_t),
]
assert sizeof(ucontext) == 32, sizeof(ucontext)
assert alignment(ucontext) == 4, alignment(ucontext)
__darwin_ucontext_t = ucontext
class ucontext64(Structure):
    pass
ucontext64._fields_ = [
    ('uc_onstack', c_int),
    ('uc_sigmask', __darwin_sigset_t),
    ('uc_stack', __darwin_stack_t),
    ('uc_link', POINTER(ucontext64)),
    ('uc_mcsize', __darwin_size_t),
    ('uc_mcontext64', __darwin_mcontext64_t),
]
assert sizeof(ucontext64) == 32, sizeof(ucontext64)
assert alignment(ucontext64) == 4, alignment(ucontext64)
__darwin_ucontext64_t = ucontext64
class timeval(Structure):
    pass
timeval._fields_ = [
    ('tv_sec', __darwin_time_t),
    ('tv_usec', __darwin_suseconds_t),
]
assert sizeof(timeval) == 8, sizeof(timeval)
assert alignment(timeval) == 4, alignment(timeval)
rlim_t = __int64_t
class rusage(Structure):
    pass
rusage._fields_ = [
    ('ru_utime', timeval),
    ('ru_stime', timeval),
    ('ru_maxrss', c_long),
    ('ru_ixrss', c_long),
    ('ru_idrss', c_long),
    ('ru_isrss', c_long),
    ('ru_minflt', c_long),
    ('ru_majflt', c_long),
    ('ru_nswap', c_long),
    ('ru_inblock', c_long),
    ('ru_oublock', c_long),
    ('ru_msgsnd', c_long),
    ('ru_msgrcv', c_long),
    ('ru_nsignals', c_long),
    ('ru_nvcsw', c_long),
    ('ru_nivcsw', c_long),
]
assert sizeof(rusage) == 72, sizeof(rusage)
assert alignment(rusage) == 4, alignment(rusage)
class rlimit(Structure):
    pass
rlimit._pack_ = 4
rlimit._fields_ = [
    ('rlim_cur', rlim_t),
    ('rlim_max', rlim_t),
]
assert sizeof(rlimit) == 16, sizeof(rlimit)
assert alignment(rlimit) == 4, alignment(rlimit)
mcontext_t = __darwin_mcontext_t
mcontext64_t = __darwin_mcontext64_t
pthread_attr_t = __darwin_pthread_attr_t
sigset_t = __darwin_sigset_t
ucontext_t = __darwin_ucontext_t
ucontext64_t = __darwin_ucontext64_t
uid_t = __darwin_uid_t
class sigval(Union):
    pass
sigval._fields_ = [
    ('sival_int', c_int),
    ('sival_ptr', c_void_p),
]
assert sizeof(sigval) == 4, sizeof(sigval)
assert alignment(sigval) == 4, alignment(sigval)
class sigevent(Structure):
    pass
sigevent._fields_ = [
    ('sigev_notify', c_int),
    ('sigev_signo', c_int),
    ('sigev_value', sigval),
    ('sigev_notify_function', CFUNCTYPE(None, sigval)),
    ('sigev_notify_attributes', POINTER(pthread_attr_t)),
]
assert sizeof(sigevent) == 20, sizeof(sigevent)
assert alignment(sigevent) == 4, alignment(sigevent)
class __siginfo(Structure):
    pass
pid_t = __darwin_pid_t
__siginfo._fields_ = [
    ('si_signo', c_int),
    ('si_errno', c_int),
    ('si_code', c_int),
    ('si_pid', pid_t),
    ('si_uid', uid_t),
    ('si_status', c_int),
    ('si_addr', c_void_p),
    ('si_value', sigval),
    ('si_band', c_long),
    ('pad', c_ulong * 7),
]
assert sizeof(__siginfo) == 64, sizeof(__siginfo)
assert alignment(__siginfo) == 4, alignment(__siginfo)
siginfo_t = __siginfo
class __sigaction_u(Union):
    pass
__sigaction_u._fields_ = [
    ('__sa_handler', CFUNCTYPE(None, c_int)),
    ('__sa_sigaction', CFUNCTYPE(None, c_int, POINTER(__siginfo), c_void_p)),
]
assert sizeof(__sigaction_u) == 4, sizeof(__sigaction_u)
assert alignment(__sigaction_u) == 4, alignment(__sigaction_u)
class __sigaction(Structure):
    pass
__sigaction._fields_ = [
    ('__sigaction_u', __sigaction_u),
    ('sa_tramp', CFUNCTYPE(None, c_void_p, c_int, c_int, POINTER(siginfo_t), c_void_p)),
    ('sa_mask', sigset_t),
    ('sa_flags', c_int),
]
assert sizeof(__sigaction) == 16, sizeof(__sigaction)
assert alignment(__sigaction) == 4, alignment(__sigaction)
class sigaction(Structure):
    pass
sigaction._fields_ = [
    ('__sigaction_u', __sigaction_u),
    ('sa_mask', sigset_t),
    ('sa_flags', c_int),
]
assert sizeof(sigaction) == 12, sizeof(sigaction)
assert alignment(sigaction) == 4, alignment(sigaction)
sig_t = CFUNCTYPE(None, c_int)
stack_t = __darwin_stack_t
class sigvec(Structure):
    pass
sigvec._fields_ = [
    ('sv_handler', CFUNCTYPE(None, c_int)),
    ('sv_mask', c_int),
    ('sv_flags', c_int),
]
assert sizeof(sigvec) == 12, sizeof(sigvec)
assert alignment(sigvec) == 4, alignment(sigvec)
class sigstack(Structure):
    pass
sigstack._fields_ = [
    ('ss_sp', STRING),
    ('ss_onstack', c_int),
]
assert sizeof(sigstack) == 8, sizeof(sigstack)
assert alignment(sigstack) == 4, alignment(sigstack)
u_char = c_ubyte
u_short = c_ushort
u_int = c_uint
u_long = c_ulong
ushort = c_ushort
uint = c_uint
u_quad_t = u_int64_t
quad_t = int64_t
qaddr_t = POINTER(quad_t)
caddr_t = STRING
daddr_t = int32_t
fixpt_t = u_int32_t
blkcnt_t = __darwin_blkcnt_t
blksize_t = __darwin_blksize_t
gid_t = __darwin_gid_t
in_addr_t = __uint32_t
in_port_t = __uint16_t
ino_t = __darwin_ino_t
key_t = __int32_t
nlink_t = __uint16_t
off_t = __darwin_off_t
segsz_t = int32_t
swblk_t = int32_t
clock_t = __darwin_clock_t
ssize_t = __darwin_ssize_t
useconds_t = __darwin_useconds_t
suseconds_t = __darwin_suseconds_t
fd_mask = __int32_t
class fd_set(Structure):
    pass
fd_set._fields_ = [
    ('fds_bits', __int32_t * 32),
]
assert sizeof(fd_set) == 128, sizeof(fd_set)
assert alignment(fd_set) == 4, alignment(fd_set)
pthread_cond_t = __darwin_pthread_cond_t
pthread_condattr_t = __darwin_pthread_condattr_t
pthread_mutex_t = __darwin_pthread_mutex_t
pthread_mutexattr_t = __darwin_pthread_mutexattr_t
pthread_once_t = __darwin_pthread_once_t
pthread_rwlock_t = __darwin_pthread_rwlock_t
pthread_rwlockattr_t = __darwin_pthread_rwlockattr_t
pthread_t = __darwin_pthread_t
pthread_key_t = __darwin_pthread_key_t
fsblkcnt_t = __darwin_fsblkcnt_t
fsfilcnt_t = __darwin_fsfilcnt_t

# values for enumeration 'idtype_t'
idtype_t = c_int # enum
id_t = __darwin_id_t
class wait(Union):
    pass
class N4wait3DOLLAR_3E(Structure):
    pass
N4wait3DOLLAR_3E._fields_ = [
    ('w_Termsig', c_uint, 7),
    ('w_Coredump', c_uint, 1),
    ('w_Retcode', c_uint, 8),
    ('w_Filler', c_uint, 16),
]
assert sizeof(N4wait3DOLLAR_3E) == 4, sizeof(N4wait3DOLLAR_3E)
assert alignment(N4wait3DOLLAR_3E) == 4, alignment(N4wait3DOLLAR_3E)
class N4wait3DOLLAR_4E(Structure):
    pass
N4wait3DOLLAR_4E._fields_ = [
    ('w_Stopval', c_uint, 8),
    ('w_Stopsig', c_uint, 8),
    ('w_Filler', c_uint, 16),
]
assert sizeof(N4wait3DOLLAR_4E) == 4, sizeof(N4wait3DOLLAR_4E)
assert alignment(N4wait3DOLLAR_4E) == 4, alignment(N4wait3DOLLAR_4E)
wait._fields_ = [
    ('w_status', c_int),
    ('w_T', N4wait3DOLLAR_3E),
    ('w_S', N4wait3DOLLAR_4E),
]
assert sizeof(wait) == 4, sizeof(wait)
assert alignment(wait) == 4, alignment(wait)
class timespec(Structure):
    pass
timespec._fields_ = [
    ('tv_sec', time_t),
    ('tv_nsec', c_long),
]
assert sizeof(timespec) == 8, sizeof(timespec)
assert alignment(timespec) == 4, alignment(timespec)
class tm(Structure):
    pass
tm._fields_ = [
    ('tm_sec', c_int),
    ('tm_min', c_int),
    ('tm_hour', c_int),
    ('tm_mday', c_int),
    ('tm_mon', c_int),
    ('tm_year', c_int),
    ('tm_wday', c_int),
    ('tm_yday', c_int),
    ('tm_isdst', c_int),
    ('tm_gmtoff', c_long),
    ('tm_zone', STRING),
]
assert sizeof(tm) == 44, sizeof(tm)
assert alignment(tm) == 4, alignment(tm)
__gnuc_va_list = STRING
ptrdiff_t = c_int
int8_t = c_byte
int16_t = c_short
uint8_t = c_ubyte
uint16_t = c_ushort
uint32_t = c_uint
uint64_t = c_ulonglong
int_least8_t = int8_t
int_least16_t = int16_t
int_least32_t = int32_t
int_least64_t = int64_t
uint_least8_t = uint8_t
uint_least16_t = uint16_t
uint_least32_t = uint32_t
uint_least64_t = uint64_t
int_fast8_t = int8_t
int_fast16_t = int16_t
int_fast32_t = int32_t
int_fast64_t = int64_t
uint_fast8_t = uint8_t
uint_fast16_t = uint16_t
uint_fast32_t = uint32_t
uint_fast64_t = uint64_t
intptr_t = c_long
uintptr_t = c_ulong
intmax_t = c_longlong
uintmax_t = c_ulonglong
__all__ = ['ENGINE', 'pkcs7_enc_content_st', '__int16_t',
           'X509_REVOKED', 'SSL_CTX', 'UIT_BOOLEAN',
           '__darwin_time_t', 'ucontext64_t', 'int_fast32_t',
           'pem_ctx_st', 'uint8_t', 'fpos_t', 'X509', 'COMP_CTX',
           'tm', 'N10pem_ctx_st4DOLLAR_17E', 'swblk_t',
           'ASN1_TEMPLATE', '__darwin_pthread_t', 'fixpt_t',
           'BIO_METHOD', 'ASN1_PRINTABLESTRING', 'EVP_ENCODE_CTX',
           'dh_method', 'bio_f_buffer_ctx_struct', 'in_port_t',
           'X509_SIG', '__darwin_ssize_t', '__darwin_sigset_t',
           'wait', 'uint_fast16_t', 'N12asn1_type_st4DOLLAR_11E',
           'uint_least8_t', 'pthread_rwlock_t', 'ASN1_IA5STRING',
           'fsfilcnt_t', 'ucontext', '__uint64_t', 'timespec',
           'x509_cinf_st', 'COMP_METHOD', 'MD5_CTX', 'buf_mem_st',
           'ASN1_ENCODING_st', 'PBEPARAM', 'X509_NAME_ENTRY',
           '__darwin_va_list', 'ucontext_t', 'lhash_st',
           'N4wait3DOLLAR_4E', '__darwin_uuid_t',
           '_ossl_old_des_ks_struct', 'id_t', 'ASN1_BIT_STRING',
           'va_list', '__darwin_wchar_t', 'pthread_key_t',
           'pkcs7_signer_info_st', 'ASN1_METHOD', 'DSA_SIG', 'DSA',
           'UIT_NONE', 'pthread_t', '__darwin_useconds_t',
           'uint_fast8_t', 'UI_STRING', 'DES_cblock',
           '__darwin_mcontext64_t', 'rlim_t', 'PEM_Encode_Seal_st',
           'SHAstate_st', 'u_quad_t', 'openssl_fptr',
           '_opaque_pthread_rwlockattr_t',
           'N18x509_attributes_st4DOLLAR_13E',
           '__darwin_pthread_rwlock_t', 'daddr_t', 'ui_string_st',
           'x509_file_st', 'X509_req_info_st', 'int_least64_t',
           'evp_Encode_Ctx_st', 'X509_OBJECTS', 'CRYPTO_EX_DATA',
           '__int8_t', 'AUTHORITY_KEYID_st', '_opaque_pthread_attr_t',
           'sigstack', 'EVP_CIPHER_CTX', 'X509_extension_st', 'pid_t',
           'RSA_METHOD', 'PEM_USER', 'pem_recip_st', 'env_md_ctx_st',
           'rc5_key_st', 'ui_st', 'X509_PUBKEY', 'u_int8_t',
           'ASN1_ITEM_st', 'pkcs7_recip_info_st', 'ssl2_state_st',
           'off_t', 'N10ssl_ctx_st4DOLLAR_18E', 'crypto_ex_data_st',
           'ui_method_st', '__darwin_pthread_rwlockattr_t',
           'CRYPTO_EX_dup', '__darwin_ino_t', '__sFILE',
           'OSUnknownByteOrder', 'BN_MONT_CTX', 'ASN1_NULL', 'time_t',
           'CRYPTO_EX_new', 'asn1_type_st', 'CRYPTO_EX_DATA_FUNCS',
           'user_time_t', 'BIGNUM', 'pthread_rwlockattr_t',
           'ASN1_VALUE_st', 'DH_METHOD', '__darwin_off_t',
           '_opaque_pthread_t', 'bn_blinding_st', 'RSA', 'ssize_t',
           'mcontext64_t', 'user_long_t', 'fsblkcnt_t', 'cert_st',
           '__darwin_pthread_condattr_t', 'X509_PKEY',
           '__darwin_id_t', '__darwin_nl_item', 'SSL2_STATE', 'FILE',
           'pthread_mutexattr_t', 'size_t',
           '_ossl_old_des_key_schedule', 'pkcs7_issuer_and_serial_st',
           'sigval', 'CRYPTO_MEM_LEAK_CB', 'X509_NAME', 'blkcnt_t',
           'uint_least16_t', '__darwin_dev_t', 'evp_cipher_info_st',
           'BN_BLINDING', 'ssl3_state_st', 'uint_least64_t',
           'user_addr_t', 'DES_key_schedule', 'RIPEMD160_CTX',
           'u_char', 'X509_algor_st', 'uid_t', 'sess_cert_st',
           'u_int64_t', 'u_int16_t', 'sigset_t', '__darwin_ptrdiff_t',
           'ASN1_CTX', 'STACK', '__int32_t', 'UI_METHOD',
           'NETSCAPE_SPKI', 'UIT_PROMPT', 'st_CRYPTO_EX_DATA_IMPL',
           'cast_key_st', 'X509_HASH_DIR_CTX', 'sigevent',
           'user_ssize_t', 'clock_t', 'aes_key_st',
           '__darwin_socklen_t', '__darwin_intptr_t', 'int_fast64_t',
           'asn1_string_table_st', 'uint_fast32_t',
           'ASN1_VISIBLESTRING', 'DSA_SIG_st', 'obj_name_st',
           'X509_LOOKUP_METHOD', 'u_int32_t', 'EVP_CIPHER_INFO',
           '__gnuc_va_list', 'AES_KEY', 'PKCS7_ISSUER_AND_SERIAL',
           'BN_CTX', '__darwin_blkcnt_t', 'key_t', 'SHA_CTX',
           'pkcs7_signed_st', 'SSL', 'N10pem_ctx_st4DOLLAR_16E',
           'pthread_attr_t', 'EVP_MD', 'uint', 'ASN1_BOOLEAN',
           'ino_t', '__darwin_clock_t', 'ASN1_OCTET_STRING',
           'asn1_ctx_st', 'BIO_F_BUFFER_CTX', 'bn_mont_ctx_st',
           'X509_REQ_INFO', 'PEM_CTX', 'sigvec',
           '__darwin_pthread_mutexattr_t', 'x509_attributes_st',
           'stack_t', '__darwin_mode_t', '__mbstate_t',
           'asn1_object_st', 'ASN1_ENCODING', '__uint8_t',
           'LHASH_NODE', 'PKCS7_SIGNER_INFO', 'asn1_method_st',
           'stack_st', 'bio_info_cb', 'div_t', 'UIT_VERIFY',
           'PBEPARAM_st', 'N4wait3DOLLAR_3E', 'quad_t', '__siginfo',
           '__darwin_mbstate_t', 'rsa_st', 'ASN1_UNIVERSALSTRING',
           'uint64_t', 'ssl_comp_st', 'X509_OBJECT', 'pthread_cond_t',
           'DH', '__darwin_wctype_t', 'PKCS7_ENVELOPE', 'ASN1_TLC_st',
           'sig_atomic_t', 'BIO', 'nlink_t', 'BUF_MEM', 'SSL3_RECORD',
           'bio_method_st', 'timeval', 'UI_string_types', 'BIO_dummy',
           'ssl_ctx_st', 'NETSCAPE_CERT_SEQUENCE',
           'BIT_STRING_BITNAME_st', '__darwin_pthread_attr_t',
           'int8_t', '__darwin_wint_t', 'OBJ_NAME',
           'PKCS8_PRIV_KEY_INFO', 'PBE2PARAM_st',
           'LHASH_DOALL_FN_TYPE', 'x509_st', 'X509_VAL', 'dev_t',
           'ASN1_TEMPLATE_st', 'MD5state_st', '__uint16_t',
           'LHASH_DOALL_ARG_FN_TYPE', 'mdc2_ctx_st', 'SSL3_STATE',
           'ssl3_buffer_st', 'ASN1_ITEM_EXP',
           '_opaque_pthread_condattr_t', 'mode_t', 'ASN1_VALUE',
           'qaddr_t', '__darwin_gid_t', 'EVP_PKEY', 'CRYPTO_EX_free',
           '_ossl_old_des_cblock', 'X509_INFO', 'asn1_string_st',
           'intptr_t', 'UIT_INFO', 'int_fast8_t', 'sigaltstack',
           'env_md_st', 'LHASH', '__darwin_ucontext_t',
           'PKCS7_SIGN_ENVELOPE', '__darwin_mcontext_t', 'ct_rune_t',
           'MD2_CTX', 'pthread_once_t', 'SSL3_BUFFER', 'fd_mask',
           'ASN1_TYPE', 'PKCS7_SIGNED', 'ssl3_record_st', 'BF_KEY',
           'MD4state_st', 'MD4_CTX', 'int16_t', 'SSL_CIPHER',
           'rune_t', 'X509_TRUST', 'siginfo_t', 'X509_STORE',
           '__sbuf', 'X509_STORE_CTX', '__darwin_blksize_t', 'ldiv_t',
           'ASN1_TIME', 'SSL_METHOD', 'X509_LOOKUP',
           'Netscape_spki_st', 'P_PID', 'sigaction', 'sig_t',
           'hostent', 'x509_cert_aux_st', '_opaque_pthread_cond_t',
           'segsz_t', 'ushort', '__darwin_ct_rune_t', 'fd_set',
           'BN_RECP_CTX', 'x509_lookup_st', 'uint16_t', 'pkcs7_st',
           'asn1_header_st', '__darwin_pthread_key_t',
           'x509_trust_st', '__darwin_pthread_handler_rec', 'int32_t',
           'X509_CRL_INFO', 'N11evp_pkey_st4DOLLAR_12E', 'MDC2_CTX',
           'N23_ossl_old_des_ks_struct4DOLLAR_10E', 'ASN1_HEADER',
           'X509_crl_info_st', 'LHASH_HASH_FN_TYPE',
           '_opaque_pthread_mutexattr_t', 'ssl_st',
           'N8pkcs7_st4DOLLAR_15E', 'evp_pkey_st',
           'pkcs7_signedandenveloped_st', '__darwin_mach_port_t',
           'EVP_PBE_KEYGEN', '_opaque_pthread_mutex_t',
           'ASN1_UTCTIME', 'mcontext', 'crypto_ex_data_func_st',
           'u_long', 'PBKDF2PARAM_st', 'rc4_key_st', 'DSA_METHOD',
           'EVP_CIPHER', 'BIT_STRING_BITNAME', 'PKCS7_RECIP_INFO',
           'ssl3_enc_method', 'X509_CERT_AUX', 'uintmax_t',
           'int_fast16_t', 'RC5_32_KEY', 'ucontext64', 'ASN1_INTEGER',
           'u_short', 'N14x509_object_st4DOLLAR_14E', 'mcontext64',
           'X509_sig_st', 'ASN1_GENERALSTRING', 'PKCS7', '__sFILEX',
           'X509_name_entry_st', 'ssl_session_st', 'caddr_t',
           'bignum_st', 'X509_CINF', '__darwin_pthread_cond_t',
           'ASN1_TLC', 'PKCS7_ENCRYPT', 'NETSCAPE_SPKAC',
           'Netscape_spkac_st', 'idtype_t', 'UIT_ERROR',
           'uint_fast64_t', 'in_addr_t', 'pthread_mutex_t',
           '__int64_t', 'ASN1_BMPSTRING', 'uint32_t',
           'PEM_ENCODE_SEAL_CTX', 'suseconds_t', 'ASN1_OBJECT',
           'X509_val_st', 'private_key_st', 'CRYPTO_dynlock',
           'X509_objects_st', 'CRYPTO_EX_DATA_IMPL',
           'pthread_condattr_t', 'PKCS7_DIGEST', 'uint_least32_t',
           'ASN1_STRING', '__uint32_t', 'P_PGID', 'rsa_meth_st',
           'X509_crl_st', 'RC2_KEY', '__darwin_fsfilcnt_t',
           'X509_revoked_st', 'PBE2PARAM', 'blksize_t',
           'Netscape_certificate_sequence', 'ssl_cipher_st',
           'bignum_ctx', 'register_t', 'ASN1_UTF8STRING',
           'pkcs7_encrypted_st', 'RC4_KEY', '__darwin_ucontext64_t',
           'N13ssl2_state_st4DOLLAR_19E', 'bn_recp_ctx_st',
           'CAST_KEY', 'X509_ATTRIBUTE', '__darwin_suseconds_t',
           '__sigaction', 'user_ulong_t', 'syscall_arg_t',
           'evp_cipher_ctx_st', 'X509_ALGOR', 'mcontext_t',
           'const_DES_cblock', '__darwin_fsblkcnt_t', 'dsa_st',
           'int_least8_t', 'MD2state_st', 'X509_EXTENSION',
           'GEN_SESSION_CB', 'int_least16_t', '__darwin_wctrans_t',
           'PBKDF2PARAM', 'x509_lookup_method_st', 'pem_password_cb',
           'X509_info_st', 'x509_store_st', '__darwin_natural_t',
           'X509_pubkey_st', 'pkcs7_digest_st', '__darwin_size_t',
           'ASN1_STRING_TABLE', 'OSLittleEndian', 'RIPEMD160state_st',
           'pkcs7_enveloped_st', 'UI', 'ptrdiff_t', 'X509_REQ',
           'CRYPTO_dynlock_value', 'X509_req_st', 'x509_store_ctx_st',
           'N13ssl3_state_st4DOLLAR_20E', 'lhash_node_st',
           '__darwin_pthread_mutex_t', 'LHASH_COMP_FN_TYPE',
           '__darwin_rune_t', 'rlimit', '__darwin_pthread_once_t',
           'OSBigEndian', 'uintptr_t', '__darwin_uid_t', 'u_int',
           'ASN1_T61STRING', 'gid_t', 'ssl_method_st', 'ASN1_ITEM',
           'ASN1_ENUMERATED', '_opaque_pthread_rwlock_t',
           'pkcs8_priv_key_info_st', 'intmax_t', 'sigcontext',
           'X509_CRL', 'rc2_key_st', 'engine_st', 'x509_object_st',
           '_opaque_pthread_once_t', 'DES_ks', 'SSL_COMP',
           'dsa_method', 'int64_t', 'bio_st', 'bf_key_st',
           'ASN1_GENERALIZEDTIME', 'PKCS7_ENC_CONTENT',
           '__darwin_pid_t', 'lldiv_t', 'comp_method_st',
           'EVP_MD_CTX', 'evp_cipher_st', 'X509_name_st',
           'x509_hash_dir_st', '__darwin_mach_port_name_t',
           'useconds_t', 'user_size_t', 'SSL_SESSION', 'rusage',
           'ssl_crock_st', 'int_least32_t', '__sigaction_u', 'dh_st',
           'P_ALL', '__darwin_stack_t', 'N6DES_ks3DOLLAR_9E',
           'comp_ctx_st', 'X509_CERT_FILE_CTX']

external/cpython_master/Lib/lib2to3/tests/data/infinite_recursion.py
DragonHPC/dragon
#
# Analogue of `multiprocessing.connection` which uses queues instead of sockets
#
# multiprocessing/dummy/connection.py
#
# Copyright (c) 2006-2008, R Oudkerk
# Licensed to PSF under a Contributor Agreement.
#

__all__ = [ 'Client', 'Listener', 'Pipe' ]

from queue import Queue


families = [None]


class Listener(object):

    def __init__(self, address=None, family=None, backlog=1):
        self._backlog_queue = Queue(backlog)

    def accept(self):
        return Connection(*self._backlog_queue.get())

    def close(self):
        self._backlog_queue = None

    @property
    def address(self):
        return self._backlog_queue

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.close()


def Client(address):
    _in, _out = Queue(), Queue()
    address.put((_out, _in))
    return Connection(_in, _out)


def Pipe(duplex=True):
    a, b = Queue(), Queue()
    return Connection(a, b), Connection(b, a)


class Connection(object):

    def __init__(self, _in, _out):
        self._out = _out
        self._in = _in
        self.send = self.send_bytes = _out.put
        self.recv = self.recv_bytes = _in.get

    def poll(self, timeout=0.0):
        if self._in.qsize() > 0:
            return True
        if timeout <= 0.0:
            return False
        with self._in.not_empty:
            self._in.not_empty.wait(timeout)
        return self._in.qsize() > 0

    def close(self):
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_tb):
        self.close()

external/cpython_master/Lib/multiprocessing/dummy/connection.py
DragonHPC/dragon
#-*- coding: iso-8859-1 -*-
# pysqlite2/test/transactions.py: tests transactions
#
# Copyright (C) 2005-2007 Gerhard Häring <gh@ghaering.de>
#
# This file is part of pysqlite.
#
# This software is provided 'as-is', without any express or implied
# warranty.  In no event will the authors be held liable for any damages
# arising from the use of this software.
#
# Permission is granted to anyone to use this software for any purpose,
# including commercial applications, and to alter it and redistribute it
# freely, subject to the following restrictions:
#
# 1. The origin of this software must not be misrepresented; you must not
#    claim that you wrote the original software. If you use this software
#    in a product, an acknowledgment in the product documentation would be
#    appreciated but is not required.
# 2. Altered source versions must be plainly marked as such, and must not be
#    misrepresented as being the original software.
# 3. This notice may not be removed or altered from any source distribution.

import os, unittest
import sqlite3 as sqlite

def get_db_path():
    return "sqlite_testdb"

class TransactionTests(unittest.TestCase):
    def setUp(self):
        try:
            os.remove(get_db_path())
        except OSError:
            pass

        self.con1 = sqlite.connect(get_db_path(), timeout=0.1)
        self.cur1 = self.con1.cursor()

        self.con2 = sqlite.connect(get_db_path(), timeout=0.1)
        self.cur2 = self.con2.cursor()

    def tearDown(self):
        self.cur1.close()
        self.con1.close()

        self.cur2.close()
        self.con2.close()

        try:
            os.unlink(get_db_path())
        except OSError:
            pass

    def CheckDMLDoesNotAutoCommitBefore(self):
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        self.cur1.execute("create table test2(j)")
        self.cur2.execute("select i from test")
        res = self.cur2.fetchall()
        self.assertEqual(len(res), 0)

    def CheckInsertStartsTransaction(self):
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        self.cur2.execute("select i from test")
        res = self.cur2.fetchall()
        self.assertEqual(len(res), 0)

    def CheckUpdateStartsTransaction(self):
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        self.con1.commit()
        self.cur1.execute("update test set i=6")
        self.cur2.execute("select i from test")
        res = self.cur2.fetchone()[0]
        self.assertEqual(res, 5)

    def CheckDeleteStartsTransaction(self):
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        self.con1.commit()
        self.cur1.execute("delete from test")
        self.cur2.execute("select i from test")
        res = self.cur2.fetchall()
        self.assertEqual(len(res), 1)

    def CheckReplaceStartsTransaction(self):
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        self.con1.commit()
        self.cur1.execute("replace into test(i) values (6)")
        self.cur2.execute("select i from test")
        res = self.cur2.fetchall()
        self.assertEqual(len(res), 1)
        self.assertEqual(res[0][0], 5)

    def CheckToggleAutoCommit(self):
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        self.con1.isolation_level = None
        self.assertEqual(self.con1.isolation_level, None)
        self.cur2.execute("select i from test")
        res = self.cur2.fetchall()
        self.assertEqual(len(res), 1)

        self.con1.isolation_level = "DEFERRED"
        self.assertEqual(self.con1.isolation_level , "DEFERRED")
        self.cur1.execute("insert into test(i) values (5)")
        self.cur2.execute("select i from test")
        res = self.cur2.fetchall()
        self.assertEqual(len(res), 1)

    @unittest.skipIf(sqlite.sqlite_version_info < (3, 2, 2),
                     'test hangs on sqlite versions older than 3.2.2')
    def CheckRaiseTimeout(self):
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        with self.assertRaises(sqlite.OperationalError):
            self.cur2.execute("insert into test(i) values (5)")

    @unittest.skipIf(sqlite.sqlite_version_info < (3, 2, 2),
                     'test hangs on sqlite versions older than 3.2.2')
    def CheckLocking(self):
        """
        This tests the improved concurrency with pysqlite 2.3.4. You needed
        to roll back con2 before you could commit con1.
        """
        self.cur1.execute("create table test(i)")
        self.cur1.execute("insert into test(i) values (5)")
        with self.assertRaises(sqlite.OperationalError):
            self.cur2.execute("insert into test(i) values (5)")
        # NO self.con2.rollback() HERE!!!
        self.con1.commit()

    def CheckRollbackCursorConsistency(self):
        """
        Checks if cursors on the connection are set into a "reset" state
        when a rollback is done on the connection.
        """
        con = sqlite.connect(":memory:")
        cur = con.cursor()
        cur.execute("create table test(x)")
        cur.execute("insert into test(x) values (5)")
        cur.execute("select 1 union select 2 union select 3")

        con.rollback()
        with self.assertRaises(sqlite.InterfaceError):
            cur.fetchall()

class SpecialCommandTests(unittest.TestCase):
    def setUp(self):
        self.con = sqlite.connect(":memory:")
        self.cur = self.con.cursor()

    def CheckDropTable(self):
        self.cur.execute("create table test(i)")
        self.cur.execute("insert into test(i) values (5)")
        self.cur.execute("drop table test")

    def CheckPragma(self):
        self.cur.execute("create table test(i)")
        self.cur.execute("insert into test(i) values (5)")
        self.cur.execute("pragma count_changes=1")

    def tearDown(self):
        self.cur.close()
        self.con.close()

class TransactionalDDL(unittest.TestCase):
    def setUp(self):
        self.con = sqlite.connect(":memory:")

    def CheckDdlDoesNotAutostartTransaction(self):
        # For backwards compatibility reasons, DDL statements should not
        # implicitly start a transaction.
        self.con.execute("create table test(i)")
        self.con.rollback()
        result = self.con.execute("select * from test").fetchall()
        self.assertEqual(result, [])

    def CheckImmediateTransactionalDDL(self):
        # You can achieve transactional DDL by issuing a BEGIN
        # statement manually.
        self.con.execute("begin immediate")
        self.con.execute("create table test(i)")
        self.con.rollback()
        with self.assertRaises(sqlite.OperationalError):
            self.con.execute("select * from test")

    def CheckTransactionalDDL(self):
        # You can achieve transactional DDL by issuing a BEGIN
        # statement manually.
        self.con.execute("begin")
        self.con.execute("create table test(i)")
        self.con.rollback()
        with self.assertRaises(sqlite.OperationalError):
            self.con.execute("select * from test")

    def tearDown(self):
        self.con.close()

def suite():
    default_suite = unittest.makeSuite(TransactionTests, "Check")
    special_command_suite = unittest.makeSuite(SpecialCommandTests, "Check")
    ddl_suite = unittest.makeSuite(TransactionalDDL, "Check")
    return unittest.TestSuite((default_suite, special_command_suite, ddl_suite))

def test():
    runner = unittest.TextTestRunner()
    runner.run(suite())

if __name__ == "__main__":
    test()

external/cpython_master/Lib/sqlite3/test/transactions.py
DragonHPC/dragon
# This should be equivalent to running regrtest.py from the cmdline.
# It can be especially handy if you're in an interactive shell, e.g.,
# from test import autotest.
from test.libregrtest import main
main()

external/cpython_master/Lib/test/autotest.py
DragonHPC/dragon
import unittest
import __future__

GOOD_SERIALS = ("alpha", "beta", "candidate", "final")

features = __future__.all_feature_names

class FutureTest(unittest.TestCase):

    def test_names(self):
        # Verify that all_feature_names appears correct.
        given_feature_names = features[:]
        for name in dir(__future__):
            obj = getattr(__future__, name, None)
            if obj is not None and isinstance(obj, __future__._Feature):
                self.assertTrue(
                    name in given_feature_names,
                    "%r should have been in all_feature_names" % name
                )
                given_feature_names.remove(name)
        self.assertEqual(len(given_feature_names), 0,
               "all_feature_names has too much: %r" % given_feature_names)

    def test_attributes(self):
        for feature in features:
            value = getattr(__future__, feature)

            optional = value.getOptionalRelease()
            mandatory = value.getMandatoryRelease()

            a = self.assertTrue
            e = self.assertEqual
            def check(t, name):
                a(isinstance(t, tuple), "%s isn't tuple" % name)
                e(len(t), 5, "%s isn't 5-tuple" % name)
                (major, minor, micro, level, serial) = t
                a(isinstance(major, int), "%s major isn't int"  % name)
                a(isinstance(minor, int), "%s minor isn't int" % name)
                a(isinstance(micro, int), "%s micro isn't int" % name)
                a(isinstance(level, str),
                    "%s level isn't string" % name)
                a(level in GOOD_SERIALS,
                       "%s level string has unknown value" % name)
                a(isinstance(serial, int), "%s serial isn't int" % name)

            check(optional, "optional")
            if mandatory is not None:
                check(mandatory, "mandatory")
                a(optional < mandatory,
                       "optional not less than mandatory, and mandatory not None")

            a(hasattr(value, "compiler_flag"),
                   "feature is missing a .compiler_flag attr")
            # Make sure the compile accepts the flag.
            compile("", "<test>", "exec", value.compiler_flag)
            a(isinstance(getattr(value, "compiler_flag"), int),
                   ".compiler_flag isn't int")


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test___future__.py
DragonHPC/dragon
import os
from test.support import load_package_tests, import_module

# Skip tests if we don't have concurrent.futures.
import_module('concurrent.futures')

def load_tests(*args):
    return load_package_tests(os.path.dirname(__file__), *args)

external/cpython_master/Lib/test/test_asyncio/__init__.py
DragonHPC/dragon
import unittest
from test import support

import io # C implementation.
import _pyio as pyio # Python implementation.

# Simple test to ensure that optimizations in the IO library deliver the
# expected results.  For best testing, run this under a debug-build Python too
# (to exercise asserts in the C code).

lengths = list(range(1, 257)) + [512, 1000, 1024, 2048, 4096, 8192, 10000,
                                 16384, 32768, 65536, 1000000]

class BufferSizeTest:
    def try_one(self, s):
        # Write s + "\n" + s to file, then open it and ensure that successive
        # .readline()s deliver what we wrote.

        # Ensure we can open TESTFN for writing.
        support.unlink(support.TESTFN)

        # Since C doesn't guarantee we can write/read arbitrary bytes in text
        # files, use binary mode.
        f = self.open(support.TESTFN, "wb")
        try:
            # write once with \n and once without
            f.write(s)
            f.write(b"\n")
            f.write(s)
            f.close()
            f = open(support.TESTFN, "rb")
            line = f.readline()
            self.assertEqual(line, s + b"\n")
            line = f.readline()
            self.assertEqual(line, s)
            line = f.readline()
            self.assertFalse(line) # Must be at EOF
            f.close()
        finally:
            support.unlink(support.TESTFN)

    def drive_one(self, pattern):
        for length in lengths:
            # Repeat string 'pattern' as often as needed to reach total length
            # 'length'.  Then call try_one with that string, a string one larger
            # than that, and a string one smaller than that.  Try this with all
            # small sizes and various powers of 2, so we exercise all likely
            # stdio buffer sizes, and "off by one" errors on both sides.
            q, r = divmod(length, len(pattern))
            teststring = pattern * q + pattern[:r]
            self.assertEqual(len(teststring), length)
            self.try_one(teststring)
            self.try_one(teststring + b"x")
            self.try_one(teststring[:-1])

    def test_primepat(self):
        # A pattern with prime length, to avoid simple relationships with
        # stdio buffer sizes.
        self.drive_one(b"1234567890\00\01\02\03\04\05\06")

    def test_nullpat(self):
        self.drive_one(b'\0' * 1000)


class CBufferSizeTest(BufferSizeTest, unittest.TestCase):
    open = io.open

class PyBufferSizeTest(BufferSizeTest, unittest.TestCase):
    open = staticmethod(pyio.open)


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_bufio.py
DragonHPC/dragon
"""
Test script for the 'cmd' module
Original by Michael Schneider
"""


import cmd
import sys
import unittest
import io
from test import support

class samplecmdclass(cmd.Cmd):
    """
    Instance the sampleclass:
    >>> mycmd = samplecmdclass()

    Test for the function parseline():
    >>> mycmd.parseline("")
    (None, None, '')
    >>> mycmd.parseline("?")
    ('help', '', 'help ')
    >>> mycmd.parseline("?help")
    ('help', 'help', 'help help')
    >>> mycmd.parseline("!")
    ('shell', '', 'shell ')
    >>> mycmd.parseline("!command")
    ('shell', 'command', 'shell command')
    >>> mycmd.parseline("func")
    ('func', '', 'func')
    >>> mycmd.parseline("func arg1")
    ('func', 'arg1', 'func arg1')


    Test for the function onecmd():
    >>> mycmd.onecmd("")
    >>> mycmd.onecmd("add 4 5")
    9
    >>> mycmd.onecmd("")
    9
    >>> mycmd.onecmd("test")
    *** Unknown syntax: test

    Test for the function emptyline():
    >>> mycmd.emptyline()
    *** Unknown syntax: test

    Test for the function default():
    >>> mycmd.default("default")
    *** Unknown syntax: default

    Test for the function completedefault():
    >>> mycmd.completedefault()
    This is the completedefault method
    >>> mycmd.completenames("a")
    ['add']

    Test for the function completenames():
    >>> mycmd.completenames("12")
    []
    >>> mycmd.completenames("help")
    ['help']

    Test for the function complete_help():
    >>> mycmd.complete_help("a")
    ['add']
    >>> mycmd.complete_help("he")
    ['help']
    >>> mycmd.complete_help("12")
    []
    >>> sorted(mycmd.complete_help(""))
    ['add', 'exit', 'help', 'shell']

    Test for the function do_help():
    >>> mycmd.do_help("testet")
    *** No help on testet
    >>> mycmd.do_help("add")
    help text for add
    >>> mycmd.onecmd("help add")
    help text for add
    >>> mycmd.do_help("")
    <BLANKLINE>
    Documented commands (type help <topic>):
    ========================================
    add  help
    <BLANKLINE>
    Undocumented commands:
    ======================
    exit  shell
    <BLANKLINE>

    Test for the function print_topics():
    >>> mycmd.print_topics("header", ["command1", "command2"], 2 ,10)
    header
    ======
    command1
    command2
    <BLANKLINE>

    Test for the function columnize():
    >>> mycmd.columnize([str(i) for i in range(20)])
    0  1  2  3  4  5  6  7  8  9  10  11  12  13  14  15  16  17  18  19
    >>> mycmd.columnize([str(i) for i in range(20)], 10)
    0  7   14
    1  8   15
    2  9   16
    3  10  17
    4  11  18
    5  12  19
    6  13

    This is an interactive test, put some commands in the cmdqueue attribute
    and let it execute
    This test includes the preloop(), postloop(), default(), emptyline(),
    parseline(), do_help() functions
    >>> mycmd.use_rawinput=0
    >>> mycmd.cmdqueue=["", "add", "add 4 5", "help", "help add","exit"]
    >>> mycmd.cmdloop()
    Hello from preloop
    help text for add
    *** invalid number of arguments
    9
    <BLANKLINE>
    Documented commands (type help <topic>):
    ========================================
    add  help
    <BLANKLINE>
    Undocumented commands:
    ======================
    exit  shell
    <BLANKLINE>
    help text for add
    Hello from postloop
    """

    def preloop(self):
        print("Hello from preloop")

    def postloop(self):
        print("Hello from postloop")

    def completedefault(self, *ignored):
        print("This is the completedefault method")

    def complete_command(self):
        print("complete command")

    def do_shell(self, s):
        pass

    def do_add(self, s):
        l = s.split()
        if len(l) != 2:
            print("*** invalid number of arguments")
            return
        try:
            l = [int(i) for i in l]
        except ValueError:
            print("*** arguments should be numbers")
            return
        print(l[0]+l[1])

    def help_add(self):
        print("help text for add")
        return

    def do_exit(self, arg):
        return True


class TestAlternateInput(unittest.TestCase):

    class simplecmd(cmd.Cmd):

        def do_print(self, args):
            print(args, file=self.stdout)

        def do_EOF(self, args):
            return True


    class simplecmd2(simplecmd):

        def do_EOF(self, args):
            print('*** Unknown syntax: EOF', file=self.stdout)
            return True


    def test_file_with_missing_final_nl(self):
        input = io.StringIO("print test\nprint test2")
        output = io.StringIO()
        cmd = self.simplecmd(stdin=input, stdout=output)
        cmd.use_rawinput = False
        cmd.cmdloop()
        self.assertMultiLineEqual(output.getvalue(),
            ("(Cmd) test\n"
             "(Cmd) test2\n"
             "(Cmd) "))


    def test_input_reset_at_EOF(self):
        input = io.StringIO("print test\nprint test2")
        output = io.StringIO()
        cmd = self.simplecmd2(stdin=input, stdout=output)
        cmd.use_rawinput = False
        cmd.cmdloop()
        self.assertMultiLineEqual(output.getvalue(),
            ("(Cmd) test\n"
             "(Cmd) test2\n"
             "(Cmd) *** Unknown syntax: EOF\n"))
        input = io.StringIO("print \n\n")
        output = io.StringIO()
        cmd.stdin = input
        cmd.stdout = output
        cmd.cmdloop()
        self.assertMultiLineEqual(output.getvalue(),
            ("(Cmd) \n"
             "(Cmd) \n"
             "(Cmd) *** Unknown syntax: EOF\n"))


def test_main(verbose=None):
    from test import test_cmd
    support.run_doctest(test_cmd, verbose)
    support.run_unittest(TestAlternateInput)

def test_coverage(coverdir):
    trace = support.import_module('trace')
    tracer=trace.Trace(ignoredirs=[sys.base_prefix, sys.base_exec_prefix,],
                        trace=0, count=1)
    tracer.run('import importlib; importlib.reload(cmd); test_main()')
    r=tracer.results()
    print("Writing coverage results...")
    r.write_results(show_missing=True, summary=True, coverdir=coverdir)

if __name__ == "__main__":
    if "-c" in sys.argv:
        test_coverage('/tmp/cmd.cover')
    elif "-i" in sys.argv:
        samplecmdclass().cmdloop()
    else:
        test_main()

external/cpython_master/Lib/test/test_cmd.py
DragonHPC/dragon
"Test InteractiveConsole and InteractiveInterpreter from code module"
import sys
import unittest
from textwrap import dedent
from contextlib import ExitStack
from unittest import mock
from test import support

code = support.import_module('code')


class TestInteractiveConsole(unittest.TestCase):

    def setUp(self):
        self.console = code.InteractiveConsole()
        self.mock_sys()

    def mock_sys(self):
        "Mock system environment for InteractiveConsole"
        # use exit stack to match patch context managers to addCleanup
        stack = ExitStack()
        self.addCleanup(stack.close)
        self.infunc = stack.enter_context(mock.patch('code.input',
                                          create=True))
        self.stdout = stack.enter_context(mock.patch('code.sys.stdout'))
        self.stderr = stack.enter_context(mock.patch('code.sys.stderr'))
        prepatch = mock.patch('code.sys', wraps=code.sys, spec=code.sys)
        self.sysmod = stack.enter_context(prepatch)
        if sys.excepthook is sys.__excepthook__:
            self.sysmod.excepthook = self.sysmod.__excepthook__
        del self.sysmod.ps1
        del self.sysmod.ps2

    def test_ps1(self):
        self.infunc.side_effect = EOFError('Finished')
        self.console.interact()
        self.assertEqual(self.sysmod.ps1, '>>> ')
        self.sysmod.ps1 = 'custom1> '
        self.console.interact()
        self.assertEqual(self.sysmod.ps1, 'custom1> ')

    def test_ps2(self):
        self.infunc.side_effect = EOFError('Finished')
        self.console.interact()
        self.assertEqual(self.sysmod.ps2, '... ')
        self.sysmod.ps1 = 'custom2> '
        self.console.interact()
        self.assertEqual(self.sysmod.ps1, 'custom2> ')

    def test_console_stderr(self):
        self.infunc.side_effect = ["'antioch'", "", EOFError('Finished')]
        self.console.interact()
        for call in list(self.stdout.method_calls):
            if 'antioch' in ''.join(call[1]):
                break
        else:
            raise AssertionError("no console stdout")

    def test_syntax_error(self):
        self.infunc.side_effect = ["undefined", EOFError('Finished')]
        self.console.interact()
        for call in self.stderr.method_calls:
            if 'NameError' in ''.join(call[1]):
                break
        else:
            raise AssertionError("No syntax error from console")

    def test_sysexcepthook(self):
        self.infunc.side_effect = ["raise ValueError('')",
                                    EOFError('Finished')]
        hook = mock.Mock()
        self.sysmod.excepthook = hook
        self.console.interact()
        self.assertTrue(hook.called)

    def test_banner(self):
        # with banner
        self.infunc.side_effect = EOFError('Finished')
        self.console.interact(banner='Foo')
        self.assertEqual(len(self.stderr.method_calls), 3)
        banner_call = self.stderr.method_calls[0]
        self.assertEqual(banner_call, ['write', ('Foo\n',), {}])

        # no banner
        self.stderr.reset_mock()
        self.infunc.side_effect = EOFError('Finished')
        self.console.interact(banner='')
        self.assertEqual(len(self.stderr.method_calls), 2)

    def test_exit_msg(self):
        # default exit message
        self.infunc.side_effect = EOFError('Finished')
        self.console.interact(banner='')
        self.assertEqual(len(self.stderr.method_calls), 2)
        err_msg = self.stderr.method_calls[1]
        expected = 'now exiting InteractiveConsole...\n'
        self.assertEqual(err_msg, ['write', (expected,), {}])

        # no exit message
        self.stderr.reset_mock()
        self.infunc.side_effect = EOFError('Finished')
        self.console.interact(banner='', exitmsg='')
        self.assertEqual(len(self.stderr.method_calls), 1)

        # custom exit message
        self.stderr.reset_mock()
        message = (
            'bye! \N{GREEK SMALL LETTER ZETA}\N{CYRILLIC SMALL LETTER ZHE}'
            )
        self.infunc.side_effect = EOFError('Finished')
        self.console.interact(banner='', exitmsg=message)
        self.assertEqual(len(self.stderr.method_calls), 2)
        err_msg = self.stderr.method_calls[1]
        expected = message + '\n'
        self.assertEqual(err_msg, ['write', (expected,), {}])


    def test_cause_tb(self):
        self.infunc.side_effect = ["raise ValueError('') from AttributeError",
                                    EOFError('Finished')]
        self.console.interact()
        output = ''.join(''.join(call[1]) for call in self.stderr.method_calls)
        expected = dedent("""
        AttributeError

        The above exception was the direct cause of the following exception:

        Traceback (most recent call last):
          File "<console>", line 1, in <module>
        ValueError
        """)
        self.assertIn(expected, output)

    def test_context_tb(self):
        self.infunc.side_effect = ["try: ham\nexcept: eggs\n",
                                    EOFError('Finished')]
        self.console.interact()
        output = ''.join(''.join(call[1]) for call in self.stderr.method_calls)
        expected = dedent("""
        Traceback (most recent call last):
          File "<console>", line 1, in <module>
        NameError: name 'ham' is not defined

        During handling of the above exception, another exception occurred:

        Traceback (most recent call last):
          File "<console>", line 2, in <module>
        NameError: name 'eggs' is not defined
        """)
        self.assertIn(expected, output)


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_code_module.py
DragonHPC/dragon
#
# test_codecencodings_kr.py
#   Codec encoding tests for ROK encodings.
#

from test import multibytecodec_support
import unittest

class Test_CP949(multibytecodec_support.TestBase, unittest.TestCase):
    encoding = 'cp949'
    tstring = multibytecodec_support.load_teststring('cp949')
    codectests = (
        # invalid bytes
        (b"abc\x80\x80\xc1\xc4", "strict",  None),
        (b"abc\xc8", "strict",  None),
        (b"abc\x80\x80\xc1\xc4", "replace", "abc\ufffd\ufffd\uc894"),
        (b"abc\x80\x80\xc1\xc4\xc8", "replace", "abc\ufffd\ufffd\uc894\ufffd"),
        (b"abc\x80\x80\xc1\xc4", "ignore",  "abc\uc894"),
    )

class Test_EUCKR(multibytecodec_support.TestBase, unittest.TestCase):
    encoding = 'euc_kr'
    tstring = multibytecodec_support.load_teststring('euc_kr')
    codectests = (
        # invalid bytes
        (b"abc\x80\x80\xc1\xc4", "strict",  None),
        (b"abc\xc8", "strict",  None),
        (b"abc\x80\x80\xc1\xc4", "replace", 'abc\ufffd\ufffd\uc894'),
        (b"abc\x80\x80\xc1\xc4\xc8", "replace", "abc\ufffd\ufffd\uc894\ufffd"),
        (b"abc\x80\x80\xc1\xc4", "ignore",  "abc\uc894"),

        # composed make-up sequence errors
        (b"\xa4\xd4", "strict", None),
        (b"\xa4\xd4\xa4", "strict", None),
        (b"\xa4\xd4\xa4\xb6", "strict", None),
        (b"\xa4\xd4\xa4\xb6\xa4", "strict", None),
        (b"\xa4\xd4\xa4\xb6\xa4\xd0", "strict", None),
        (b"\xa4\xd4\xa4\xb6\xa4\xd0\xa4", "strict", None),
        (b"\xa4\xd4\xa4\xb6\xa4\xd0\xa4\xd4", "strict", "\uc4d4"),
        (b"\xa4\xd4\xa4\xb6\xa4\xd0\xa4\xd4x", "strict", "\uc4d4x"),
        (b"a\xa4\xd4\xa4\xb6\xa4", "replace", 'a\ufffd'),
        (b"\xa4\xd4\xa3\xb6\xa4\xd0\xa4\xd4", "strict", None),
        (b"\xa4\xd4\xa4\xb6\xa3\xd0\xa4\xd4", "strict", None),
        (b"\xa4\xd4\xa4\xb6\xa4\xd0\xa3\xd4", "strict", None),
        (b"\xa4\xd4\xa4\xff\xa4\xd0\xa4\xd4", "replace", '\ufffd\u6e21\ufffd\u3160\ufffd'),
        (b"\xa4\xd4\xa4\xb6\xa4\xff\xa4\xd4", "replace", '\ufffd\u6e21\ub544\ufffd\ufffd'),
        (b"\xa4\xd4\xa4\xb6\xa4\xd0\xa4\xff", "replace", '\ufffd\u6e21\ub544\u572d\ufffd'),
        (b"\xa4\xd4\xff\xa4\xd4\xa4\xb6\xa4\xd0\xa4\xd4", "replace", '\ufffd\ufffd\ufffd\uc4d4'),
        (b"\xc1\xc4", "strict", "\uc894"),
    )

class Test_JOHAB(multibytecodec_support.TestBase, unittest.TestCase):
    encoding = 'johab'
    tstring = multibytecodec_support.load_teststring('johab')
    codectests = (
        # invalid bytes
        (b"abc\x80\x80\xc1\xc4", "strict",  None),
        (b"abc\xc8", "strict",  None),
        (b"abc\x80\x80\xc1\xc4", "replace", "abc\ufffd\ufffd\ucd27"),
        (b"abc\x80\x80\xc1\xc4\xc8", "replace", "abc\ufffd\ufffd\ucd27\ufffd"),
        (b"abc\x80\x80\xc1\xc4", "ignore",  "abc\ucd27"),
        (b"\xD8abc", "replace",  "\uFFFDabc"),
        (b"\xD8\xFFabc", "replace",  "\uFFFD\uFFFDabc"),
        (b"\x84bxy", "replace",  "\uFFFDbxy"),
        (b"\x8CBxy", "replace",  "\uFFFDBxy"),
    )

if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_codecencodings_kr.py
DragonHPC/dragon
"""Test the errno module
   Roger E. Masse
"""

import errno
import unittest

std_c_errors = frozenset(['EDOM', 'ERANGE'])

class ErrnoAttributeTests(unittest.TestCase):

    def test_for_improper_attributes(self):
        # No unexpected attributes should be on the module.
        for error_code in std_c_errors:
            self.assertTrue(hasattr(errno, error_code),
                            "errno is missing %s" % error_code)

    def test_using_errorcode(self):
        # Every key value in errno.errorcode should be on the module.
        for value in errno.errorcode.values():
            self.assertTrue(hasattr(errno, value),
                            'no %s attr in errno' % value)


class ErrorcodeTests(unittest.TestCase):

    def test_attributes_in_errorcode(self):
        for attribute in errno.__dict__.keys():
            if attribute.isupper():
                self.assertIn(getattr(errno, attribute), errno.errorcode,
                              'no %s attr in errno.errorcode' % attribute)


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_errno.py
DragonHPC/dragon
# -*- coding: utf-8 -*-
# There are tests here with unicode string literals and
# identifiers. There's a code in ast.c that was added because of a
# failure with a non-ascii-only expression.  So, I have tests for
# that.  There are workarounds that would let me run tests for that
# code without unicode identifiers and strings, but just using them
# directly seems like the easiest and therefore safest thing to do.
# Unicode identifiers in tests is allowed by PEP 3131.

import ast
import os
import re
import types
import decimal
import unittest
from test.support import temp_cwd, use_old_parser
from test.support.script_helper import assert_python_failure

a_global = 'global variable'

# You could argue that I'm too strict in looking for specific error
#  values with assertRaisesRegex, but without it it's way too easy to
#  make a syntax error in the test strings. Especially with all of the
#  triple quotes, raw strings, backslashes, etc. I think it's a
#  worthwhile tradeoff. When I switched to this method, I found many
#  examples where I wasn't testing what I thought I was.

class TestCase(unittest.TestCase):
    def assertAllRaise(self, exception_type, regex, error_strings):
        for str in error_strings:
            with self.subTest(str=str):
                with self.assertRaisesRegex(exception_type, regex):
                    eval(str)

    def test__format__lookup(self):
        # Make sure __format__ is looked up on the type, not the instance.
        class X:
            def __format__(self, spec):
                return 'class'

        x = X()

        # Add a bound __format__ method to the 'y' instance, but not
        #  the 'x' instance.
        y = X()
        y.__format__ = types.MethodType(lambda self, spec: 'instance', y)

        self.assertEqual(f'{y}', format(y))
        self.assertEqual(f'{y}', 'class')
        self.assertEqual(format(x), format(y))

        # __format__ is not called this way, but still make sure it
        #  returns what we expect (so we can make sure we're bypassing
        #  it).
        self.assertEqual(x.__format__(''), 'class')
        self.assertEqual(y.__format__(''), 'instance')

        # This is how __format__ is actually called.
        self.assertEqual(type(x).__format__(x, ''), 'class')
        self.assertEqual(type(y).__format__(y, ''), 'class')

    def test_ast(self):
        # Inspired by http://bugs.python.org/issue24975
        class X:
            def __init__(self):
                self.called = False
            def __call__(self):
                self.called = True
                return 4
        x = X()
        expr = """
a = 10
f'{a * x()}'"""
        t = ast.parse(expr)
        c = compile(t, '', 'exec')

        # Make sure x was not called.
        self.assertFalse(x.called)

        # Actually run the code.
        exec(c)

        # Make sure x was called.
        self.assertTrue(x.called)

    def test_ast_line_numbers(self):
        expr = """
a = 10
f'{a * x()}'"""
        t = ast.parse(expr)
        self.assertEqual(type(t), ast.Module)
        self.assertEqual(len(t.body), 2)
        # check `a = 10`
        self.assertEqual(type(t.body[0]), ast.Assign)
        self.assertEqual(t.body[0].lineno, 2)
        # check `f'...'`
        self.assertEqual(type(t.body[1]), ast.Expr)
        self.assertEqual(type(t.body[1].value), ast.JoinedStr)
        self.assertEqual(len(t.body[1].value.values), 1)
        self.assertEqual(type(t.body[1].value.values[0]), ast.FormattedValue)
        self.assertEqual(t.body[1].lineno, 3)
        self.assertEqual(t.body[1].value.lineno, 3)
        self.assertEqual(t.body[1].value.values[0].lineno, 3)
        # check the binop location
        binop = t.body[1].value.values[0].value
        self.assertEqual(type(binop), ast.BinOp)
        self.assertEqual(type(binop.left), ast.Name)
        self.assertEqual(type(binop.op), ast.Mult)
        self.assertEqual(type(binop.right), ast.Call)
        self.assertEqual(binop.lineno, 3)
        self.assertEqual(binop.left.lineno, 3)
        self.assertEqual(binop.right.lineno, 3)
        self.assertEqual(binop.col_offset, 3)
        self.assertEqual(binop.left.col_offset, 3)
        self.assertEqual(binop.right.col_offset, 7)

    def test_ast_line_numbers_multiple_formattedvalues(self):
        expr = """
f'no formatted values'
f'eggs {a * x()} spam {b + y()}'"""
        t = ast.parse(expr)
        self.assertEqual(type(t), ast.Module)
        self.assertEqual(len(t.body), 2)
        # check `f'no formatted value'`
        self.assertEqual(type(t.body[0]), ast.Expr)
        self.assertEqual(type(t.body[0].value), ast.JoinedStr)
        self.assertEqual(t.body[0].lineno, 2)
        # check `f'...'`
        self.assertEqual(type(t.body[1]), ast.Expr)
        self.assertEqual(type(t.body[1].value), ast.JoinedStr)
        self.assertEqual(len(t.body[1].value.values), 4)
        self.assertEqual(type(t.body[1].value.values[0]), ast.Constant)
        self.assertEqual(type(t.body[1].value.values[0].value), str)
        self.assertEqual(type(t.body[1].value.values[1]), ast.FormattedValue)
        self.assertEqual(type(t.body[1].value.values[2]), ast.Constant)
        self.assertEqual(type(t.body[1].value.values[2].value), str)
        self.assertEqual(type(t.body[1].value.values[3]), ast.FormattedValue)
        self.assertEqual(t.body[1].lineno, 3)
        self.assertEqual(t.body[1].value.lineno, 3)
        self.assertEqual(t.body[1].value.values[0].lineno, 3)
        self.assertEqual(t.body[1].value.values[1].lineno, 3)
        self.assertEqual(t.body[1].value.values[2].lineno, 3)
        self.assertEqual(t.body[1].value.values[3].lineno, 3)
        # check the first binop location
        binop1 = t.body[1].value.values[1].value
        self.assertEqual(type(binop1), ast.BinOp)
        self.assertEqual(type(binop1.left), ast.Name)
        self.assertEqual(type(binop1.op), ast.Mult)
        self.assertEqual(type(binop1.right), ast.Call)
        self.assertEqual(binop1.lineno, 3)
        self.assertEqual(binop1.left.lineno, 3)
        self.assertEqual(binop1.right.lineno, 3)
        self.assertEqual(binop1.col_offset, 8)
        self.assertEqual(binop1.left.col_offset, 8)
        self.assertEqual(binop1.right.col_offset, 12)
        # check the second binop location
        binop2 = t.body[1].value.values[3].value
        self.assertEqual(type(binop2), ast.BinOp)
        self.assertEqual(type(binop2.left), ast.Name)
        self.assertEqual(type(binop2.op), ast.Add)
        self.assertEqual(type(binop2.right), ast.Call)
        self.assertEqual(binop2.lineno, 3)
        self.assertEqual(binop2.left.lineno, 3)
        self.assertEqual(binop2.right.lineno, 3)
        self.assertEqual(binop2.col_offset, 23)
        self.assertEqual(binop2.left.col_offset, 23)
        self.assertEqual(binop2.right.col_offset, 27)

    def test_ast_line_numbers_nested(self):
        expr = """
a = 10
f'{a * f"-{x()}-"}'"""
        t = ast.parse(expr)
        self.assertEqual(type(t), ast.Module)
        self.assertEqual(len(t.body), 2)
        # check `a = 10`
        self.assertEqual(type(t.body[0]), ast.Assign)
        self.assertEqual(t.body[0].lineno, 2)
        # check `f'...'`
        self.assertEqual(type(t.body[1]), ast.Expr)
        self.assertEqual(type(t.body[1].value), ast.JoinedStr)
        self.assertEqual(len(t.body[1].value.values), 1)
        self.assertEqual(type(t.body[1].value.values[0]), ast.FormattedValue)
        self.assertEqual(t.body[1].lineno, 3)
        self.assertEqual(t.body[1].value.lineno, 3)
        self.assertEqual(t.body[1].value.values[0].lineno, 3)
        # check the binop location
        binop = t.body[1].value.values[0].value
        self.assertEqual(type(binop), ast.BinOp)
        self.assertEqual(type(binop.left), ast.Name)
        self.assertEqual(type(binop.op), ast.Mult)
        self.assertEqual(type(binop.right), ast.JoinedStr)
        self.assertEqual(binop.lineno, 3)
        self.assertEqual(binop.left.lineno, 3)
        self.assertEqual(binop.right.lineno, 3)
        self.assertEqual(binop.col_offset, 3)
        self.assertEqual(binop.left.col_offset, 3)
        self.assertEqual(binop.right.col_offset, 7)
        # check the nested call location
        self.assertEqual(len(binop.right.values), 3)
        self.assertEqual(type(binop.right.values[0]), ast.Constant)
        self.assertEqual(type(binop.right.values[0].value), str)
        self.assertEqual(type(binop.right.values[1]), ast.FormattedValue)
        self.assertEqual(type(binop.right.values[2]), ast.Constant)
        self.assertEqual(type(binop.right.values[2].value), str)
        self.assertEqual(binop.right.values[0].lineno, 3)
        self.assertEqual(binop.right.values[1].lineno, 3)
        self.assertEqual(binop.right.values[2].lineno, 3)
        call = binop.right.values[1].value
        self.assertEqual(type(call), ast.Call)
        self.assertEqual(call.lineno, 3)
        self.assertEqual(call.col_offset, 11)

    def test_ast_line_numbers_duplicate_expression(self):
        """Duplicate expression

        NOTE: this is currently broken, always sets location of the first
        expression.
        """
        expr = """
a = 10
f'{a * x()} {a * x()} {a * x()}'
"""
        t = ast.parse(expr)
        self.assertEqual(type(t), ast.Module)
        self.assertEqual(len(t.body), 2)
        # check `a = 10`
        self.assertEqual(type(t.body[0]), ast.Assign)
        self.assertEqual(t.body[0].lineno, 2)
        # check `f'...'`
        self.assertEqual(type(t.body[1]), ast.Expr)
        self.assertEqual(type(t.body[1].value), ast.JoinedStr)
        self.assertEqual(len(t.body[1].value.values), 5)
        self.assertEqual(type(t.body[1].value.values[0]), ast.FormattedValue)
        self.assertEqual(type(t.body[1].value.values[1]), ast.Constant)
        self.assertEqual(type(t.body[1].value.values[1].value), str)
        self.assertEqual(type(t.body[1].value.values[2]), ast.FormattedValue)
        self.assertEqual(type(t.body[1].value.values[3]), ast.Constant)
        self.assertEqual(type(t.body[1].value.values[3].value), str)
        self.assertEqual(type(t.body[1].value.values[4]), ast.FormattedValue)
        self.assertEqual(t.body[1].lineno, 3)
        self.assertEqual(t.body[1].value.lineno, 3)
        self.assertEqual(t.body[1].value.values[0].lineno, 3)
        self.assertEqual(t.body[1].value.values[1].lineno, 3)
        self.assertEqual(t.body[1].value.values[2].lineno, 3)
        self.assertEqual(t.body[1].value.values[3].lineno, 3)
        self.assertEqual(t.body[1].value.values[4].lineno, 3)
        # check the first binop location
        binop = t.body[1].value.values[0].value
        self.assertEqual(type(binop), ast.BinOp)
        self.assertEqual(type(binop.left), ast.Name)
        self.assertEqual(type(binop.op), ast.Mult)
        self.assertEqual(type(binop.right), ast.Call)
        self.assertEqual(binop.lineno, 3)
        self.assertEqual(binop.left.lineno, 3)
        self.assertEqual(binop.right.lineno, 3)
        self.assertEqual(binop.col_offset, 3)
        self.assertEqual(binop.left.col_offset, 3)
        self.assertEqual(binop.right.col_offset, 7)
        # check the second binop location
        binop = t.body[1].value.values[2].value
        self.assertEqual(type(binop), ast.BinOp)
        self.assertEqual(type(binop.left), ast.Name)
        self.assertEqual(type(binop.op), ast.Mult)
        self.assertEqual(type(binop.right), ast.Call)
        self.assertEqual(binop.lineno, 3)
        self.assertEqual(binop.left.lineno, 3)
        self.assertEqual(binop.right.lineno, 3)
        self.assertEqual(binop.col_offset, 3)  # FIXME: this is wrong
        self.assertEqual(binop.left.col_offset, 3)  # FIXME: this is wrong
        self.assertEqual(binop.right.col_offset, 7)  # FIXME: this is wrong
        # check the third binop location
        binop = t.body[1].value.values[4].value
        self.assertEqual(type(binop), ast.BinOp)
        self.assertEqual(type(binop.left), ast.Name)
        self.assertEqual(type(binop.op), ast.Mult)
        self.assertEqual(type(binop.right), ast.Call)
        self.assertEqual(binop.lineno, 3)
        self.assertEqual(binop.left.lineno, 3)
        self.assertEqual(binop.right.lineno, 3)
        self.assertEqual(binop.col_offset, 3)  # FIXME: this is wrong
        self.assertEqual(binop.left.col_offset, 3)  # FIXME: this is wrong
        self.assertEqual(binop.right.col_offset, 7)  # FIXME: this is wrong

    def test_ast_line_numbers_multiline_fstring(self):
        # See bpo-30465 for details.
        expr = """
a = 10
f'''
  {a
     *
       x()}
non-important content
'''
"""
        t = ast.parse(expr)
        self.assertEqual(type(t), ast.Module)
        self.assertEqual(len(t.body), 2)
        # check `a = 10`
        self.assertEqual(type(t.body[0]), ast.Assign)
        self.assertEqual(t.body[0].lineno, 2)
        # check `f'...'`
        self.assertEqual(type(t.body[1]), ast.Expr)
        self.assertEqual(type(t.body[1].value), ast.JoinedStr)
        self.assertEqual(len(t.body[1].value.values), 3)
        self.assertEqual(type(t.body[1].value.values[0]), ast.Constant)
        self.assertEqual(type(t.body[1].value.values[0].value), str)
        self.assertEqual(type(t.body[1].value.values[1]), ast.FormattedValue)
        self.assertEqual(type(t.body[1].value.values[2]), ast.Constant)
        self.assertEqual(type(t.body[1].value.values[2].value), str)
        self.assertEqual(t.body[1].lineno, 3)
        self.assertEqual(t.body[1].value.lineno, 3)
        self.assertEqual(t.body[1].value.values[0].lineno, 3)
        self.assertEqual(t.body[1].value.values[1].lineno, 3)
        self.assertEqual(t.body[1].value.values[2].lineno, 3)
        self.assertEqual(t.body[1].col_offset, 0)
        self.assertEqual(t.body[1].value.col_offset, 0)
        self.assertEqual(t.body[1].value.values[0].col_offset, 0)
        self.assertEqual(t.body[1].value.values[1].col_offset, 0)
        self.assertEqual(t.body[1].value.values[2].col_offset, 0)
        # NOTE: the following lineno information and col_offset is correct for
        # expressions within FormattedValues.
        binop = t.body[1].value.values[1].value
        self.assertEqual(type(binop), ast.BinOp)
        self.assertEqual(type(binop.left), ast.Name)
        self.assertEqual(type(binop.op), ast.Mult)
        self.assertEqual(type(binop.right), ast.Call)
        self.assertEqual(binop.lineno, 4)
        self.assertEqual(binop.left.lineno, 4)
        self.assertEqual(binop.right.lineno, 6)
        self.assertEqual(binop.col_offset, 4)
        self.assertEqual(binop.left.col_offset, 4)
        self.assertEqual(binop.right.col_offset, 7)

    def test_ast_line_numbers_with_parentheses(self):
        expr = """
x = (
    f" {test(t)}"
)"""
        t = ast.parse(expr)
        self.assertEqual(type(t), ast.Module)
        self.assertEqual(len(t.body), 1)
        # check the test(t) location
        call = t.body[0].value.values[1].value
        self.assertEqual(type(call), ast.Call)
        self.assertEqual(call.lineno, 3)
        self.assertEqual(call.end_lineno, 3)
        self.assertEqual(call.col_offset, 8)
        self.assertEqual(call.end_col_offset, 15)

        expr = """
x = (
        'PERL_MM_OPT', (
            f'wat'
            f'some_string={f(x)} '
            f'wat'
        ),
)
"""
        t = ast.parse(expr)
        self.assertEqual(type(t), ast.Module)
        self.assertEqual(len(t.body), 1)
        # check the fstring
        fstring = t.body[0].value.elts[1]
        self.assertEqual(type(fstring), ast.JoinedStr)
        self.assertEqual(len(fstring.values), 3)
        wat1, middle, wat2 = fstring.values
        # check the first wat
        self.assertEqual(type(wat1), ast.Constant)
        self.assertEqual(wat1.lineno, 4)
        self.assertEqual(wat1.end_lineno, 6)
        self.assertEqual(wat1.col_offset, 12)
        self.assertEqual(wat1.end_col_offset, 18)
        # check the call
        call = middle.value
        self.assertEqual(type(call), ast.Call)
        self.assertEqual(call.lineno, 5)
        self.assertEqual(call.end_lineno, 5)
        self.assertEqual(call.col_offset, 27)
        self.assertEqual(call.end_col_offset, 31)
        # check the second wat
        self.assertEqual(type(wat2), ast.Constant)
        self.assertEqual(wat2.lineno, 4)
        self.assertEqual(wat2.end_lineno, 6)
        self.assertEqual(wat2.col_offset, 12)
        self.assertEqual(wat2.end_col_offset, 18)

    def test_docstring(self):
        def f():
            f'''Not a docstring'''
        self.assertIsNone(f.__doc__)
        def g():
            '''Not a docstring''' \
            f''
        self.assertIsNone(g.__doc__)

    def test_literal_eval(self):
        with self.assertRaisesRegex(ValueError, 'malformed node or string'):
            ast.literal_eval("f'x'")

    def test_ast_compile_time_concat(self):
        x = ['']

        expr = """x[0] = 'foo' f'{3}'"""
        t = ast.parse(expr)
        c = compile(t, '', 'exec')
        exec(c)
        self.assertEqual(x[0], 'foo3')

    def test_compile_time_concat_errors(self):
        self.assertAllRaise(SyntaxError,
                            'cannot mix bytes and nonbytes literals',
                            [r"""f'' b''""",
                             r"""b'' f''""",
                             ])

    def test_literal(self):
        self.assertEqual(f'', '')
        self.assertEqual(f'a', 'a')
        self.assertEqual(f' ', ' ')

    def test_unterminated_string(self):
        self.assertAllRaise(SyntaxError, 'f-string: unterminated string',
                            [r"""f'{"x'""",
                             r"""f'{"x}'""",
                             r"""f'{("x'""",
                             r"""f'{("x}'""",
                             ])

    def test_mismatched_parens(self):
        self.assertAllRaise(SyntaxError, r"f-string: closing parenthesis '\}' "
                            r"does not match opening parenthesis '\('",
                            ["f'{((}'",
                             ])
        self.assertAllRaise(SyntaxError, r"f-string: closing parenthesis '\)' "
                            r"does not match opening parenthesis '\['",
                            ["f'{a[4)}'",
                            ])
        self.assertAllRaise(SyntaxError, r"f-string: closing parenthesis '\]' "
                            r"does not match opening parenthesis '\('",
                            ["f'{a(4]}'",
                            ])
        self.assertAllRaise(SyntaxError, r"f-string: closing parenthesis '\}' "
                            r"does not match opening parenthesis '\['",
                            ["f'{a[4}'",
                            ])
        self.assertAllRaise(SyntaxError, r"f-string: closing parenthesis '\}' "
                            r"does not match opening parenthesis '\('",
                            ["f'{a(4}'",
                            ])
        self.assertRaises(SyntaxError, eval, "f'{" + "("*500 + "}'")

    def test_double_braces(self):
        self.assertEqual(f'{{', '{')
        self.assertEqual(f'a{{', 'a{')
        self.assertEqual(f'{{b', '{b')
        self.assertEqual(f'a{{b', 'a{b')
        self.assertEqual(f'}}', '}')
        self.assertEqual(f'a}}', 'a}')
        self.assertEqual(f'}}b', '}b')
        self.assertEqual(f'a}}b', 'a}b')
        self.assertEqual(f'{{}}', '{}')
        self.assertEqual(f'a{{}}', 'a{}')
        self.assertEqual(f'{{b}}', '{b}')
        self.assertEqual(f'{{}}c', '{}c')
        self.assertEqual(f'a{{b}}', 'a{b}')
        self.assertEqual(f'a{{}}c', 'a{}c')
        self.assertEqual(f'{{b}}c', '{b}c')
        self.assertEqual(f'a{{b}}c', 'a{b}c')

        self.assertEqual(f'{{{10}', '{10')
        self.assertEqual(f'}}{10}', '}10')
        self.assertEqual(f'}}{{{10}', '}{10')
        self.assertEqual(f'}}a{{{10}', '}a{10')

        self.assertEqual(f'{10}{{', '10{')
        self.assertEqual(f'{10}}}', '10}')
        self.assertEqual(f'{10}}}{{', '10}{')
        self.assertEqual(f'{10}}}a{{' '}', '10}a{}')

        # Inside of strings, don't interpret doubled brackets.
        self.assertEqual(f'{"{{}}"}', '{{}}')

        self.assertAllRaise(TypeError, 'unhashable type',
                            ["f'{ {{}} }'", # dict in a set
                             ])

    def test_compile_time_concat(self):
        x = 'def'
        self.assertEqual('abc' f'## {x}ghi', 'abc## defghi')
        self.assertEqual('abc' f'{x}' 'ghi', 'abcdefghi')
        self.assertEqual('abc' f'{x}' 'gh' f'i{x:4}', 'abcdefghidef ')
        self.assertEqual('{x}' f'{x}', '{x}def')
        self.assertEqual('{x' f'{x}', '{xdef')
        self.assertEqual('{x}' f'{x}', '{x}def')
        self.assertEqual('{{x}}' f'{x}', '{{x}}def')
        self.assertEqual('{{x' f'{x}', '{{xdef')
        self.assertEqual('x}}' f'{x}', 'x}}def')
        self.assertEqual(f'{x}' 'x}}', 'defx}}')
        self.assertEqual(f'{x}' '', 'def')
        self.assertEqual('' f'{x}' '', 'def')
        self.assertEqual('' f'{x}', 'def')
        self.assertEqual(f'{x}' '2', 'def2')
        self.assertEqual('1' f'{x}' '2', '1def2')
        self.assertEqual('1' f'{x}', '1def')
        self.assertEqual(f'{x}' f'-{x}', 'def-def')
        self.assertEqual('' f'', '')
        self.assertEqual('' f'' '', '')
        self.assertEqual('' f'' '' f'', '')
        self.assertEqual(f'', '')
        self.assertEqual(f'' '', '')
        self.assertEqual(f'' '' f'', '')
        self.assertEqual(f'' '' f'' '', '')

        self.assertAllRaise(SyntaxError, "f-string: expecting '}'",
                            ["f'{3' f'}'",  # can't concat to get a valid f-string
                             ])

    def test_comments(self):
        # These aren't comments, since they're in strings.
        d = {'#': 'hash'}
        self.assertEqual(f'{"#"}', '#')
        self.assertEqual(f'{d["#"]}', 'hash')

        self.assertAllRaise(SyntaxError, "f-string expression part cannot include '#'",
                            ["f'{1#}'",   # error because the expression becomes "(1#)"
                             "f'{3(#)}'",
                             "f'{#}'",
                             ])
        self.assertAllRaise(SyntaxError, r"f-string: unmatched '\)'",
                            ["f'{)#}'",   # When wrapped in parens, this becomes
                                          #  '()#)'.  Make sure that doesn't compile.
                             ])

    def test_many_expressions(self):
        # Create a string with many expressions in it. Note that
        #  because we have a space in here as a literal, we're actually
        #  going to use twice as many ast nodes: one for each literal
        #  plus one for each expression.
        def build_fstr(n, extra=''):
            return "f'" + ('{x} ' * n) + extra + "'"

        x = 'X'
        width = 1

        # Test around 256.
        for i in range(250, 260):
            self.assertEqual(eval(build_fstr(i)), (x+' ')*i)

        # Test concatenating 2 largs fstrings.
        self.assertEqual(eval(build_fstr(255)*256), (x+' ')*(255*256))

        s = build_fstr(253, '{x:{width}} ')
        self.assertEqual(eval(s), (x+' ')*254)

        # Test lots of expressions and constants, concatenated.
        s = "f'{1}' 'x' 'y'" * 1024
        self.assertEqual(eval(s), '1xy' * 1024)

    def test_format_specifier_expressions(self):
        width = 10
        precision = 4
        value = decimal.Decimal('12.34567')
        self.assertEqual(f'result: {value:{width}.{precision}}', 'result:      12.35')
        self.assertEqual(f'result: {value:{width!r}.{precision}}', 'result:      12.35')
        self.assertEqual(f'result: {value:{width:0}.{precision:1}}', 'result:      12.35')
        self.assertEqual(f'result: {value:{1}{0:0}.{precision:1}}', 'result:      12.35')
        self.assertEqual(f'result: {value:{ 1}{ 0:0}.{ precision:1}}', 'result:      12.35')
        self.assertEqual(f'{10:#{1}0x}', '       0xa')
        self.assertEqual(f'{10:{"#"}1{0}{"x"}}', '       0xa')
        self.assertEqual(f'{-10:-{"#"}1{0}x}', '      -0xa')
        self.assertEqual(f'{-10:{"-"}#{1}0{"x"}}', '      -0xa')
        self.assertEqual(f'{10:#{3 != {4:5} and width}x}', '       0xa')

        self.assertAllRaise(SyntaxError, "f-string: expecting '}'",
                            ["""f'{"s"!r{":10"}}'""",

                             # This looks like a nested format spec.
                             ])

        err_msg = "invalid syntax" if use_old_parser() else "f-string: invalid syntax"
        self.assertAllRaise(SyntaxError, err_msg,
                            [# Invalid syntax inside a nested spec.
                             "f'{4:{/5}}'",
                             ])

        self.assertAllRaise(SyntaxError, "f-string: expressions nested too deeply",
                            [# Can't nest format specifiers.
                             "f'result: {value:{width:{0}}.{precision:1}}'",
                             ])

        self.assertAllRaise(SyntaxError, 'f-string: invalid conversion character',
                            [# No expansion inside conversion or for
                             #  the : or ! itself.
                             """f'{"s"!{"r"}}'""",
                             ])

    def test_side_effect_order(self):
        class X:
            def __init__(self):
                self.i = 0
            def __format__(self, spec):
                self.i += 1
                return str(self.i)

        x = X()
        self.assertEqual(f'{x} {x}', '1 2')

    def test_missing_expression(self):
        self.assertAllRaise(SyntaxError, 'f-string: empty expression not allowed',
                            ["f'{}'",
                             "f'{ }'"
                             "f' {} '",
                             "f'{!r}'",
                             "f'{ !r}'",
                             "f'{10:{ }}'",
                             "f' { } '",

                             # The Python parser ignores also the following
                             # whitespace characters in additional to a space.
                             "f'''{\t\f\r\n}'''",

                             # Catch the empty expression before the
                             #  invalid conversion.
                             "f'{!x}'",
                             "f'{ !xr}'",
                             "f'{!x:}'",
                             "f'{!x:a}'",
                             "f'{ !xr:}'",
                             "f'{ !xr:a}'",

                             "f'{!}'",
                             "f'{:}'",

                             # We find the empty expression before the
                             #  missing closing brace.
                             "f'{!'",
                             "f'{!s:'",
                             "f'{:'",
                             "f'{:x'",
                             ])

        # Different error message is raised for other whitespace characters.
        self.assertAllRaise(SyntaxError, r"invalid non-printable character U\+00A0",
                            ["f'''{\xa0}'''",
                             "\xa0",
                             ])

    def test_parens_in_expressions(self):
        self.assertEqual(f'{3,}', '(3,)')

        # Add these because when an expression is evaluated, parens
        #  are added around it. But we shouldn't go from an invalid
        #  expression to a valid one. The added parens are just
        #  supposed to allow whitespace (including newlines).
        err_msg = "invalid syntax" if use_old_parser() else "f-string: invalid syntax"
        self.assertAllRaise(SyntaxError, err_msg,
                            ["f'{,}'",
                             "f'{,}'",  # this is (,), which is an error
                             ])

        self.assertAllRaise(SyntaxError, r"f-string: unmatched '\)'",
                            ["f'{3)+(4}'",
                             ])

        self.assertAllRaise(SyntaxError, 'EOL while scanning string literal',
                            ["f'{\n}'",
                             ])

    def test_backslashes_in_string_part(self):
        self.assertEqual(f'\t', '\t')
        self.assertEqual(r'\t', '\\t')
        self.assertEqual(rf'\t', '\\t')
        self.assertEqual(f'{2}\t', '2\t')
        self.assertEqual(f'{2}\t{3}', '2\t3')
        self.assertEqual(f'\t{3}', '\t3')

        self.assertEqual(f'\u0394', '\u0394')
        self.assertEqual(r'\u0394', '\\u0394')
        self.assertEqual(rf'\u0394', '\\u0394')
        self.assertEqual(f'{2}\u0394', '2\u0394')
        self.assertEqual(f'{2}\u0394{3}', '2\u03943')
        self.assertEqual(f'\u0394{3}', '\u03943')

        self.assertEqual(f'\U00000394', '\u0394')
        self.assertEqual(r'\U00000394', '\\U00000394')
        self.assertEqual(rf'\U00000394', '\\U00000394')
        self.assertEqual(f'{2}\U00000394', '2\u0394')
        self.assertEqual(f'{2}\U00000394{3}', '2\u03943')
        self.assertEqual(f'\U00000394{3}', '\u03943')

        self.assertEqual(f'\N{GREEK CAPITAL LETTER DELTA}', '\u0394')
        self.assertEqual(f'{2}\N{GREEK CAPITAL LETTER DELTA}', '2\u0394')
        self.assertEqual(f'{2}\N{GREEK CAPITAL LETTER DELTA}{3}', '2\u03943')
        self.assertEqual(f'\N{GREEK CAPITAL LETTER DELTA}{3}', '\u03943')
        self.assertEqual(f'2\N{GREEK CAPITAL LETTER DELTA}', '2\u0394')
        self.assertEqual(f'2\N{GREEK CAPITAL LETTER DELTA}3', '2\u03943')
        self.assertEqual(f'\N{GREEK CAPITAL LETTER DELTA}3', '\u03943')

        self.assertEqual(f'\x20', ' ')
        self.assertEqual(r'\x20', '\\x20')
        self.assertEqual(rf'\x20', '\\x20')
        self.assertEqual(f'{2}\x20', '2 ')
        self.assertEqual(f'{2}\x20{3}', '2 3')
        self.assertEqual(f'\x20{3}', ' 3')

        self.assertEqual(f'2\x20', '2 ')
        self.assertEqual(f'2\x203', '2 3')
        self.assertEqual(f'\x203', ' 3')

        with self.assertWarns(DeprecationWarning):  # invalid escape sequence
            value = eval(r"f'\{6*7}'")
        self.assertEqual(value, '\\42')
        self.assertEqual(f'\\{6*7}', '\\42')
        self.assertEqual(fr'\{6*7}', '\\42')

        AMPERSAND = 'spam'
        # Get the right unicode character (&), or pick up local variable
        # depending on the number of backslashes.
        self.assertEqual(f'\N{AMPERSAND}', '&')
        self.assertEqual(f'\\N{AMPERSAND}', '\\Nspam')
        self.assertEqual(fr'\N{AMPERSAND}', '\\Nspam')
        self.assertEqual(f'\\\N{AMPERSAND}', '\\&')

    def test_misformed_unicode_character_name(self):
        # These test are needed because unicode names are parsed
        # differently inside f-strings.
        self.assertAllRaise(SyntaxError, r"\(unicode error\) 'unicodeescape' codec can't decode bytes in position .*: malformed \\N character escape",
                            [r"f'\N'",
                             r"f'\N{'",
                             r"f'\N{GREEK CAPITAL LETTER DELTA'",

                             # Here are the non-f-string versions,
                             #  which should give the same errors.
                             r"'\N'",
                             r"'\N{'",
                             r"'\N{GREEK CAPITAL LETTER DELTA'",
                             ])

    def test_no_backslashes_in_expression_part(self):
        self.assertAllRaise(SyntaxError, 'f-string expression part cannot include a backslash',
                            [r"f'{\'a\'}'",
                             r"f'{\t3}'",
                             r"f'{\}'",
                             r"rf'{\'a\'}'",
                             r"rf'{\t3}'",
                             r"rf'{\}'",
                             r"""rf'{"\N{LEFT CURLY BRACKET}"}'""",
                             r"f'{\n}'",
                             ])

    def test_no_escapes_for_braces(self):
        """
        Only literal curly braces begin an expression.
        """
        # \x7b is '{'.
        self.assertEqual(f'\x7b1+1}}', '{1+1}')
        self.assertEqual(f'\x7b1+1', '{1+1')
        self.assertEqual(f'\u007b1+1', '{1+1')
        self.assertEqual(f'\N{LEFT CURLY BRACKET}1+1\N{RIGHT CURLY BRACKET}', '{1+1}')

    def test_newlines_in_expressions(self):
        self.assertEqual(f'{0}', '0')
        self.assertEqual(rf'''{3+
4}''', '7')

    def test_lambda(self):
        x = 5
        self.assertEqual(f'{(lambda y:x*y)("8")!r}', "'88888'")
        self.assertEqual(f'{(lambda y:x*y)("8")!r:10}', "'88888'   ")
        self.assertEqual(f'{(lambda y:x*y)("8"):10}', "88888     ")

        # lambda doesn't work without parens, because the colon
        #  makes the parser think it's a format_spec
        err_msg = "invalid syntax" if use_old_parser() else "f-string: invalid syntax"
        self.assertAllRaise(SyntaxError, err_msg,
                            ["f'{lambda x:x}'",
                             ])

    def test_yield(self):
        # Not terribly useful, but make sure the yield turns
        #  a function into a generator
        def fn(y):
            f'y:{yield y*2}'
            f'{yield}'

        g = fn(4)
        self.assertEqual(next(g), 8)
        self.assertEqual(next(g), None)

    def test_yield_send(self):
        def fn(x):
            yield f'x:{yield (lambda i: x * i)}'

        g = fn(10)
        the_lambda = next(g)
        self.assertEqual(the_lambda(4), 40)
        self.assertEqual(g.send('string'), 'x:string')

    def test_expressions_with_triple_quoted_strings(self):
        self.assertEqual(f"{'''x'''}", 'x')
        self.assertEqual(f"{'''eric's'''}", "eric's")

        # Test concatenation within an expression
        self.assertEqual(f'{"x" """eric"s""" "y"}', 'xeric"sy')
        self.assertEqual(f'{"x" """eric"s"""}', 'xeric"s')
        self.assertEqual(f'{"""eric"s""" "y"}', 'eric"sy')
        self.assertEqual(f'{"""x""" """eric"s""" "y"}', 'xeric"sy')
        self.assertEqual(f'{"""x""" """eric"s""" """y"""}', 'xeric"sy')
        self.assertEqual(f'{r"""x""" """eric"s""" """y"""}', 'xeric"sy')

    def test_multiple_vars(self):
        x = 98
        y = 'abc'
        self.assertEqual(f'{x}{y}', '98abc')

        self.assertEqual(f'X{x}{y}', 'X98abc')
        self.assertEqual(f'{x}X{y}', '98Xabc')
        self.assertEqual(f'{x}{y}X', '98abcX')

        self.assertEqual(f'X{x}Y{y}', 'X98Yabc')
        self.assertEqual(f'X{x}{y}Y', 'X98abcY')
        self.assertEqual(f'{x}X{y}Y', '98XabcY')

        self.assertEqual(f'X{x}Y{y}Z', 'X98YabcZ')

    def test_closure(self):
        def outer(x):
            def inner():
                return f'x:{x}'
            return inner

        self.assertEqual(outer('987')(), 'x:987')
        self.assertEqual(outer(7)(), 'x:7')

    def test_arguments(self):
        y = 2
        def f(x, width):
            return f'x={x*y:{width}}'

        self.assertEqual(f('foo', 10), 'x=foofoo    ')
        x = 'bar'
        self.assertEqual(f(10, 10), 'x=        20')

    def test_locals(self):
        value = 123
        self.assertEqual(f'v:{value}', 'v:123')

    def test_missing_variable(self):
        with self.assertRaises(NameError):
            f'v:{value}'

    def test_missing_format_spec(self):
        class O:
            def __format__(self, spec):
                if not spec:
                    return '*'
                return spec

        self.assertEqual(f'{O():x}', 'x')
        self.assertEqual(f'{O()}', '*')
        self.assertEqual(f'{O():}', '*')

        self.assertEqual(f'{3:}', '3')
        self.assertEqual(f'{3!s:}', '3')

    def test_global(self):
        self.assertEqual(f'g:{a_global}', 'g:global variable')
        self.assertEqual(f'g:{a_global!r}', "g:'global variable'")

        a_local = 'local variable'
        self.assertEqual(f'g:{a_global} l:{a_local}',
                         'g:global variable l:local variable')
        self.assertEqual(f'g:{a_global!r}',
                         "g:'global variable'")
        self.assertEqual(f'g:{a_global} l:{a_local!r}',
                         "g:global variable l:'local variable'")

        self.assertIn("module 'unittest' from", f'{unittest}')

    def test_shadowed_global(self):
        a_global = 'really a local'
        self.assertEqual(f'g:{a_global}', 'g:really a local')
        self.assertEqual(f'g:{a_global!r}', "g:'really a local'")

        a_local = 'local variable'
        self.assertEqual(f'g:{a_global} l:{a_local}',
                         'g:really a local l:local variable')
        self.assertEqual(f'g:{a_global!r}',
                         "g:'really a local'")
        self.assertEqual(f'g:{a_global} l:{a_local!r}',
                         "g:really a local l:'local variable'")

    def test_call(self):
        def foo(x):
            return 'x=' + str(x)

        self.assertEqual(f'{foo(10)}', 'x=10')

    def test_nested_fstrings(self):
        y = 5
        self.assertEqual(f'{f"{0}"*3}', '000')
        self.assertEqual(f'{f"{y}"*3}', '555')

    def test_invalid_string_prefixes(self):
        single_quote_cases = ["fu''",
                             "uf''",
                             "Fu''",
                             "fU''",
                             "Uf''",
                             "uF''",
                             "ufr''",
                             "urf''",
                             "fur''",
                             "fru''",
                             "rfu''",
                             "ruf''",
                             "FUR''",
                             "Fur''",
                             "fb''",
                             "fB''",
                             "Fb''",
                             "FB''",
                             "bf''",
                             "bF''",
                             "Bf''",
                             "BF''",]
        double_quote_cases = [case.replace("'", '"') for case in single_quote_cases]
        error_msg = (
            'invalid syntax'
            if use_old_parser()
            else 'unexpected EOF while parsing'
        )
        self.assertAllRaise(SyntaxError, error_msg,
                            single_quote_cases + double_quote_cases)

    def test_leading_trailing_spaces(self):
        self.assertEqual(f'{ 3}', '3')
        self.assertEqual(f'{  3}', '3')
        self.assertEqual(f'{3 }', '3')
        self.assertEqual(f'{3  }', '3')

        self.assertEqual(f'expr={ {x: y for x, y in [(1, 2), ]}}',
                         'expr={1: 2}')
        self.assertEqual(f'expr={ {x: y for x, y in [(1, 2), ]} }',
                         'expr={1: 2}')

    def test_not_equal(self):
        # There's a special test for this because there's a special
        #  case in the f-string parser to look for != as not ending an
        #  expression. Normally it would, while looking for !s or !r.

        self.assertEqual(f'{3!=4}', 'True')
        self.assertEqual(f'{3!=4:}', 'True')
        self.assertEqual(f'{3!=4!s}', 'True')
        self.assertEqual(f'{3!=4!s:.3}', 'Tru')

    def test_equal_equal(self):
        # Because an expression ending in = has special meaning,
        # there's a special test for ==. Make sure it works.

        self.assertEqual(f'{0==1}', 'False')

    def test_conversions(self):
        self.assertEqual(f'{3.14:10.10}', '      3.14')
        self.assertEqual(f'{3.14!s:10.10}', '3.14      ')
        self.assertEqual(f'{3.14!r:10.10}', '3.14      ')
        self.assertEqual(f'{3.14!a:10.10}', '3.14      ')

        self.assertEqual(f'{"a"}', 'a')
        self.assertEqual(f'{"a"!r}', "'a'")
        self.assertEqual(f'{"a"!a}', "'a'")

        # Not a conversion.
        self.assertEqual(f'{"a!r"}', "a!r")

        # Not a conversion, but show that ! is allowed in a format spec.
        self.assertEqual(f'{3.14:!<10.10}', '3.14!!!!!!')

        self.assertAllRaise(SyntaxError, 'f-string: invalid conversion character',
                            ["f'{3!g}'",
                             "f'{3!A}'",
                             "f'{3!3}'",
                             "f'{3!G}'",
                             "f'{3!!}'",
                             "f'{3!:}'",
                             "f'{3! s}'",  # no space before conversion char
                             ])

        self.assertAllRaise(SyntaxError, "f-string: expecting '}'",
                            ["f'{x!s{y}}'",
                             "f'{3!ss}'",
                             "f'{3!ss:}'",
                             "f'{3!ss:s}'",
                             ])

    def test_assignment(self):
        self.assertAllRaise(SyntaxError, 'invalid syntax',
                            ["f'' = 3",
                             "f'{0}' = x",
                             "f'{x}' = x",
                             ])

    def test_del(self):
        self.assertAllRaise(SyntaxError, 'invalid syntax',
                            ["del f''",
                             "del '' f''",
                             ])

    def test_mismatched_braces(self):
        self.assertAllRaise(SyntaxError, "f-string: single '}' is not allowed",
                            ["f'{{}'",
                             "f'{{}}}'",
                             "f'}'",
                             "f'x}'",
                             "f'x}x'",
                             r"f'\u007b}'",

                             # Can't have { or } in a format spec.
                             "f'{3:}>10}'",
                             "f'{3:}}>10}'",
                             ])

        self.assertAllRaise(SyntaxError, "f-string: expecting '}'",
                            ["f'{3:{{>10}'",
                             "f'{3'",
                             "f'{3!'",
                             "f'{3:'",
                             "f'{3!s'",
                             "f'{3!s:'",
                             "f'{3!s:3'",
                             "f'x{'",
                             "f'x{x'",
                             "f'{x'",
                             "f'{3:s'",
                             "f'{{{'",
                             "f'{{}}{'",
                             "f'{'",
                             ])

        # But these are just normal strings.
        self.assertEqual(f'{"{"}', '{')
        self.assertEqual(f'{"}"}', '}')
        self.assertEqual(f'{3:{"}"}>10}', '}}}}}}}}}3')
        self.assertEqual(f'{2:{"{"}>10}', '{{{{{{{{{2')

    def test_if_conditional(self):
        # There's special logic in compile.c to test if the
        #  conditional for an if (and while) are constants. Exercise
        #  that code.

        def test_fstring(x, expected):
            flag = 0
            if f'{x}':
                flag = 1
            else:
                flag = 2
            self.assertEqual(flag, expected)

        def test_concat_empty(x, expected):
            flag = 0
            if '' f'{x}':
                flag = 1
            else:
                flag = 2
            self.assertEqual(flag, expected)

        def test_concat_non_empty(x, expected):
            flag = 0
            if ' ' f'{x}':
                flag = 1
            else:
                flag = 2
            self.assertEqual(flag, expected)

        test_fstring('', 2)
        test_fstring(' ', 1)

        test_concat_empty('', 2)
        test_concat_empty(' ', 1)

        test_concat_non_empty('', 1)
        test_concat_non_empty(' ', 1)

    def test_empty_format_specifier(self):
        x = 'test'
        self.assertEqual(f'{x}', 'test')
        self.assertEqual(f'{x:}', 'test')
        self.assertEqual(f'{x!s:}', 'test')
        self.assertEqual(f'{x!r:}', "'test'")

    def test_str_format_differences(self):
        d = {'a': 'string',
             0: 'integer',
             }
        a = 0
        self.assertEqual(f'{d[0]}', 'integer')
        self.assertEqual(f'{d["a"]}', 'string')
        self.assertEqual(f'{d[a]}', 'integer')
        self.assertEqual('{d[a]}'.format(d=d), 'string')
        self.assertEqual('{d[0]}'.format(d=d), 'integer')

    def test_errors(self):
        # see issue 26287
        self.assertAllRaise(TypeError, 'unsupported',
                            [r"f'{(lambda: 0):x}'",
                             r"f'{(0,):x}'",
                             ])
        self.assertAllRaise(ValueError, 'Unknown format code',
                            [r"f'{1000:j}'",
                             r"f'{1000:j}'",
                            ])

    @unittest.skipIf(use_old_parser(), "The old parser only supports <fstring> as the filename")
    def test_filename_in_syntaxerror(self):
        # see issue 38964
        with temp_cwd() as cwd:
            file_path = os.path.join(cwd, 't.py')
            with open(file_path, 'w') as f:
                f.write('f"{a b}"') # This generates a SyntaxError
            _, _, stderr = assert_python_failure(file_path,
                                                 PYTHONIOENCODING='ascii')
        self.assertIn(file_path.encode('ascii', 'backslashreplace'), stderr)

    def test_loop(self):
        for i in range(1000):
            self.assertEqual(f'i:{i}', 'i:' + str(i))

    def test_dict(self):
        d = {'"': 'dquote',
             "'": 'squote',
             'foo': 'bar',
             }
        self.assertEqual(f'''{d["'"]}''', 'squote')
        self.assertEqual(f"""{d['"']}""", 'dquote')

        self.assertEqual(f'{d["foo"]}', 'bar')
        self.assertEqual(f"{d['foo']}", 'bar')

    def test_backslash_char(self):
        # Check eval of a backslash followed by a control char.
        # See bpo-30682: this used to raise an assert in pydebug mode.
        self.assertEqual(eval('f"\\\n"'), '')
        self.assertEqual(eval('f"\\\r"'), '')

    def test_debug_conversion(self):
        x = 'A string'
        self.assertEqual(f'{x=}', 'x=' + repr(x))
        self.assertEqual(f'{x =}', 'x =' + repr(x))
        self.assertEqual(f'{x=!s}', 'x=' + str(x))
        self.assertEqual(f'{x=!r}', 'x=' + repr(x))
        self.assertEqual(f'{x=!a}', 'x=' + ascii(x))

        x = 2.71828
        self.assertEqual(f'{x=:.2f}', 'x=' + format(x, '.2f'))
        self.assertEqual(f'{x=:}', 'x=' + format(x, ''))
        self.assertEqual(f'{x=!r:^20}', 'x=' + format(repr(x), '^20'))
        self.assertEqual(f'{x=!s:^20}', 'x=' + format(str(x), '^20'))
        self.assertEqual(f'{x=!a:^20}', 'x=' + format(ascii(x), '^20'))

        x = 9
        self.assertEqual(f'{3*x+15=}', '3*x+15=42')

        # There is code in ast.c that deals with non-ascii expression values.  So,
        # use a unicode identifier to trigger that.
        tenπ = 31.4
        self.assertEqual(f'{tenπ=:.2f}', 'tenπ=31.40')

        # Also test with Unicode in non-identifiers.
        self.assertEqual(f'{"Σ"=}', '"Σ"=\'Σ\'')

        # Make sure nested fstrings still work.
        self.assertEqual(f'{f"{3.1415=:.1f}":*^20}', '*****3.1415=3.1*****')

        # Make sure text before and after an expression with = works
        # correctly.
        pi = 'π'
        self.assertEqual(f'alpha α {pi=} ω omega', "alpha α pi='π' ω omega")

        # Check multi-line expressions.
        self.assertEqual(f'''{
3
=}''', '\n3\n=3')

        # Since = is handled specially, make sure all existing uses of
        # it still work.

        self.assertEqual(f'{0==1}', 'False')
        self.assertEqual(f'{0!=1}', 'True')
        self.assertEqual(f'{0<=1}', 'True')
        self.assertEqual(f'{0>=1}', 'False')
        self.assertEqual(f'{(x:="5")}', '5')
        self.assertEqual(x, '5')
        self.assertEqual(f'{(x:=5)}', '5')
        self.assertEqual(x, 5)
        self.assertEqual(f'{"="}', '=')

        x = 20
        # This isn't an assignment expression, it's 'x', with a format
        # spec of '=10'.  See test_walrus: you need to use parens.
        self.assertEqual(f'{x:=10}', '        20')

        # Test named function parameters, to make sure '=' parsing works
        # there.
        def f(a):
            nonlocal x
            oldx = x
            x = a
            return oldx
        x = 0
        self.assertEqual(f'{f(a="3=")}', '0')
        self.assertEqual(x, '3=')
        self.assertEqual(f'{f(a=4)}', '3=')
        self.assertEqual(x, 4)

        # Make sure __format__ is being called.
        class C:
            def __format__(self, s):
                return f'FORMAT-{s}'
            def __repr__(self):
                return 'REPR'

        self.assertEqual(f'{C()=}', 'C()=REPR')
        self.assertEqual(f'{C()=!r}', 'C()=REPR')
        self.assertEqual(f'{C()=:}', 'C()=FORMAT-')
        self.assertEqual(f'{C()=: }', 'C()=FORMAT- ')
        self.assertEqual(f'{C()=:x}', 'C()=FORMAT-x')
        self.assertEqual(f'{C()=!r:*^20}', 'C()=********REPR********')

        self.assertRaises(SyntaxError, eval, "f'{C=]'")

        # Make sure leading and following text works.
        x = 'foo'
        self.assertEqual(f'X{x=}Y', 'Xx='+repr(x)+'Y')

        # Make sure whitespace around the = works.
        self.assertEqual(f'X{x  =}Y', 'Xx  ='+repr(x)+'Y')
        self.assertEqual(f'X{x=  }Y', 'Xx=  '+repr(x)+'Y')
        self.assertEqual(f'X{x  =  }Y', 'Xx  =  '+repr(x)+'Y')

        # These next lines contains tabs.  Backslash escapes don't
        # work in f-strings.
        # patchcheck doesn't like these tabs.  So the only way to test
        # this will be to dynamically created and exec the f-strings.  But
        # that's such a hassle I'll save it for another day.  For now, convert
        # the tabs to spaces just to shut up patchcheck.
        #self.assertEqual(f'X{x =}Y', 'Xx\t='+repr(x)+'Y')
        #self.assertEqual(f'X{x =       }Y', 'Xx\t=\t'+repr(x)+'Y')

    def test_walrus(self):
        x = 20
        # This isn't an assignment expression, it's 'x', with a format
        # spec of '=10'.
        self.assertEqual(f'{x:=10}', '        20')

        # This is an assignment expression, which requires parens.
        self.assertEqual(f'{(x:=10)}', '10')
        self.assertEqual(x, 10)

    def test_invalid_syntax_error_message(self):
        err_msg = "invalid syntax" if use_old_parser() else "f-string: invalid syntax"
        with self.assertRaisesRegex(SyntaxError, err_msg):
            compile("f'{a $ b}'", "?", "exec")

    def test_with_two_commas_in_format_specifier(self):
        error_msg = re.escape("Cannot specify ',' with ','.")
        with self.assertRaisesRegex(ValueError, error_msg):
            f'{1:,,}'

    def test_with_two_underscore_in_format_specifier(self):
        error_msg = re.escape("Cannot specify '_' with '_'.")
        with self.assertRaisesRegex(ValueError, error_msg):
            f'{1:__}'

    def test_with_a_commas_and_an_underscore_in_format_specifier(self):
        error_msg = re.escape("Cannot specify both ',' and '_'.")
        with self.assertRaisesRegex(ValueError, error_msg):
            f'{1:,_}'

    def test_with_an_underscore_and_a_comma_in_format_specifier(self):
        error_msg = re.escape("Cannot specify both ',' and '_'.")
        with self.assertRaisesRegex(ValueError, error_msg):
            f'{1:_,}'

if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_fstring.py
DragonHPC/dragon
from __future__ import unicode_literals
import unittest


class Tests(unittest.TestCase):
    def test_unicode_literals(self):
        self.assertIsInstance("literal", str)


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_future4.py
DragonHPC/dragon
# test_getopt.py
# David Goodger <dgoodger@bigfoot.com> 2000-08-19

from test.support import verbose, run_doctest, EnvironmentVarGuard
import unittest

import getopt

sentinel = object()

class GetoptTests(unittest.TestCase):
    def setUp(self):
        self.env = EnvironmentVarGuard()
        if "POSIXLY_CORRECT" in self.env:
            del self.env["POSIXLY_CORRECT"]

    def tearDown(self):
        self.env.__exit__()
        del self.env

    def assertError(self, *args, **kwargs):
        self.assertRaises(getopt.GetoptError, *args, **kwargs)

    def test_short_has_arg(self):
        self.assertTrue(getopt.short_has_arg('a', 'a:'))
        self.assertFalse(getopt.short_has_arg('a', 'a'))
        self.assertError(getopt.short_has_arg, 'a', 'b')

    def test_long_has_args(self):
        has_arg, option = getopt.long_has_args('abc', ['abc='])
        self.assertTrue(has_arg)
        self.assertEqual(option, 'abc')

        has_arg, option = getopt.long_has_args('abc', ['abc'])
        self.assertFalse(has_arg)
        self.assertEqual(option, 'abc')

        has_arg, option = getopt.long_has_args('abc', ['abcd'])
        self.assertFalse(has_arg)
        self.assertEqual(option, 'abcd')

        self.assertError(getopt.long_has_args, 'abc', ['def'])
        self.assertError(getopt.long_has_args, 'abc', [])
        self.assertError(getopt.long_has_args, 'abc', ['abcd','abcde'])

    def test_do_shorts(self):
        opts, args = getopt.do_shorts([], 'a', 'a', [])
        self.assertEqual(opts, [('-a', '')])
        self.assertEqual(args, [])

        opts, args = getopt.do_shorts([], 'a1', 'a:', [])
        self.assertEqual(opts, [('-a', '1')])
        self.assertEqual(args, [])

        #opts, args = getopt.do_shorts([], 'a=1', 'a:', [])
        #self.assertEqual(opts, [('-a', '1')])
        #self.assertEqual(args, [])

        opts, args = getopt.do_shorts([], 'a', 'a:', ['1'])
        self.assertEqual(opts, [('-a', '1')])
        self.assertEqual(args, [])

        opts, args = getopt.do_shorts([], 'a', 'a:', ['1', '2'])
        self.assertEqual(opts, [('-a', '1')])
        self.assertEqual(args, ['2'])

        self.assertError(getopt.do_shorts, [], 'a1', 'a', [])
        self.assertError(getopt.do_shorts, [], 'a', 'a:', [])

    def test_do_longs(self):
        opts, args = getopt.do_longs([], 'abc', ['abc'], [])
        self.assertEqual(opts, [('--abc', '')])
        self.assertEqual(args, [])

        opts, args = getopt.do_longs([], 'abc=1', ['abc='], [])
        self.assertEqual(opts, [('--abc', '1')])
        self.assertEqual(args, [])

        opts, args = getopt.do_longs([], 'abc=1', ['abcd='], [])
        self.assertEqual(opts, [('--abcd', '1')])
        self.assertEqual(args, [])

        opts, args = getopt.do_longs([], 'abc', ['ab', 'abc', 'abcd'], [])
        self.assertEqual(opts, [('--abc', '')])
        self.assertEqual(args, [])

        # Much like the preceding, except with a non-alpha character ("-") in
        # option name that precedes "="; failed in
        # http://python.org/sf/126863
        opts, args = getopt.do_longs([], 'foo=42', ['foo-bar', 'foo=',], [])
        self.assertEqual(opts, [('--foo', '42')])
        self.assertEqual(args, [])

        self.assertError(getopt.do_longs, [], 'abc=1', ['abc'], [])
        self.assertError(getopt.do_longs, [], 'abc', ['abc='], [])

    def test_getopt(self):
        # note: the empty string between '-a' and '--beta' is significant:
        # it simulates an empty string option argument ('-a ""') on the
        # command line.
        cmdline = ['-a', '1', '-b', '--alpha=2', '--beta', '-a', '3', '-a',
                   '', '--beta', 'arg1', 'arg2']

        opts, args = getopt.getopt(cmdline, 'a:b', ['alpha=', 'beta'])
        self.assertEqual(opts, [('-a', '1'), ('-b', ''),
                                ('--alpha', '2'), ('--beta', ''),
                                ('-a', '3'), ('-a', ''), ('--beta', '')])
        # Note ambiguity of ('-b', '') and ('-a', '') above. This must be
        # accounted for in the code that calls getopt().
        self.assertEqual(args, ['arg1', 'arg2'])

        self.assertError(getopt.getopt, cmdline, 'a:b', ['alpha', 'beta'])

    def test_gnu_getopt(self):
        # Test handling of GNU style scanning mode.
        cmdline = ['-a', 'arg1', '-b', '1', '--alpha', '--beta=2']

        # GNU style
        opts, args = getopt.gnu_getopt(cmdline, 'ab:', ['alpha', 'beta='])
        self.assertEqual(args, ['arg1'])
        self.assertEqual(opts, [('-a', ''), ('-b', '1'),
                                ('--alpha', ''), ('--beta', '2')])

        # recognize "-" as an argument
        opts, args = getopt.gnu_getopt(['-a', '-', '-b', '-'], 'ab:', [])
        self.assertEqual(args, ['-'])
        self.assertEqual(opts, [('-a', ''), ('-b', '-')])

        # Posix style via +
        opts, args = getopt.gnu_getopt(cmdline, '+ab:', ['alpha', 'beta='])
        self.assertEqual(opts, [('-a', '')])
        self.assertEqual(args, ['arg1', '-b', '1', '--alpha', '--beta=2'])

        # Posix style via POSIXLY_CORRECT
        self.env["POSIXLY_CORRECT"] = "1"
        opts, args = getopt.gnu_getopt(cmdline, 'ab:', ['alpha', 'beta='])
        self.assertEqual(opts, [('-a', '')])
        self.assertEqual(args, ['arg1', '-b', '1', '--alpha', '--beta=2'])

    def test_libref_examples(self):
        s = """
        Examples from the Library Reference:  Doc/lib/libgetopt.tex

        An example using only Unix style options:


        >>> import getopt
        >>> args = '-a -b -cfoo -d bar a1 a2'.split()
        >>> args
        ['-a', '-b', '-cfoo', '-d', 'bar', 'a1', 'a2']
        >>> optlist, args = getopt.getopt(args, 'abc:d:')
        >>> optlist
        [('-a', ''), ('-b', ''), ('-c', 'foo'), ('-d', 'bar')]
        >>> args
        ['a1', 'a2']

        Using long option names is equally easy:


        >>> s = '--condition=foo --testing --output-file abc.def -x a1 a2'
        >>> args = s.split()
        >>> args
        ['--condition=foo', '--testing', '--output-file', 'abc.def', '-x', 'a1', 'a2']
        >>> optlist, args = getopt.getopt(args, 'x', [
        ...     'condition=', 'output-file=', 'testing'])
        >>> optlist
        [('--condition', 'foo'), ('--testing', ''), ('--output-file', 'abc.def'), ('-x', '')]
        >>> args
        ['a1', 'a2']
        """

        import types
        m = types.ModuleType("libreftest", s)
        run_doctest(m, verbose)

    def test_issue4629(self):
        longopts, shortopts = getopt.getopt(['--help='], '', ['help='])
        self.assertEqual(longopts, [('--help', '')])
        longopts, shortopts = getopt.getopt(['--help=x'], '', ['help='])
        self.assertEqual(longopts, [('--help', 'x')])
        self.assertRaises(getopt.GetoptError, getopt.getopt, ['--help='], '', ['help'])

if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_getopt.py
DragonHPC/dragon
from .subpkg import util
from . import rebinding
util = util.util

external/cpython_master/Lib/test/test_import/data/circular_imports/rebinding2.py
DragonHPC/dragon
import package.submodule
package.submodule

external/cpython_master/Lib/test/test_import/data/package/__init__.py
DragonHPC/dragon
import sys
sys.modules.pop(__package__, None)
from . import submodule2

external/cpython_master/Lib/test/test_import/data/package2/submodule1.py
DragonHPC/dragon
from .. import abc
from .. import util

machinery = util.import_importlib('importlib.machinery')

from test.support import captured_stdout
import types
import unittest
import warnings


class ExecModuleTests(abc.LoaderTests):

    def exec_module(self, name):
        with util.uncache(name), captured_stdout() as stdout:
            spec = self.machinery.ModuleSpec(
                    name, self.machinery.FrozenImporter, origin='frozen',
                    is_package=self.machinery.FrozenImporter.is_package(name))
            module = types.ModuleType(name)
            module.__spec__ = spec
            assert not hasattr(module, 'initialized')
            self.machinery.FrozenImporter.exec_module(module)
            self.assertTrue(module.initialized)
            self.assertTrue(hasattr(module, '__spec__'))
            self.assertEqual(module.__spec__.origin, 'frozen')
            return module, stdout.getvalue()

    def test_module(self):
        name = '__hello__'
        module, output = self.exec_module(name)
        check = {'__name__': name}
        for attr, value in check.items():
            self.assertEqual(getattr(module, attr), value)
        self.assertEqual(output, 'Hello world!\n')
        self.assertTrue(hasattr(module, '__spec__'))

    def test_package(self):
        name = '__phello__'
        module, output = self.exec_module(name)
        check = {'__name__': name}
        for attr, value in check.items():
            attr_value = getattr(module, attr)
            self.assertEqual(attr_value, value,
                        'for {name}.{attr}, {given!r} != {expected!r}'.format(
                                 name=name, attr=attr, given=attr_value,
                                 expected=value))
        self.assertEqual(output, 'Hello world!\n')

    def test_lacking_parent(self):
        name = '__phello__.spam'
        with util.uncache('__phello__'):
            module, output = self.exec_module(name)
            check = {'__name__': name}
            for attr, value in check.items():
                attr_value = getattr(module, attr)
                self.assertEqual(attr_value, value,
                        'for {name}.{attr}, {given} != {expected!r}'.format(
                                 name=name, attr=attr, given=attr_value,
                                 expected=value))
            self.assertEqual(output, 'Hello world!\n')

    def test_module_repr(self):
        name = '__hello__'
        module, output = self.exec_module(name)
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', DeprecationWarning)
            repr_str = self.machinery.FrozenImporter.module_repr(module)
        self.assertEqual(repr_str,
                         "<module '__hello__' (frozen)>")

    def test_module_repr_indirect(self):
        name = '__hello__'
        module, output = self.exec_module(name)
        self.assertEqual(repr(module),
                         "<module '__hello__' (frozen)>")

    # No way to trigger an error in a frozen module.
    test_state_after_failure = None

    def test_unloadable(self):
        assert self.machinery.FrozenImporter.find_module('_not_real') is None
        with self.assertRaises(ImportError) as cm:
            self.exec_module('_not_real')
        self.assertEqual(cm.exception.name, '_not_real')


(Frozen_ExecModuleTests,
 Source_ExecModuleTests
 ) = util.test_both(ExecModuleTests, machinery=machinery)


class LoaderTests(abc.LoaderTests):

    def test_module(self):
        with util.uncache('__hello__'), captured_stdout() as stdout:
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', DeprecationWarning)
                module = self.machinery.FrozenImporter.load_module('__hello__')
            check = {'__name__': '__hello__',
                    '__package__': '',
                    '__loader__': self.machinery.FrozenImporter,
                    }
            for attr, value in check.items():
                self.assertEqual(getattr(module, attr), value)
            self.assertEqual(stdout.getvalue(), 'Hello world!\n')
            self.assertFalse(hasattr(module, '__file__'))

    def test_package(self):
        with util.uncache('__phello__'),  captured_stdout() as stdout:
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', DeprecationWarning)
                module = self.machinery.FrozenImporter.load_module('__phello__')
            check = {'__name__': '__phello__',
                     '__package__': '__phello__',
                     '__path__': [],
                     '__loader__': self.machinery.FrozenImporter,
                     }
            for attr, value in check.items():
                attr_value = getattr(module, attr)
                self.assertEqual(attr_value, value,
                                 "for __phello__.%s, %r != %r" %
                                 (attr, attr_value, value))
            self.assertEqual(stdout.getvalue(), 'Hello world!\n')
            self.assertFalse(hasattr(module, '__file__'))

    def test_lacking_parent(self):
        with util.uncache('__phello__', '__phello__.spam'), \
             captured_stdout() as stdout:
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', DeprecationWarning)
                module = self.machinery.FrozenImporter.load_module('__phello__.spam')
            check = {'__name__': '__phello__.spam',
                    '__package__': '__phello__',
                    '__loader__': self.machinery.FrozenImporter,
                    }
            for attr, value in check.items():
                attr_value = getattr(module, attr)
                self.assertEqual(attr_value, value,
                                 "for __phello__.spam.%s, %r != %r" %
                                 (attr, attr_value, value))
            self.assertEqual(stdout.getvalue(), 'Hello world!\n')
            self.assertFalse(hasattr(module, '__file__'))

    def test_module_reuse(self):
        with util.uncache('__hello__'), captured_stdout() as stdout:
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', DeprecationWarning)
                module1 = self.machinery.FrozenImporter.load_module('__hello__')
                module2 = self.machinery.FrozenImporter.load_module('__hello__')
            self.assertIs(module1, module2)
            self.assertEqual(stdout.getvalue(),
                             'Hello world!\nHello world!\n')

    def test_module_repr(self):
        with util.uncache('__hello__'), captured_stdout():
            with warnings.catch_warnings():
                warnings.simplefilter('ignore', DeprecationWarning)
                module = self.machinery.FrozenImporter.load_module('__hello__')
                repr_str = self.machinery.FrozenImporter.module_repr(module)
            self.assertEqual(repr_str,
                             "<module '__hello__' (frozen)>")

    def test_module_repr_indirect(self):
        with util.uncache('__hello__'), captured_stdout():
            module = self.machinery.FrozenImporter.load_module('__hello__')
        self.assertEqual(repr(module),
                         "<module '__hello__' (frozen)>")

    # No way to trigger an error in a frozen module.
    test_state_after_failure = None

    def test_unloadable(self):
        assert self.machinery.FrozenImporter.find_module('_not_real') is None
        with self.assertRaises(ImportError) as cm:
            self.machinery.FrozenImporter.load_module('_not_real')
        self.assertEqual(cm.exception.name, '_not_real')


(Frozen_LoaderTests,
 Source_LoaderTests
 ) = util.test_both(LoaderTests, machinery=machinery)


class InspectLoaderTests:

    """Tests for the InspectLoader methods for FrozenImporter."""

    def test_get_code(self):
        # Make sure that the code object is good.
        name = '__hello__'
        with captured_stdout() as stdout:
            code = self.machinery.FrozenImporter.get_code(name)
            mod = types.ModuleType(name)
            exec(code, mod.__dict__)
            self.assertTrue(hasattr(mod, 'initialized'))
            self.assertEqual(stdout.getvalue(), 'Hello world!\n')

    def test_get_source(self):
        # Should always return None.
        result = self.machinery.FrozenImporter.get_source('__hello__')
        self.assertIsNone(result)

    def test_is_package(self):
        # Should be able to tell what is a package.
        test_for = (('__hello__', False), ('__phello__', True),
                    ('__phello__.spam', False))
        for name, is_package in test_for:
            result = self.machinery.FrozenImporter.is_package(name)
            self.assertEqual(bool(result), is_package)

    def test_failure(self):
        # Raise ImportError for modules that are not frozen.
        for meth_name in ('get_code', 'get_source', 'is_package'):
            method = getattr(self.machinery.FrozenImporter, meth_name)
            with self.assertRaises(ImportError) as cm:
                method('importlib')
            self.assertEqual(cm.exception.name, 'importlib')

(Frozen_ILTests,
 Source_ILTests
 ) = util.test_both(InspectLoaderTests, machinery=machinery)


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_importlib/frozen/test_loader.py
DragonHPC/dragon
"""Test that the semantics relating to the 'fromlist' argument are correct."""
from .. import util
import warnings
import unittest


class ReturnValue:

    """The use of fromlist influences what import returns.

    If direct ``import ...`` statement is used, the root module or package is
    returned [import return]. But if fromlist is set, then the specified module
    is actually returned (whether it is a relative import or not)
    [from return].

    """

    def test_return_from_import(self):
        # [import return]
        with util.mock_spec('pkg.__init__', 'pkg.module') as importer:
            with util.import_state(meta_path=[importer]):
                module = self.__import__('pkg.module')
                self.assertEqual(module.__name__, 'pkg')

    def test_return_from_from_import(self):
        # [from return]
        with util.mock_modules('pkg.__init__', 'pkg.module')as importer:
            with util.import_state(meta_path=[importer]):
                module = self.__import__('pkg.module', fromlist=['attr'])
                self.assertEqual(module.__name__, 'pkg.module')


(Frozen_ReturnValue,
 Source_ReturnValue
 ) = util.test_both(ReturnValue, __import__=util.__import__)


class HandlingFromlist:

    """Using fromlist triggers different actions based on what is being asked
    of it.

    If fromlist specifies an object on a module, nothing special happens
    [object case]. This is even true if the object does not exist [bad object].

    If a package is being imported, then what is listed in fromlist may be
    treated as a module to be imported [module]. And this extends to what is
    contained in __all__ when '*' is imported [using *]. And '*' does not need
    to be the only name in the fromlist [using * with others].

    """

    def test_object(self):
        # [object case]
        with util.mock_modules('module') as importer:
            with util.import_state(meta_path=[importer]):
                module = self.__import__('module', fromlist=['attr'])
                self.assertEqual(module.__name__, 'module')

    def test_nonexistent_object(self):
        # [bad object]
        with util.mock_modules('module') as importer:
            with util.import_state(meta_path=[importer]):
                module = self.__import__('module', fromlist=['non_existent'])
                self.assertEqual(module.__name__, 'module')
                self.assertFalse(hasattr(module, 'non_existent'))

    def test_module_from_package(self):
        # [module]
        with util.mock_modules('pkg.__init__', 'pkg.module') as importer:
            with util.import_state(meta_path=[importer]):
                module = self.__import__('pkg', fromlist=['module'])
                self.assertEqual(module.__name__, 'pkg')
                self.assertTrue(hasattr(module, 'module'))
                self.assertEqual(module.module.__name__, 'pkg.module')

    def test_nonexistent_from_package(self):
        with util.mock_modules('pkg.__init__') as importer:
            with util.import_state(meta_path=[importer]):
                module = self.__import__('pkg', fromlist=['non_existent'])
                self.assertEqual(module.__name__, 'pkg')
                self.assertFalse(hasattr(module, 'non_existent'))

    def test_module_from_package_triggers_ModuleNotFoundError(self):
        # If a submodule causes an ModuleNotFoundError because it tries
        # to import a module which doesn't exist, that should let the
        # ModuleNotFoundError propagate.
        def module_code():
            import i_do_not_exist
        with util.mock_modules('pkg.__init__', 'pkg.mod',
                               module_code={'pkg.mod': module_code}) as importer:
            with util.import_state(meta_path=[importer]):
                with self.assertRaises(ModuleNotFoundError) as exc:
                    self.__import__('pkg', fromlist=['mod'])
                self.assertEqual('i_do_not_exist', exc.exception.name)

    def test_empty_string(self):
        with util.mock_modules('pkg.__init__', 'pkg.mod') as importer:
            with util.import_state(meta_path=[importer]):
                module = self.__import__('pkg.mod', fromlist=[''])
                self.assertEqual(module.__name__, 'pkg.mod')

    def basic_star_test(self, fromlist=['*']):
        # [using *]
        with util.mock_modules('pkg.__init__', 'pkg.module') as mock:
            with util.import_state(meta_path=[mock]):
                mock['pkg'].__all__ = ['module']
                module = self.__import__('pkg', fromlist=fromlist)
                self.assertEqual(module.__name__, 'pkg')
                self.assertTrue(hasattr(module, 'module'))
                self.assertEqual(module.module.__name__, 'pkg.module')

    def test_using_star(self):
        # [using *]
        self.basic_star_test()

    def test_fromlist_as_tuple(self):
        self.basic_star_test(('*',))

    def test_star_with_others(self):
        # [using * with others]
        context = util.mock_modules('pkg.__init__', 'pkg.module1', 'pkg.module2')
        with context as mock:
            with util.import_state(meta_path=[mock]):
                mock['pkg'].__all__ = ['module1']
                module = self.__import__('pkg', fromlist=['module2', '*'])
                self.assertEqual(module.__name__, 'pkg')
                self.assertTrue(hasattr(module, 'module1'))
                self.assertTrue(hasattr(module, 'module2'))
                self.assertEqual(module.module1.__name__, 'pkg.module1')
                self.assertEqual(module.module2.__name__, 'pkg.module2')

    def test_nonexistent_in_all(self):
        with util.mock_modules('pkg.__init__') as importer:
            with util.import_state(meta_path=[importer]):
                importer['pkg'].__all__ = ['non_existent']
                module = self.__import__('pkg', fromlist=['*'])
                self.assertEqual(module.__name__, 'pkg')
                self.assertFalse(hasattr(module, 'non_existent'))

    def test_star_in_all(self):
        with util.mock_modules('pkg.__init__') as importer:
            with util.import_state(meta_path=[importer]):
                importer['pkg'].__all__ = ['*']
                module = self.__import__('pkg', fromlist=['*'])
                self.assertEqual(module.__name__, 'pkg')
                self.assertFalse(hasattr(module, '*'))

    def test_invalid_type(self):
        with util.mock_modules('pkg.__init__') as importer:
            with util.import_state(meta_path=[importer]), \
                 warnings.catch_warnings():
                warnings.simplefilter('error', BytesWarning)
                with self.assertRaisesRegex(TypeError, r'\bfrom\b'):
                    self.__import__('pkg', fromlist=[b'attr'])
                with self.assertRaisesRegex(TypeError, r'\bfrom\b'):
                    self.__import__('pkg', fromlist=iter([b'attr']))

    def test_invalid_type_in_all(self):
        with util.mock_modules('pkg.__init__') as importer:
            with util.import_state(meta_path=[importer]), \
                 warnings.catch_warnings():
                warnings.simplefilter('error', BytesWarning)
                importer['pkg'].__all__ = [b'attr']
                with self.assertRaisesRegex(TypeError, r'\bpkg\.__all__\b'):
                    self.__import__('pkg', fromlist=['*'])


(Frozen_FromList,
 Source_FromList
 ) = util.test_both(HandlingFromlist, __import__=util.__import__)


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_importlib/import_/test_fromlist.py
DragonHPC/dragon
from .. import util
import sys
import unittest
from test import support


class ParentModuleTests:

    """Importing a submodule should import the parent modules."""

    def test_import_parent(self):
        with util.mock_spec('pkg.__init__', 'pkg.module') as mock:
            with util.import_state(meta_path=[mock]):
                module = self.__import__('pkg.module')
                self.assertIn('pkg', sys.modules)

    def test_bad_parent(self):
        with util.mock_spec('pkg.module') as mock:
            with util.import_state(meta_path=[mock]):
                with self.assertRaises(ImportError) as cm:
                    self.__import__('pkg.module')
                self.assertEqual(cm.exception.name, 'pkg')

    def test_raising_parent_after_importing_child(self):
        def __init__():
            import pkg.module
            1/0
        mock = util.mock_spec('pkg.__init__', 'pkg.module',
                                 module_code={'pkg': __init__})
        with mock:
            with util.import_state(meta_path=[mock]):
                with self.assertRaises(ZeroDivisionError):
                    self.__import__('pkg')
                self.assertNotIn('pkg', sys.modules)
                self.assertIn('pkg.module', sys.modules)
                with self.assertRaises(ZeroDivisionError):
                    self.__import__('pkg.module')
                self.assertNotIn('pkg', sys.modules)
                self.assertIn('pkg.module', sys.modules)

    def test_raising_parent_after_relative_importing_child(self):
        def __init__():
            from . import module
            1/0
        mock = util.mock_spec('pkg.__init__', 'pkg.module',
                                 module_code={'pkg': __init__})
        with mock:
            with util.import_state(meta_path=[mock]):
                with self.assertRaises((ZeroDivisionError, ImportError)):
                    # This raises ImportError on the "from . import module"
                    # line, not sure why.
                    self.__import__('pkg')
                self.assertNotIn('pkg', sys.modules)
                with self.assertRaises((ZeroDivisionError, ImportError)):
                    self.__import__('pkg.module')
                self.assertNotIn('pkg', sys.modules)
                # XXX False
                #self.assertIn('pkg.module', sys.modules)

    def test_raising_parent_after_double_relative_importing_child(self):
        def __init__():
            from ..subpkg import module
            1/0
        mock = util.mock_spec('pkg.__init__', 'pkg.subpkg.__init__',
                                 'pkg.subpkg.module',
                                 module_code={'pkg.subpkg': __init__})
        with mock:
            with util.import_state(meta_path=[mock]):
                with self.assertRaises((ZeroDivisionError, ImportError)):
                    # This raises ImportError on the "from ..subpkg import module"
                    # line, not sure why.
                    self.__import__('pkg.subpkg')
                self.assertNotIn('pkg.subpkg', sys.modules)
                with self.assertRaises((ZeroDivisionError, ImportError)):
                    self.__import__('pkg.subpkg.module')
                self.assertNotIn('pkg.subpkg', sys.modules)
                # XXX False
                #self.assertIn('pkg.subpkg.module', sys.modules)

    def test_module_not_package(self):
        # Try to import a submodule from a non-package should raise ImportError.
        assert not hasattr(sys, '__path__')
        with self.assertRaises(ImportError) as cm:
            self.__import__('sys.no_submodules_here')
        self.assertEqual(cm.exception.name, 'sys.no_submodules_here')

    def test_module_not_package_but_side_effects(self):
        # If a module injects something into sys.modules as a side-effect, then
        # pick up on that fact.
        name = 'mod'
        subname = name + '.b'
        def module_injection():
            sys.modules[subname] = 'total bunk'
        mock_spec = util.mock_spec('mod',
                                         module_code={'mod': module_injection})
        with mock_spec as mock:
            with util.import_state(meta_path=[mock]):
                try:
                    submodule = self.__import__(subname)
                finally:
                    support.unload(subname)


(Frozen_ParentTests,
 Source_ParentTests
 ) = util.test_both(ParentModuleTests, __import__=util.__import__)


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_importlib/import_/test_packages.py
DragonHPC/dragon
from . import util as test_util

init = test_util.import_importlib('importlib')
machinery = test_util.import_importlib('importlib.machinery')
util = test_util.import_importlib('importlib.util')

import os.path
import pathlib
from test.support import CleanImport
import unittest
import sys
import warnings



class TestLoader:

    def __init__(self, path=None, is_package=None):
        self.path = path
        self.package = is_package

    def __repr__(self):
        return '<TestLoader object>'

    def __getattr__(self, name):
        if name == 'get_filename' and self.path is not None:
            return self._get_filename
        if name == 'is_package':
            return self._is_package
        raise AttributeError(name)

    def _get_filename(self, name):
        return self.path

    def _is_package(self, name):
        return self.package

    def create_module(self, spec):
        return None


class NewLoader(TestLoader):

    EGGS = 1

    def exec_module(self, module):
        module.eggs = self.EGGS


class LegacyLoader(TestLoader):

    HAM = -1

    with warnings.catch_warnings():
        warnings.simplefilter("ignore", DeprecationWarning)

        frozen_util = util['Frozen']

        @frozen_util.module_for_loader
        def load_module(self, module):
            module.ham = self.HAM
            return module


class ModuleSpecTests:

    def setUp(self):
        self.name = 'spam'
        self.path = 'spam.py'
        self.cached = self.util.cache_from_source(self.path)
        self.loader = TestLoader()
        self.spec = self.machinery.ModuleSpec(self.name, self.loader)
        self.loc_spec = self.machinery.ModuleSpec(self.name, self.loader,
                                                  origin=self.path)
        self.loc_spec._set_fileattr = True

    def test_default(self):
        spec = self.machinery.ModuleSpec(self.name, self.loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_default_no_loader(self):
        spec = self.machinery.ModuleSpec(self.name, None)

        self.assertEqual(spec.name, self.name)
        self.assertIs(spec.loader, None)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_default_is_package_false(self):
        spec = self.machinery.ModuleSpec(self.name, self.loader,
                                         is_package=False)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_default_is_package_true(self):
        spec = self.machinery.ModuleSpec(self.name, self.loader,
                                         is_package=True)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, [])
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_has_location_setter(self):
        spec = self.machinery.ModuleSpec(self.name, self.loader,
                                         origin='somewhere')
        self.assertFalse(spec.has_location)
        spec.has_location = True
        self.assertTrue(spec.has_location)

    def test_equality(self):
        other = type(sys.implementation)(name=self.name,
                                         loader=self.loader,
                                         origin=None,
                                         submodule_search_locations=None,
                                         has_location=False,
                                         cached=None,
                                         )

        self.assertTrue(self.spec == other)

    def test_equality_location(self):
        other = type(sys.implementation)(name=self.name,
                                         loader=self.loader,
                                         origin=self.path,
                                         submodule_search_locations=None,
                                         has_location=True,
                                         cached=self.cached,
                                         )

        self.assertEqual(self.loc_spec, other)

    def test_inequality(self):
        other = type(sys.implementation)(name='ham',
                                         loader=self.loader,
                                         origin=None,
                                         submodule_search_locations=None,
                                         has_location=False,
                                         cached=None,
                                         )

        self.assertNotEqual(self.spec, other)

    def test_inequality_incomplete(self):
        other = type(sys.implementation)(name=self.name,
                                         loader=self.loader,
                                         )

        self.assertNotEqual(self.spec, other)

    def test_package(self):
        spec = self.machinery.ModuleSpec('spam.eggs', self.loader)

        self.assertEqual(spec.parent, 'spam')

    def test_package_is_package(self):
        spec = self.machinery.ModuleSpec('spam.eggs', self.loader,
                                         is_package=True)

        self.assertEqual(spec.parent, 'spam.eggs')

    # cached

    def test_cached_set(self):
        before = self.spec.cached
        self.spec.cached = 'there'
        after = self.spec.cached

        self.assertIs(before, None)
        self.assertEqual(after, 'there')

    def test_cached_no_origin(self):
        spec = self.machinery.ModuleSpec(self.name, self.loader)

        self.assertIs(spec.cached, None)

    def test_cached_with_origin_not_location(self):
        spec = self.machinery.ModuleSpec(self.name, self.loader,
                                         origin=self.path)

        self.assertIs(spec.cached, None)

    def test_cached_source(self):
        expected = self.util.cache_from_source(self.path)

        self.assertEqual(self.loc_spec.cached, expected)

    def test_cached_source_unknown_suffix(self):
        self.loc_spec.origin = 'spam.spamspamspam'

        self.assertIs(self.loc_spec.cached, None)

    def test_cached_source_missing_cache_tag(self):
        original = sys.implementation.cache_tag
        sys.implementation.cache_tag = None
        try:
            cached = self.loc_spec.cached
        finally:
            sys.implementation.cache_tag = original

        self.assertIs(cached, None)

    def test_cached_sourceless(self):
        self.loc_spec.origin = 'spam.pyc'

        self.assertEqual(self.loc_spec.cached, 'spam.pyc')


(Frozen_ModuleSpecTests,
 Source_ModuleSpecTests
 ) = test_util.test_both(ModuleSpecTests, util=util, machinery=machinery)


class ModuleSpecMethodsTests:

    @property
    def bootstrap(self):
        return self.init._bootstrap

    def setUp(self):
        self.name = 'spam'
        self.path = 'spam.py'
        self.cached = self.util.cache_from_source(self.path)
        self.loader = TestLoader()
        self.spec = self.machinery.ModuleSpec(self.name, self.loader)
        self.loc_spec = self.machinery.ModuleSpec(self.name, self.loader,
                                                  origin=self.path)
        self.loc_spec._set_fileattr = True

    # exec()

    def test_exec(self):
        self.spec.loader = NewLoader()
        module = self.util.module_from_spec(self.spec)
        sys.modules[self.name] = module
        self.assertFalse(hasattr(module, 'eggs'))
        self.bootstrap._exec(self.spec, module)

        self.assertEqual(module.eggs, 1)

    # load()

    def test_load(self):
        self.spec.loader = NewLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)
            installed = sys.modules[self.spec.name]

        self.assertEqual(loaded.eggs, 1)
        self.assertIs(loaded, installed)

    def test_load_replaced(self):
        replacement = object()
        class ReplacingLoader(TestLoader):
            def exec_module(self, module):
                sys.modules[module.__name__] = replacement
        self.spec.loader = ReplacingLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)
            installed = sys.modules[self.spec.name]

        self.assertIs(loaded, replacement)
        self.assertIs(installed, replacement)

    def test_load_failed(self):
        class FailedLoader(TestLoader):
            def exec_module(self, module):
                raise RuntimeError
        self.spec.loader = FailedLoader()
        with CleanImport(self.spec.name):
            with self.assertRaises(RuntimeError):
                loaded = self.bootstrap._load(self.spec)
            self.assertNotIn(self.spec.name, sys.modules)

    def test_load_failed_removed(self):
        class FailedLoader(TestLoader):
            def exec_module(self, module):
                del sys.modules[module.__name__]
                raise RuntimeError
        self.spec.loader = FailedLoader()
        with CleanImport(self.spec.name):
            with self.assertRaises(RuntimeError):
                loaded = self.bootstrap._load(self.spec)
            self.assertNotIn(self.spec.name, sys.modules)

    def test_load_legacy(self):
        self.spec.loader = LegacyLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)

        self.assertEqual(loaded.ham, -1)

    def test_load_legacy_attributes(self):
        self.spec.loader = LegacyLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)

        self.assertIs(loaded.__loader__, self.spec.loader)
        self.assertEqual(loaded.__package__, self.spec.parent)
        self.assertIs(loaded.__spec__, self.spec)

    def test_load_legacy_attributes_immutable(self):
        module = object()
        class ImmutableLoader(TestLoader):
            def load_module(self, name):
                sys.modules[name] = module
                return module
        self.spec.loader = ImmutableLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)

            self.assertIs(sys.modules[self.spec.name], module)

    # reload()

    def test_reload(self):
        self.spec.loader = NewLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)
            reloaded = self.bootstrap._exec(self.spec, loaded)
            installed = sys.modules[self.spec.name]

        self.assertEqual(loaded.eggs, 1)
        self.assertIs(reloaded, loaded)
        self.assertIs(installed, loaded)

    def test_reload_modified(self):
        self.spec.loader = NewLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)
            loaded.eggs = 2
            reloaded = self.bootstrap._exec(self.spec, loaded)

        self.assertEqual(loaded.eggs, 1)
        self.assertIs(reloaded, loaded)

    def test_reload_extra_attributes(self):
        self.spec.loader = NewLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)
            loaded.available = False
            reloaded = self.bootstrap._exec(self.spec, loaded)

        self.assertFalse(loaded.available)
        self.assertIs(reloaded, loaded)

    def test_reload_init_module_attrs(self):
        self.spec.loader = NewLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)
            loaded.__name__ = 'ham'
            del loaded.__loader__
            del loaded.__package__
            del loaded.__spec__
            self.bootstrap._exec(self.spec, loaded)

        self.assertEqual(loaded.__name__, self.spec.name)
        self.assertIs(loaded.__loader__, self.spec.loader)
        self.assertEqual(loaded.__package__, self.spec.parent)
        self.assertIs(loaded.__spec__, self.spec)
        self.assertFalse(hasattr(loaded, '__path__'))
        self.assertFalse(hasattr(loaded, '__file__'))
        self.assertFalse(hasattr(loaded, '__cached__'))

    def test_reload_legacy(self):
        self.spec.loader = LegacyLoader()
        with CleanImport(self.spec.name):
            loaded = self.bootstrap._load(self.spec)
            reloaded = self.bootstrap._exec(self.spec, loaded)
            installed = sys.modules[self.spec.name]

        self.assertEqual(loaded.ham, -1)
        self.assertIs(reloaded, loaded)
        self.assertIs(installed, loaded)


(Frozen_ModuleSpecMethodsTests,
 Source_ModuleSpecMethodsTests
 ) = test_util.test_both(ModuleSpecMethodsTests, init=init, util=util,
                         machinery=machinery)


class ModuleReprTests:

    @property
    def bootstrap(self):
        return self.init._bootstrap

    def setUp(self):
        self.module = type(os)('spam')
        self.spec = self.machinery.ModuleSpec('spam', TestLoader())

    def test_module___loader___module_repr(self):
        class Loader:
            def module_repr(self, module):
                return '<delicious {}>'.format(module.__name__)
        self.module.__loader__ = Loader()
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr, '<delicious spam>')

    def test_module___loader___module_repr_bad(self):
        class Loader(TestLoader):
            def module_repr(self, module):
                raise Exception
        self.module.__loader__ = Loader()
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr,
                         '<module {!r} (<TestLoader object>)>'.format('spam'))

    def test_module___spec__(self):
        origin = 'in a hole, in the ground'
        self.spec.origin = origin
        self.module.__spec__ = self.spec
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr, '<module {!r} ({})>'.format('spam', origin))

    def test_module___spec___location(self):
        location = 'in_a_galaxy_far_far_away.py'
        self.spec.origin = location
        self.spec._set_fileattr = True
        self.module.__spec__ = self.spec
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr,
                         '<module {!r} from {!r}>'.format('spam', location))

    def test_module___spec___no_origin(self):
        self.spec.loader = TestLoader()
        self.module.__spec__ = self.spec
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr,
                         '<module {!r} (<TestLoader object>)>'.format('spam'))

    def test_module___spec___no_origin_no_loader(self):
        self.spec.loader = None
        self.module.__spec__ = self.spec
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr, '<module {!r}>'.format('spam'))

    def test_module_no_name(self):
        del self.module.__name__
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr, '<module {!r}>'.format('?'))

    def test_module_with_file(self):
        filename = 'e/i/e/i/o/spam.py'
        self.module.__file__ = filename
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr,
                         '<module {!r} from {!r}>'.format('spam', filename))

    def test_module_no_file(self):
        self.module.__loader__ = TestLoader()
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr,
                         '<module {!r} (<TestLoader object>)>'.format('spam'))

    def test_module_no_file_no_loader(self):
        modrepr = self.bootstrap._module_repr(self.module)

        self.assertEqual(modrepr, '<module {!r}>'.format('spam'))


(Frozen_ModuleReprTests,
 Source_ModuleReprTests
 ) = test_util.test_both(ModuleReprTests, init=init, util=util,
                         machinery=machinery)


class FactoryTests:

    def setUp(self):
        self.name = 'spam'
        self.path = 'spam.py'
        self.cached = self.util.cache_from_source(self.path)
        self.loader = TestLoader()
        self.fileloader = TestLoader(self.path)
        self.pkgloader = TestLoader(self.path, True)

    # spec_from_loader()

    def test_spec_from_loader_default(self):
        spec = self.util.spec_from_loader(self.name, self.loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_default_with_bad_is_package(self):
        class Loader:
            def is_package(self, name):
                raise ImportError
        loader = Loader()
        spec = self.util.spec_from_loader(self.name, loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_origin(self):
        origin = 'somewhere over the rainbow'
        spec = self.util.spec_from_loader(self.name, self.loader,
                                          origin=origin)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, origin)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_is_package_false(self):
        spec = self.util.spec_from_loader(self.name, self.loader,
                                          is_package=False)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_is_package_true(self):
        spec = self.util.spec_from_loader(self.name, self.loader,
                                          is_package=True)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, [])
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_origin_and_is_package(self):
        origin = 'where the streets have no name'
        spec = self.util.spec_from_loader(self.name, self.loader,
                                          origin=origin, is_package=True)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertIs(spec.origin, origin)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, [])
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_is_package_with_loader_false(self):
        loader = TestLoader(is_package=False)
        spec = self.util.spec_from_loader(self.name, loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_is_package_with_loader_true(self):
        loader = TestLoader(is_package=True)
        spec = self.util.spec_from_loader(self.name, loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, loader)
        self.assertIs(spec.origin, None)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, [])
        self.assertIs(spec.cached, None)
        self.assertFalse(spec.has_location)

    def test_spec_from_loader_default_with_file_loader(self):
        spec = self.util.spec_from_loader(self.name, self.fileloader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_loader_is_package_false_with_fileloader(self):
        spec = self.util.spec_from_loader(self.name, self.fileloader,
                                          is_package=False)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_loader_is_package_true_with_fileloader(self):
        spec = self.util.spec_from_loader(self.name, self.fileloader,
                                          is_package=True)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, [''])
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    # spec_from_file_location()

    def test_spec_from_file_location_default(self):
        spec = self.util.spec_from_file_location(self.name, self.path)

        self.assertEqual(spec.name, self.name)
        # Need to use a circuitous route to get at importlib.machinery to make
        # sure the same class object is used in the isinstance() check as
        # would have been used to create the loader.
        self.assertIsInstance(spec.loader,
                              self.util.abc.machinery.SourceFileLoader)
        self.assertEqual(spec.loader.name, self.name)
        self.assertEqual(spec.loader.path, self.path)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_path_like_arg(self):
        spec = self.util.spec_from_file_location(self.name,
                                                 pathlib.PurePath(self.path))
        self.assertEqual(spec.origin, self.path)

    def test_spec_from_file_location_default_without_location(self):
        spec = self.util.spec_from_file_location(self.name)

        self.assertIs(spec, None)

    def test_spec_from_file_location_default_bad_suffix(self):
        spec = self.util.spec_from_file_location(self.name, 'spam.eggs')

        self.assertIs(spec, None)

    def test_spec_from_file_location_loader_no_location(self):
        spec = self.util.spec_from_file_location(self.name,
                                                 loader=self.fileloader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_loader_no_location_no_get_filename(self):
        spec = self.util.spec_from_file_location(self.name,
                                                 loader=self.loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.loader)
        self.assertEqual(spec.origin, '<unknown>')
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_loader_no_location_bad_get_filename(self):
        class Loader:
            def get_filename(self, name):
                raise ImportError
        loader = Loader()
        spec = self.util.spec_from_file_location(self.name, loader=loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, loader)
        self.assertEqual(spec.origin, '<unknown>')
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertIs(spec.cached, None)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_smsl_none(self):
        spec = self.util.spec_from_file_location(self.name, self.path,
                                       loader=self.fileloader,
                                       submodule_search_locations=None)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_smsl_empty(self):
        spec = self.util.spec_from_file_location(self.name, self.path,
                                       loader=self.fileloader,
                                       submodule_search_locations=[])

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, [''])
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_smsl_not_empty(self):
        spec = self.util.spec_from_file_location(self.name, self.path,
                                       loader=self.fileloader,
                                       submodule_search_locations=['eggs'])

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, ['eggs'])
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_smsl_default(self):
        spec = self.util.spec_from_file_location(self.name, self.path,
                                       loader=self.pkgloader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.pkgloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertEqual(spec.submodule_search_locations, [''])
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_smsl_default_not_package(self):
        class Loader:
            def is_package(self, name):
                return False
        loader = Loader()
        spec = self.util.spec_from_file_location(self.name, self.path,
                                                 loader=loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, loader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_smsl_default_no_is_package(self):
        spec = self.util.spec_from_file_location(self.name, self.path,
                                       loader=self.fileloader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, self.fileloader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)

    def test_spec_from_file_location_smsl_default_bad_is_package(self):
        class Loader:
            def is_package(self, name):
                raise ImportError
        loader = Loader()
        spec = self.util.spec_from_file_location(self.name, self.path,
                                                 loader=loader)

        self.assertEqual(spec.name, self.name)
        self.assertEqual(spec.loader, loader)
        self.assertEqual(spec.origin, self.path)
        self.assertIs(spec.loader_state, None)
        self.assertIs(spec.submodule_search_locations, None)
        self.assertEqual(spec.cached, self.cached)
        self.assertTrue(spec.has_location)


(Frozen_FactoryTests,
 Source_FactoryTests
 ) = test_util.test_both(FactoryTests, util=util, machinery=machinery)


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_importlib/test_spec.py
DragonHPC/dragon
""" Test Iterator Length Transparency

Some functions or methods which accept general iterable arguments have
optional, more efficient code paths if they know how many items to expect.
For instance, map(func, iterable), will pre-allocate the exact amount of
space required whenever the iterable can report its length.

The desired invariant is:  len(it)==len(list(it)).

A complication is that an iterable and iterator can be the same object. To
maintain the invariant, an iterator needs to dynamically update its length.
For instance, an iterable such as range(10) always reports its length as ten,
but it=iter(range(10)) starts at ten, and then goes to nine after next(it).
Having this capability means that map() can ignore the distinction between
map(func, iterable) and map(func, iter(iterable)).

When the iterable is immutable, the implementation can straight-forwardly
report the original length minus the cumulative number of calls to next().
This is the case for tuples, range objects, and itertools.repeat().

Some containers become temporarily immutable during iteration.  This includes
dicts, sets, and collections.deque.  Their implementation is equally simple
though they need to permanently set their length to zero whenever there is
an attempt to iterate after a length mutation.

The situation slightly more involved whenever an object allows length mutation
during iteration.  Lists and sequence iterators are dynamically updatable.
So, if a list is extended during iteration, the iterator will continue through
the new items.  If it shrinks to a point before the most recent iteration,
then no further items are available and the length is reported at zero.

Reversed objects can also be wrapped around mutable objects; however, any
appends after the current position are ignored.  Any other approach leads
to confusion and possibly returning the same item more than once.

The iterators not listed above, such as enumerate and the other itertools,
are not length transparent because they have no way to distinguish between
iterables that report static length and iterators whose length changes with
each call (i.e. the difference between enumerate('abc') and
enumerate(iter('abc')).

"""

import unittest
from itertools import repeat
from collections import deque
from operator import length_hint

n = 10


class TestInvariantWithoutMutations:

    def test_invariant(self):
        it = self.it
        for i in reversed(range(1, n+1)):
            self.assertEqual(length_hint(it), i)
            next(it)
        self.assertEqual(length_hint(it), 0)
        self.assertRaises(StopIteration, next, it)
        self.assertEqual(length_hint(it), 0)

class TestTemporarilyImmutable(TestInvariantWithoutMutations):

    def test_immutable_during_iteration(self):
        # objects such as deques, sets, and dictionaries enforce
        # length immutability  during iteration

        it = self.it
        self.assertEqual(length_hint(it), n)
        next(it)
        self.assertEqual(length_hint(it), n-1)
        self.mutate()
        self.assertRaises(RuntimeError, next, it)
        self.assertEqual(length_hint(it), 0)

## ------- Concrete Type Tests -------

class TestRepeat(TestInvariantWithoutMutations, unittest.TestCase):

    def setUp(self):
        self.it = repeat(None, n)

class TestXrange(TestInvariantWithoutMutations, unittest.TestCase):

    def setUp(self):
        self.it = iter(range(n))

class TestXrangeCustomReversed(TestInvariantWithoutMutations, unittest.TestCase):

    def setUp(self):
        self.it = reversed(range(n))

class TestTuple(TestInvariantWithoutMutations, unittest.TestCase):

    def setUp(self):
        self.it = iter(tuple(range(n)))

## ------- Types that should not be mutated during iteration -------

class TestDeque(TestTemporarilyImmutable, unittest.TestCase):

    def setUp(self):
        d = deque(range(n))
        self.it = iter(d)
        self.mutate = d.pop

class TestDequeReversed(TestTemporarilyImmutable, unittest.TestCase):

    def setUp(self):
        d = deque(range(n))
        self.it = reversed(d)
        self.mutate = d.pop

class TestDictKeys(TestTemporarilyImmutable, unittest.TestCase):

    def setUp(self):
        d = dict.fromkeys(range(n))
        self.it = iter(d)
        self.mutate = d.popitem

class TestDictItems(TestTemporarilyImmutable, unittest.TestCase):

    def setUp(self):
        d = dict.fromkeys(range(n))
        self.it = iter(d.items())
        self.mutate = d.popitem

class TestDictValues(TestTemporarilyImmutable, unittest.TestCase):

    def setUp(self):
        d = dict.fromkeys(range(n))
        self.it = iter(d.values())
        self.mutate = d.popitem

class TestSet(TestTemporarilyImmutable, unittest.TestCase):

    def setUp(self):
        d = set(range(n))
        self.it = iter(d)
        self.mutate = d.pop

## ------- Types that can mutate during iteration -------

class TestList(TestInvariantWithoutMutations, unittest.TestCase):

    def setUp(self):
        self.it = iter(range(n))

    def test_mutation(self):
        d = list(range(n))
        it = iter(d)
        next(it)
        next(it)
        self.assertEqual(length_hint(it), n - 2)
        d.append(n)
        self.assertEqual(length_hint(it), n - 1)  # grow with append
        d[1:] = []
        self.assertEqual(length_hint(it), 0)
        self.assertEqual(list(it), [])
        d.extend(range(20))
        self.assertEqual(length_hint(it), 0)


class TestListReversed(TestInvariantWithoutMutations, unittest.TestCase):

    def setUp(self):
        self.it = reversed(range(n))

    def test_mutation(self):
        d = list(range(n))
        it = reversed(d)
        next(it)
        next(it)
        self.assertEqual(length_hint(it), n - 2)
        d.append(n)
        self.assertEqual(length_hint(it), n - 2)  # ignore append
        d[1:] = []
        self.assertEqual(length_hint(it), 0)
        self.assertEqual(list(it), [])  # confirm invariant
        d.extend(range(20))
        self.assertEqual(length_hint(it), 0)

## -- Check to make sure exceptions are not suppressed by __length_hint__()


class BadLen(object):
    def __iter__(self):
        return iter(range(10))

    def __len__(self):
        raise RuntimeError('hello')


class BadLengthHint(object):
    def __iter__(self):
        return iter(range(10))

    def __length_hint__(self):
        raise RuntimeError('hello')


class NoneLengthHint(object):
    def __iter__(self):
        return iter(range(10))

    def __length_hint__(self):
        return NotImplemented


class TestLengthHintExceptions(unittest.TestCase):

    def test_issue1242657(self):
        self.assertRaises(RuntimeError, list, BadLen())
        self.assertRaises(RuntimeError, list, BadLengthHint())
        self.assertRaises(RuntimeError, [].extend, BadLen())
        self.assertRaises(RuntimeError, [].extend, BadLengthHint())
        b = bytearray(range(10))
        self.assertRaises(RuntimeError, b.extend, BadLen())
        self.assertRaises(RuntimeError, b.extend, BadLengthHint())

    def test_invalid_hint(self):
        # Make sure an invalid result doesn't muck-up the works
        self.assertEqual(list(NoneLengthHint()), list(range(10)))


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_iterlen.py
DragonHPC/dragon
import os
import json
import doctest
import unittest

from test import support

# import json with and without accelerations
cjson = support.import_fresh_module('json', fresh=['_json'])
pyjson = support.import_fresh_module('json', blocked=['_json'])
# JSONDecodeError is cached inside the _json module
cjson.JSONDecodeError = cjson.decoder.JSONDecodeError = json.JSONDecodeError

# create two base classes that will be used by the other tests
class PyTest(unittest.TestCase):
    json = pyjson
    loads = staticmethod(pyjson.loads)
    dumps = staticmethod(pyjson.dumps)
    JSONDecodeError = staticmethod(pyjson.JSONDecodeError)

@unittest.skipUnless(cjson, 'requires _json')
class CTest(unittest.TestCase):
    if cjson is not None:
        json = cjson
        loads = staticmethod(cjson.loads)
        dumps = staticmethod(cjson.dumps)
        JSONDecodeError = staticmethod(cjson.JSONDecodeError)

# test PyTest and CTest checking if the functions come from the right module
class TestPyTest(PyTest):
    def test_pyjson(self):
        self.assertEqual(self.json.scanner.make_scanner.__module__,
                         'json.scanner')
        self.assertEqual(self.json.decoder.scanstring.__module__,
                         'json.decoder')
        self.assertEqual(self.json.encoder.encode_basestring_ascii.__module__,
                         'json.encoder')

class TestCTest(CTest):
    def test_cjson(self):
        self.assertEqual(self.json.scanner.make_scanner.__module__, '_json')
        self.assertEqual(self.json.decoder.scanstring.__module__, '_json')
        self.assertEqual(self.json.encoder.c_make_encoder.__module__, '_json')
        self.assertEqual(self.json.encoder.encode_basestring_ascii.__module__,
                         '_json')


def load_tests(loader, _, pattern):
    suite = unittest.TestSuite()
    for mod in (json, json.encoder, json.decoder):
        suite.addTest(doctest.DocTestSuite(mod))
    suite.addTest(TestPyTest('test_pyjson'))
    suite.addTest(TestCTest('test_cjson'))

    pkg_dir = os.path.dirname(__file__)
    return support.load_package_tests(pkg_dir, loader, suite, pattern)

external/cpython_master/Lib/test/test_json/__init__.py
DragonHPC/dragon
from collections import OrderedDict
from test.test_json import PyTest, CTest
from test.support import bigaddrspacetest


CASES = [
    ('/\\"\ucafe\ubabe\uab98\ufcde\ubcda\uef4a\x08\x0c\n\r\t`1~!@#$%^&*()_+-=[]{}|;:\',./<>?', '"/\\\\\\"\\ucafe\\ubabe\\uab98\\ufcde\\ubcda\\uef4a\\b\\f\\n\\r\\t`1~!@#$%^&*()_+-=[]{}|;:\',./<>?"'),
    ('\u0123\u4567\u89ab\ucdef\uabcd\uef4a', '"\\u0123\\u4567\\u89ab\\ucdef\\uabcd\\uef4a"'),
    ('controls', '"controls"'),
    ('\x08\x0c\n\r\t', '"\\b\\f\\n\\r\\t"'),
    ('{"object with 1 member":["array with 1 element"]}', '"{\\"object with 1 member\\":[\\"array with 1 element\\"]}"'),
    (' s p a c e d ', '" s p a c e d "'),
    ('\U0001d120', '"\\ud834\\udd20"'),
    ('\u03b1\u03a9', '"\\u03b1\\u03a9"'),
    ("`1~!@#$%^&*()_+-={':[,]}|;.</>?", '"`1~!@#$%^&*()_+-={\':[,]}|;.</>?"'),
    ('\x08\x0c\n\r\t', '"\\b\\f\\n\\r\\t"'),
    ('\u0123\u4567\u89ab\ucdef\uabcd\uef4a', '"\\u0123\\u4567\\u89ab\\ucdef\\uabcd\\uef4a"'),
]

class TestEncodeBasestringAscii:
    def test_encode_basestring_ascii(self):
        fname = self.json.encoder.encode_basestring_ascii.__name__
        for input_string, expect in CASES:
            result = self.json.encoder.encode_basestring_ascii(input_string)
            self.assertEqual(result, expect,
                '{0!r} != {1!r} for {2}({3!r})'.format(
                    result, expect, fname, input_string))

    def test_ordered_dict(self):
        # See issue 6105
        items = [('one', 1), ('two', 2), ('three', 3), ('four', 4), ('five', 5)]
        s = self.dumps(OrderedDict(items))
        self.assertEqual(s, '{"one": 1, "two": 2, "three": 3, "four": 4, "five": 5}')

    def test_sorted_dict(self):
        items = [('one', 1), ('two', 2), ('three', 3), ('four', 4), ('five', 5)]
        s = self.dumps(dict(items), sort_keys=True)
        self.assertEqual(s, '{"five": 5, "four": 4, "one": 1, "three": 3, "two": 2}')


class TestPyEncodeBasestringAscii(TestEncodeBasestringAscii, PyTest): pass
class TestCEncodeBasestringAscii(TestEncodeBasestringAscii, CTest):
    @bigaddrspacetest
    def test_overflow(self):
        size = (2**32)//6 + 1
        s = "\x00"*size
        with self.assertRaises(OverflowError):
            self.json.encoder.encode_basestring_ascii(s)

external/cpython_master/Lib/test/test_json/test_encode_basestring_ascii.py
DragonHPC/dragon
import sys
from test.test_json import PyTest, CTest


class TestScanstring:
    def test_scanstring(self):
        scanstring = self.json.decoder.scanstring
        self.assertEqual(
            scanstring('"z\U0001d120x"', 1, True),
            ('z\U0001d120x', 5))

        self.assertEqual(
            scanstring('"\\u007b"', 1, True),
            ('{', 8))

        self.assertEqual(
            scanstring('"A JSON payload should be an object or array, not a string."', 1, True),
            ('A JSON payload should be an object or array, not a string.', 60))

        self.assertEqual(
            scanstring('["Unclosed array"', 2, True),
            ('Unclosed array', 17))

        self.assertEqual(
            scanstring('["extra comma",]', 2, True),
            ('extra comma', 14))

        self.assertEqual(
            scanstring('["double extra comma",,]', 2, True),
            ('double extra comma', 21))

        self.assertEqual(
            scanstring('["Comma after the close"],', 2, True),
            ('Comma after the close', 24))

        self.assertEqual(
            scanstring('["Extra close"]]', 2, True),
            ('Extra close', 14))

        self.assertEqual(
            scanstring('{"Extra comma": true,}', 2, True),
            ('Extra comma', 14))

        self.assertEqual(
            scanstring('{"Extra value after close": true} "misplaced quoted value"', 2, True),
            ('Extra value after close', 26))

        self.assertEqual(
            scanstring('{"Illegal expression": 1 + 2}', 2, True),
            ('Illegal expression', 21))

        self.assertEqual(
            scanstring('{"Illegal invocation": alert()}', 2, True),
            ('Illegal invocation', 21))

        self.assertEqual(
            scanstring('{"Numbers cannot have leading zeroes": 013}', 2, True),
            ('Numbers cannot have leading zeroes', 37))

        self.assertEqual(
            scanstring('{"Numbers cannot be hex": 0x14}', 2, True),
            ('Numbers cannot be hex', 24))

        self.assertEqual(
            scanstring('[[[[[[[[[[[[[[[[[[[["Too deep"]]]]]]]]]]]]]]]]]]]]', 21, True),
            ('Too deep', 30))

        self.assertEqual(
            scanstring('{"Missing colon" null}', 2, True),
            ('Missing colon', 16))

        self.assertEqual(
            scanstring('{"Double colon":: null}', 2, True),
            ('Double colon', 15))

        self.assertEqual(
            scanstring('{"Comma instead of colon", null}', 2, True),
            ('Comma instead of colon', 25))

        self.assertEqual(
            scanstring('["Colon instead of comma": false]', 2, True),
            ('Colon instead of comma', 25))

        self.assertEqual(
            scanstring('["Bad value", truth]', 2, True),
            ('Bad value', 12))

    def test_surrogates(self):
        scanstring = self.json.decoder.scanstring
        def assertScan(given, expect):
            self.assertEqual(scanstring(given, 1, True),
                             (expect, len(given)))

        assertScan('"z\\ud834\\u0079x"', 'z\ud834yx')
        assertScan('"z\\ud834\\udd20x"', 'z\U0001d120x')
        assertScan('"z\\ud834\\ud834\\udd20x"', 'z\ud834\U0001d120x')
        assertScan('"z\\ud834x"', 'z\ud834x')
        assertScan('"z\\ud834\udd20x12345"', 'z\ud834\udd20x12345')
        assertScan('"z\\udd20x"', 'z\udd20x')
        assertScan('"z\ud834\udd20x"', 'z\ud834\udd20x')
        assertScan('"z\ud834\\udd20x"', 'z\ud834\udd20x')
        assertScan('"z\ud834x"', 'z\ud834x')

    def test_bad_escapes(self):
        scanstring = self.json.decoder.scanstring
        bad_escapes = [
            '"\\"',
            '"\\x"',
            '"\\u"',
            '"\\u0"',
            '"\\u01"',
            '"\\u012"',
            '"\\uz012"',
            '"\\u0z12"',
            '"\\u01z2"',
            '"\\u012z"',
            '"\\u0x12"',
            '"\\u0X12"',
            '"\\ud834\\"',
            '"\\ud834\\u"',
            '"\\ud834\\ud"',
            '"\\ud834\\udd"',
            '"\\ud834\\udd2"',
            '"\\ud834\\uzdd2"',
            '"\\ud834\\udzd2"',
            '"\\ud834\\uddz2"',
            '"\\ud834\\udd2z"',
            '"\\ud834\\u0x20"',
            '"\\ud834\\u0X20"',
        ]
        for s in bad_escapes:
            with self.assertRaises(self.JSONDecodeError, msg=s):
                scanstring(s, 1, True)

    def test_overflow(self):
        with self.assertRaises(OverflowError):
            self.json.decoder.scanstring(b"xxx", sys.maxsize+1)


class TestPyScanstring(TestScanstring, PyTest): pass
class TestCScanstring(TestScanstring, CTest): pass

external/cpython_master/Lib/test/test_json/test_scanstring.py
DragonHPC/dragon
import unittest
import test._test_multiprocessing

import sys
from test import support

if support.PGO:
    raise unittest.SkipTest("test is not helpful for PGO")

if sys.platform == "win32":
    raise unittest.SkipTest("fork is not available on Windows")

if sys.platform == 'darwin':
    raise unittest.SkipTest("test may crash on macOS (bpo-33725)")

test._test_multiprocessing.install_tests_in_module_dict(globals(), 'fork')

if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_multiprocessing_fork.py
DragonHPC/dragon
# test interactions between int, float, Decimal and Fraction

import unittest
import random
import math
import sys
import operator

from decimal import Decimal as D
from fractions import Fraction as F

# Constants related to the hash implementation;  hash(x) is based
# on the reduction of x modulo the prime _PyHASH_MODULUS.
_PyHASH_MODULUS = sys.hash_info.modulus
_PyHASH_INF = sys.hash_info.inf

class HashTest(unittest.TestCase):
    def check_equal_hash(self, x, y):
        # check both that x and y are equal and that their hashes are equal
        self.assertEqual(hash(x), hash(y),
                         "got different hashes for {!r} and {!r}".format(x, y))
        self.assertEqual(x, y)

    def test_bools(self):
        self.check_equal_hash(False, 0)
        self.check_equal_hash(True, 1)

    def test_integers(self):
        # check that equal values hash equal

        # exact integers
        for i in range(-1000, 1000):
            self.check_equal_hash(i, float(i))
            self.check_equal_hash(i, D(i))
            self.check_equal_hash(i, F(i))

        # the current hash is based on reduction modulo 2**n-1 for some
        # n, so pay special attention to numbers of the form 2**n and 2**n-1.
        for i in range(100):
            n = 2**i - 1
            if n == int(float(n)):
                self.check_equal_hash(n, float(n))
                self.check_equal_hash(-n, -float(n))
            self.check_equal_hash(n, D(n))
            self.check_equal_hash(n, F(n))
            self.check_equal_hash(-n, D(-n))
            self.check_equal_hash(-n, F(-n))

            n = 2**i
            self.check_equal_hash(n, float(n))
            self.check_equal_hash(-n, -float(n))
            self.check_equal_hash(n, D(n))
            self.check_equal_hash(n, F(n))
            self.check_equal_hash(-n, D(-n))
            self.check_equal_hash(-n, F(-n))

        # random values of various sizes
        for _ in range(1000):
            e = random.randrange(300)
            n = random.randrange(-10**e, 10**e)
            self.check_equal_hash(n, D(n))
            self.check_equal_hash(n, F(n))
            if n == int(float(n)):
                self.check_equal_hash(n, float(n))

    def test_binary_floats(self):
        # check that floats hash equal to corresponding Fractions and Decimals

        # floats that are distinct but numerically equal should hash the same
        self.check_equal_hash(0.0, -0.0)

        # zeros
        self.check_equal_hash(0.0, D(0))
        self.check_equal_hash(-0.0, D(0))
        self.check_equal_hash(-0.0, D('-0.0'))
        self.check_equal_hash(0.0, F(0))

        # infinities and nans
        self.check_equal_hash(float('inf'), D('inf'))
        self.check_equal_hash(float('-inf'), D('-inf'))

        for _ in range(1000):
            x = random.random() * math.exp(random.random()*200.0 - 100.0)
            self.check_equal_hash(x, D.from_float(x))
            self.check_equal_hash(x, F.from_float(x))

    def test_complex(self):
        # complex numbers with zero imaginary part should hash equal to
        # the corresponding float

        test_values = [0.0, -0.0, 1.0, -1.0, 0.40625, -5136.5,
                       float('inf'), float('-inf')]

        for zero in -0.0, 0.0:
            for value in test_values:
                self.check_equal_hash(value, complex(value, zero))

    def test_decimals(self):
        # check that Decimal instances that have different representations
        # but equal values give the same hash
        zeros = ['0', '-0', '0.0', '-0.0e10', '000e-10']
        for zero in zeros:
            self.check_equal_hash(D(zero), D(0))

        self.check_equal_hash(D('1.00'), D(1))
        self.check_equal_hash(D('1.00000'), D(1))
        self.check_equal_hash(D('-1.00'), D(-1))
        self.check_equal_hash(D('-1.00000'), D(-1))
        self.check_equal_hash(D('123e2'), D(12300))
        self.check_equal_hash(D('1230e1'), D(12300))
        self.check_equal_hash(D('12300'), D(12300))
        self.check_equal_hash(D('12300.0'), D(12300))
        self.check_equal_hash(D('12300.00'), D(12300))
        self.check_equal_hash(D('12300.000'), D(12300))

    def test_fractions(self):
        # check special case for fractions where either the numerator
        # or the denominator is a multiple of _PyHASH_MODULUS
        self.assertEqual(hash(F(1, _PyHASH_MODULUS)), _PyHASH_INF)
        self.assertEqual(hash(F(-1, 3*_PyHASH_MODULUS)), -_PyHASH_INF)
        self.assertEqual(hash(F(7*_PyHASH_MODULUS, 1)), 0)
        self.assertEqual(hash(F(-_PyHASH_MODULUS, 1)), 0)

    def test_hash_normalization(self):
        # Test for a bug encountered while changing long_hash.
        #
        # Given objects x and y, it should be possible for y's
        # __hash__ method to return hash(x) in order to ensure that
        # hash(x) == hash(y).  But hash(x) is not exactly equal to the
        # result of x.__hash__(): there's some internal normalization
        # to make sure that the result fits in a C long, and is not
        # equal to the invalid hash value -1.  This internal
        # normalization must therefore not change the result of
        # hash(x) for any x.

        class HalibutProxy:
            def __hash__(self):
                return hash('halibut')
            def __eq__(self, other):
                return other == 'halibut'

        x = {'halibut', HalibutProxy()}
        self.assertEqual(len(x), 1)

class ComparisonTest(unittest.TestCase):
    def test_mixed_comparisons(self):

        # ordered list of distinct test values of various types:
        # int, float, Fraction, Decimal
        test_values = [
            float('-inf'),
            D('-1e425000000'),
            -1e308,
            F(-22, 7),
            -3.14,
            -2,
            0.0,
            1e-320,
            True,
            F('1.2'),
            D('1.3'),
            float('1.4'),
            F(275807, 195025),
            D('1.414213562373095048801688724'),
            F(114243, 80782),
            F(473596569, 84615),
            7e200,
            D('infinity'),
            ]
        for i, first in enumerate(test_values):
            for second in test_values[i+1:]:
                self.assertLess(first, second)
                self.assertLessEqual(first, second)
                self.assertGreater(second, first)
                self.assertGreaterEqual(second, first)

    def test_complex(self):
        # comparisons with complex are special:  equality and inequality
        # comparisons should always succeed, but order comparisons should
        # raise TypeError.
        z = 1.0 + 0j
        w = -3.14 + 2.7j

        for v in 1, 1.0, F(1), D(1), complex(1):
            self.assertEqual(z, v)
            self.assertEqual(v, z)

        for v in 2, 2.0, F(2), D(2), complex(2):
            self.assertNotEqual(z, v)
            self.assertNotEqual(v, z)
            self.assertNotEqual(w, v)
            self.assertNotEqual(v, w)

        for v in (1, 1.0, F(1), D(1), complex(1),
                  2, 2.0, F(2), D(2), complex(2), w):
            for op in operator.le, operator.lt, operator.ge, operator.gt:
                self.assertRaises(TypeError, op, z, v)
                self.assertRaises(TypeError, op, v, z)


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_numeric_tower.py
DragonHPC/dragon
"""
Test suite for OS X interpreter environment variables.
"""

from test.support import EnvironmentVarGuard
import subprocess
import sys
import sysconfig
import unittest

@unittest.skipUnless(sys.platform == 'darwin' and
                     sysconfig.get_config_var('WITH_NEXT_FRAMEWORK'),
                     'unnecessary on this platform')
class OSXEnvironmentVariableTestCase(unittest.TestCase):
    def _check_sys(self, ev, cond, sv, val = sys.executable + 'dummy'):
        with EnvironmentVarGuard() as evg:
            subpc = [str(sys.executable), '-c',
                'import sys; sys.exit(2 if "%s" %s %s else 3)' % (val, cond, sv)]
            # ensure environment variable does not exist
            evg.unset(ev)
            # test that test on sys.xxx normally fails
            rc = subprocess.call(subpc)
            self.assertEqual(rc, 3, "expected %s not %s %s" % (ev, cond, sv))
            # set environ variable
            evg.set(ev, val)
            # test that sys.xxx has been influenced by the environ value
            rc = subprocess.call(subpc)
            self.assertEqual(rc, 2, "expected %s %s %s" % (ev, cond, sv))

    def test_pythonexecutable_sets_sys_executable(self):
        self._check_sys('PYTHONEXECUTABLE', '==', 'sys.executable')

if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_osx_env.py
DragonHPC/dragon
# Tests for the correctly-rounded string -> float conversions
# introduced in Python 2.7 and 3.1.

import random
import unittest
import re
import sys
import test.support

if getattr(sys, 'float_repr_style', '') != 'short':
    raise unittest.SkipTest('correctly-rounded string->float conversions '
                            'not available on this system')

# Correctly rounded str -> float in pure Python, for comparison.

strtod_parser = re.compile(r"""    # A numeric string consists of:
    (?P<sign>[-+])?          # an optional sign, followed by
    (?=\d|\.\d)              # a number with at least one digit
    (?P<int>\d*)             # having a (possibly empty) integer part
    (?:\.(?P<frac>\d*))?     # followed by an optional fractional part
    (?:E(?P<exp>[-+]?\d+))?  # and an optional exponent
    \Z
""", re.VERBOSE | re.IGNORECASE).match

# Pure Python version of correctly rounded string->float conversion.
# Avoids any use of floating-point by returning the result as a hex string.
def strtod(s, mant_dig=53, min_exp = -1021, max_exp = 1024):
    """Convert a finite decimal string to a hex string representing an
    IEEE 754 binary64 float.  Return 'inf' or '-inf' on overflow.
    This function makes no use of floating-point arithmetic at any
    stage."""

    # parse string into a pair of integers 'a' and 'b' such that
    # abs(decimal value) = a/b, along with a boolean 'negative'.
    m = strtod_parser(s)
    if m is None:
        raise ValueError('invalid numeric string')
    fraction = m.group('frac') or ''
    intpart = int(m.group('int') + fraction)
    exp = int(m.group('exp') or '0') - len(fraction)
    negative = m.group('sign') == '-'
    a, b = intpart*10**max(exp, 0), 10**max(0, -exp)

    # quick return for zeros
    if not a:
        return '-0x0.0p+0' if negative else '0x0.0p+0'

    # compute exponent e for result; may be one too small in the case
    # that the rounded value of a/b lies in a different binade from a/b
    d = a.bit_length() - b.bit_length()
    d += (a >> d if d >= 0 else a << -d) >= b
    e = max(d, min_exp) - mant_dig

    # approximate a/b by number of the form q * 2**e; adjust e if necessary
    a, b = a << max(-e, 0), b << max(e, 0)
    q, r = divmod(a, b)
    if 2*r > b or 2*r == b and q & 1:
        q += 1
        if q.bit_length() == mant_dig+1:
            q //= 2
            e += 1

    # double check that (q, e) has the right form
    assert q.bit_length() <= mant_dig and e >= min_exp - mant_dig
    assert q.bit_length() == mant_dig or e == min_exp - mant_dig

    # check for overflow and underflow
    if e + q.bit_length() > max_exp:
        return '-inf' if negative else 'inf'
    if not q:
        return '-0x0.0p+0' if negative else '0x0.0p+0'

    # for hex representation, shift so # bits after point is a multiple of 4
    hexdigs = 1 + (mant_dig-2)//4
    shift = 3 - (mant_dig-2)%4
    q, e = q << shift, e - shift
    return '{}0x{:x}.{:0{}x}p{:+d}'.format(
        '-' if negative else '',
        q // 16**hexdigs,
        q % 16**hexdigs,
        hexdigs,
        e + 4*hexdigs)

TEST_SIZE = 10

class StrtodTests(unittest.TestCase):
    def check_strtod(self, s):
        """Compare the result of Python's builtin correctly rounded
        string->float conversion (using float) to a pure Python
        correctly rounded string->float implementation.  Fail if the
        two methods give different results."""

        try:
            fs = float(s)
        except OverflowError:
            got = '-inf' if s[0] == '-' else 'inf'
        except MemoryError:
            got = 'memory error'
        else:
            got = fs.hex()
        expected = strtod(s)
        self.assertEqual(expected, got,
                         "Incorrectly rounded str->float conversion for {}: "
                         "expected {}, got {}".format(s, expected, got))

    def test_short_halfway_cases(self):
        # exact halfway cases with a small number of significant digits
        for k in 0, 5, 10, 15, 20:
            # upper = smallest integer >= 2**54/5**k
            upper = -(-2**54//5**k)
            # lower = smallest odd number >= 2**53/5**k
            lower = -(-2**53//5**k)
            if lower % 2 == 0:
                lower += 1
            for i in range(TEST_SIZE):
                # Select a random odd n in [2**53/5**k,
                # 2**54/5**k). Then n * 10**k gives a halfway case
                # with small number of significant digits.
                n, e = random.randrange(lower, upper, 2), k

                # Remove any additional powers of 5.
                while n % 5 == 0:
                    n, e = n // 5, e + 1
                assert n % 10 in (1, 3, 7, 9)

                # Try numbers of the form n * 2**p2 * 10**e, p2 >= 0,
                # until n * 2**p2 has more than 20 significant digits.
                digits, exponent = n, e
                while digits < 10**20:
                    s = '{}e{}'.format(digits, exponent)
                    self.check_strtod(s)
                    # Same again, but with extra trailing zeros.
                    s = '{}e{}'.format(digits * 10**40, exponent - 40)
                    self.check_strtod(s)
                    digits *= 2

                # Try numbers of the form n * 5**p2 * 10**(e - p5), p5
                # >= 0, with n * 5**p5 < 10**20.
                digits, exponent = n, e
                while digits < 10**20:
                    s = '{}e{}'.format(digits, exponent)
                    self.check_strtod(s)
                    # Same again, but with extra trailing zeros.
                    s = '{}e{}'.format(digits * 10**40, exponent - 40)
                    self.check_strtod(s)
                    digits *= 5
                    exponent -= 1

    def test_halfway_cases(self):
        # test halfway cases for the round-half-to-even rule
        for i in range(100 * TEST_SIZE):
            # bit pattern for a random finite positive (or +0.0) float
            bits = random.randrange(2047*2**52)

            # convert bit pattern to a number of the form m * 2**e
            e, m = divmod(bits, 2**52)
            if e:
                m, e = m + 2**52, e - 1
            e -= 1074

            # add 0.5 ulps
            m, e = 2*m + 1, e - 1

            # convert to a decimal string
            if e >= 0:
                digits = m << e
                exponent = 0
            else:
                # m * 2**e = (m * 5**-e) * 10**e
                digits = m * 5**-e
                exponent = e
            s = '{}e{}'.format(digits, exponent)
            self.check_strtod(s)

    def test_boundaries(self):
        # boundaries expressed as triples (n, e, u), where
        # n*10**e is an approximation to the boundary value and
        # u*10**e is 1ulp
        boundaries = [
            (10000000000000000000, -19, 1110),   # a power of 2 boundary (1.0)
            (17976931348623159077, 289, 1995),   # overflow boundary (2.**1024)
            (22250738585072013831, -327, 4941),  # normal/subnormal (2.**-1022)
            (0, -327, 4941),                     # zero
            ]
        for n, e, u in boundaries:
            for j in range(1000):
                digits = n + random.randrange(-3*u, 3*u)
                exponent = e
                s = '{}e{}'.format(digits, exponent)
                self.check_strtod(s)
                n *= 10
                u *= 10
                e -= 1

    def test_underflow_boundary(self):
        # test values close to 2**-1075, the underflow boundary; similar
        # to boundary_tests, except that the random error doesn't scale
        # with n
        for exponent in range(-400, -320):
            base = 10**-exponent // 2**1075
            for j in range(TEST_SIZE):
                digits = base + random.randrange(-1000, 1000)
                s = '{}e{}'.format(digits, exponent)
                self.check_strtod(s)

    def test_bigcomp(self):
        for ndigs in 5, 10, 14, 15, 16, 17, 18, 19, 20, 40, 41, 50:
            dig10 = 10**ndigs
            for i in range(10 * TEST_SIZE):
                digits = random.randrange(dig10)
                exponent = random.randrange(-400, 400)
                s = '{}e{}'.format(digits, exponent)
                self.check_strtod(s)

    def test_parsing(self):
        # make '0' more likely to be chosen than other digits
        digits = '000000123456789'
        signs = ('+', '-', '')

        # put together random short valid strings
        # \d*[.\d*]?e
        for i in range(1000):
            for j in range(TEST_SIZE):
                s = random.choice(signs)
                intpart_len = random.randrange(5)
                s += ''.join(random.choice(digits) for _ in range(intpart_len))
                if random.choice([True, False]):
                    s += '.'
                    fracpart_len = random.randrange(5)
                    s += ''.join(random.choice(digits)
                                 for _ in range(fracpart_len))
                else:
                    fracpart_len = 0
                if random.choice([True, False]):
                    s += random.choice(['e', 'E'])
                    s += random.choice(signs)
                    exponent_len = random.randrange(1, 4)
                    s += ''.join(random.choice(digits)
                                 for _ in range(exponent_len))

                if intpart_len + fracpart_len:
                    self.check_strtod(s)
                else:
                    try:
                        float(s)
                    except ValueError:
                        pass
                    else:
                        assert False, "expected ValueError"

    @test.support.bigmemtest(size=test.support._2G+10, memuse=3, dry_run=False)
    def test_oversized_digit_strings(self, maxsize):
        # Input string whose length doesn't fit in an INT.
        s = "1." + "1" * maxsize
        with self.assertRaises(ValueError):
            float(s)
        del s

        s = "0." + "0" * maxsize + "1"
        with self.assertRaises(ValueError):
            float(s)
        del s

    def test_large_exponents(self):
        # Verify that the clipping of the exponent in strtod doesn't affect the
        # output values.
        def positive_exp(n):
            """ Long string with value 1.0 and exponent n"""
            return '0.{}1e+{}'.format('0'*(n-1), n)

        def negative_exp(n):
            """ Long string with value 1.0 and exponent -n"""
            return '1{}e-{}'.format('0'*n, n)

        self.assertEqual(float(positive_exp(10000)), 1.0)
        self.assertEqual(float(positive_exp(20000)), 1.0)
        self.assertEqual(float(positive_exp(30000)), 1.0)
        self.assertEqual(float(negative_exp(10000)), 1.0)
        self.assertEqual(float(negative_exp(20000)), 1.0)
        self.assertEqual(float(negative_exp(30000)), 1.0)

    def test_particular(self):
        # inputs that produced crashes or incorrectly rounded results with
        # previous versions of dtoa.c, for various reasons
        test_strings = [
            # issue 7632 bug 1, originally reported failing case
            '2183167012312112312312.23538020374420446192e-370',
            # 5 instances of issue 7632 bug 2
            '12579816049008305546974391768996369464963024663104e-357',
            '17489628565202117263145367596028389348922981857013e-357',
            '18487398785991994634182916638542680759613590482273e-357',
            '32002864200581033134358724675198044527469366773928e-358',
            '94393431193180696942841837085033647913224148539854e-358',
            '73608278998966969345824653500136787876436005957953e-358',
            '64774478836417299491718435234611299336288082136054e-358',
            '13704940134126574534878641876947980878824688451169e-357',
            '46697445774047060960624497964425416610480524760471e-358',
            # failing case for bug introduced by METD in r77451 (attempted
            # fix for issue 7632, bug 2), and fixed in r77482.
            '28639097178261763178489759107321392745108491825303e-311',
            # two numbers demonstrating a flaw in the bigcomp 'dig == 0'
            # correction block (issue 7632, bug 3)
            '1.00000000000000001e44',
            '1.0000000000000000100000000000000000000001e44',
            # dtoa.c bug for numbers just smaller than a power of 2 (issue
            # 7632, bug 4)
            '99999999999999994487665465554760717039532578546e-47',
            # failing case for off-by-one error introduced by METD in
            # r77483 (dtoa.c cleanup), fixed in r77490
            '965437176333654931799035513671997118345570045914469' #...
            '6213413350821416312194420007991306908470147322020121018368e0',
            # incorrect lsb detection for round-half-to-even when
            # bc->scale != 0 (issue 7632, bug 6).
            '104308485241983990666713401708072175773165034278685' #...
            '682646111762292409330928739751702404658197872319129' #...
            '036519947435319418387839758990478549477777586673075' #...
            '945844895981012024387992135617064532141489278815239' #...
            '849108105951619997829153633535314849999674266169258' #...
            '928940692239684771590065027025835804863585454872499' #...
            '320500023126142553932654370362024104462255244034053' #...
            '203998964360882487378334860197725139151265590832887' #...
            '433736189468858614521708567646743455601905935595381' #...
            '852723723645799866672558576993978025033590728687206' #...
            '296379801363024094048327273913079612469982585674824' #...
            '156000783167963081616214710691759864332339239688734' #...
            '656548790656486646106983450809073750535624894296242' #...
            '072010195710276073042036425579852459556183541199012' #...
            '652571123898996574563824424330960027873516082763671875e-1075',
            # demonstration that original fix for issue 7632 bug 1 was
            # buggy; the exit condition was too strong
            '247032822920623295e-341',
            # demonstrate similar problem to issue 7632 bug1: crash
            # with 'oversized quotient in quorem' message.
            '99037485700245683102805043437346965248029601286431e-373',
            '99617639833743863161109961162881027406769510558457e-373',
            '98852915025769345295749278351563179840130565591462e-372',
            '99059944827693569659153042769690930905148015876788e-373',
            '98914979205069368270421829889078356254059760327101e-372',
            # issue 7632 bug 5: the following 2 strings convert differently
            '1000000000000000000000000000000000000000e-16',
            '10000000000000000000000000000000000000000e-17',
            # issue 7632 bug 7
            '991633793189150720000000000000000000000000000000000000000e-33',
            # And another, similar, failing halfway case
            '4106250198039490000000000000000000000000000000000000000e-38',
            # issue 7632 bug 8:  the following produced 10.0
            '10.900000000000000012345678912345678912345',

            # two humongous values from issue 7743
            '116512874940594195638617907092569881519034793229385' #...
            '228569165191541890846564669771714896916084883987920' #...
            '473321268100296857636200926065340769682863349205363' #...
            '349247637660671783209907949273683040397979984107806' #...
            '461822693332712828397617946036239581632976585100633' #...
            '520260770761060725403904123144384571612073732754774' #...
            '588211944406465572591022081973828448927338602556287' #...
            '851831745419397433012491884869454462440536895047499' #...
            '436551974649731917170099387762871020403582994193439' #...
            '761933412166821484015883631622539314203799034497982' #...
            '130038741741727907429575673302461380386596501187482' #...
            '006257527709842179336488381672818798450229339123527' #...
            '858844448336815912020452294624916993546388956561522' #...
            '161875352572590420823607478788399460162228308693742' #...
            '05287663441403533948204085390898399055004119873046875e-1075',

            '525440653352955266109661060358202819561258984964913' #...
            '892256527849758956045218257059713765874251436193619' #...
            '443248205998870001633865657517447355992225852945912' #...
            '016668660000210283807209850662224417504752264995360' #...
            '631512007753855801075373057632157738752800840302596' #...
            '237050247910530538250008682272783660778181628040733' #...
            '653121492436408812668023478001208529190359254322340' #...
            '397575185248844788515410722958784640926528544043090' #...
            '115352513640884988017342469275006999104519620946430' #...
            '818767147966495485406577703972687838176778993472989' #...
            '561959000047036638938396333146685137903018376496408' #...
            '319705333868476925297317136513970189073693314710318' #...
            '991252811050501448326875232850600451776091303043715' #...
            '157191292827614046876950225714743118291034780466325' #...
            '085141343734564915193426994587206432697337118211527' #...
            '278968731294639353354774788602467795167875117481660' #...
            '4738791256853675690543663283782215866825e-1180',

            # exercise exit conditions in bigcomp comparison loop
            '2602129298404963083833853479113577253105939995688e2',
            '260212929840496308383385347911357725310593999568896e0',
            '26021292984049630838338534791135772531059399956889601e-2',
            '260212929840496308383385347911357725310593999568895e0',
            '260212929840496308383385347911357725310593999568897e0',
            '260212929840496308383385347911357725310593999568996e0',
            '260212929840496308383385347911357725310593999568866e0',
            # 2**53
            '9007199254740992.00',
            # 2**1024 - 2**970:  exact overflow boundary.  All values
            # smaller than this should round to something finite;  any value
            # greater than or equal to this one overflows.
            '179769313486231580793728971405303415079934132710037' #...
            '826936173778980444968292764750946649017977587207096' #...
            '330286416692887910946555547851940402630657488671505' #...
            '820681908902000708383676273854845817711531764475730' #...
            '270069855571366959622842914819860834936475292719074' #...
            '168444365510704342711559699508093042880177904174497792',
            # 2**1024 - 2**970 - tiny
            '179769313486231580793728971405303415079934132710037' #...
            '826936173778980444968292764750946649017977587207096' #...
            '330286416692887910946555547851940402630657488671505' #...
            '820681908902000708383676273854845817711531764475730' #...
            '270069855571366959622842914819860834936475292719074' #...
            '168444365510704342711559699508093042880177904174497791.999',
            # 2**1024 - 2**970 + tiny
            '179769313486231580793728971405303415079934132710037' #...
            '826936173778980444968292764750946649017977587207096' #...
            '330286416692887910946555547851940402630657488671505' #...
            '820681908902000708383676273854845817711531764475730' #...
            '270069855571366959622842914819860834936475292719074' #...
            '168444365510704342711559699508093042880177904174497792.001',
            # 1 - 2**-54, +-tiny
            '999999999999999944488848768742172978818416595458984375e-54',
            '9999999999999999444888487687421729788184165954589843749999999e-54',
            '9999999999999999444888487687421729788184165954589843750000001e-54',
            # Value found by Rick Regan that gives a result of 2**-968
            # under Gay's dtoa.c (as of Nov 04, 2010);  since fixed.
            # (Fixed some time ago in Python's dtoa.c.)
            '0.0000000000000000000000000000000000000000100000000' #...
            '000000000576129113423785429971690421191214034235435' #...
            '087147763178149762956868991692289869941246658073194' #...
            '51982237978882039897143840789794921875',
            ]
        for s in test_strings:
            self.check_strtod(s)

if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_strtod.py
DragonHPC/dragon
import unittest
from test import support

# Skip this test if the _testcapi module isn't available.
support.import_module('_testcapi')
from _testcapi import _test_structmembersType, \
    CHAR_MAX, CHAR_MIN, UCHAR_MAX, \
    SHRT_MAX, SHRT_MIN, USHRT_MAX, \
    INT_MAX, INT_MIN, UINT_MAX, \
    LONG_MAX, LONG_MIN, ULONG_MAX, \
    LLONG_MAX, LLONG_MIN, ULLONG_MAX, \
    PY_SSIZE_T_MAX, PY_SSIZE_T_MIN

ts=_test_structmembersType(False,  # T_BOOL
                          1,      # T_BYTE
                          2,      # T_UBYTE
                          3,      # T_SHORT
                          4,      # T_USHORT
                          5,      # T_INT
                          6,      # T_UINT
                          7,      # T_LONG
                          8,      # T_ULONG
                          23,     # T_PYSSIZET
                          9.99999,# T_FLOAT
                          10.1010101010, # T_DOUBLE
                          "hi" # T_STRING_INPLACE
                          )

class ReadWriteTests(unittest.TestCase):

    def test_bool(self):
        ts.T_BOOL = True
        self.assertEqual(ts.T_BOOL, True)
        ts.T_BOOL = False
        self.assertEqual(ts.T_BOOL, False)
        self.assertRaises(TypeError, setattr, ts, 'T_BOOL', 1)

    def test_byte(self):
        ts.T_BYTE = CHAR_MAX
        self.assertEqual(ts.T_BYTE, CHAR_MAX)
        ts.T_BYTE = CHAR_MIN
        self.assertEqual(ts.T_BYTE, CHAR_MIN)
        ts.T_UBYTE = UCHAR_MAX
        self.assertEqual(ts.T_UBYTE, UCHAR_MAX)

    def test_short(self):
        ts.T_SHORT = SHRT_MAX
        self.assertEqual(ts.T_SHORT, SHRT_MAX)
        ts.T_SHORT = SHRT_MIN
        self.assertEqual(ts.T_SHORT, SHRT_MIN)
        ts.T_USHORT = USHRT_MAX
        self.assertEqual(ts.T_USHORT, USHRT_MAX)

    def test_int(self):
        ts.T_INT = INT_MAX
        self.assertEqual(ts.T_INT, INT_MAX)
        ts.T_INT = INT_MIN
        self.assertEqual(ts.T_INT, INT_MIN)
        ts.T_UINT = UINT_MAX
        self.assertEqual(ts.T_UINT, UINT_MAX)

    def test_long(self):
        ts.T_LONG = LONG_MAX
        self.assertEqual(ts.T_LONG, LONG_MAX)
        ts.T_LONG = LONG_MIN
        self.assertEqual(ts.T_LONG, LONG_MIN)
        ts.T_ULONG = ULONG_MAX
        self.assertEqual(ts.T_ULONG, ULONG_MAX)

    def test_py_ssize_t(self):
        ts.T_PYSSIZET = PY_SSIZE_T_MAX
        self.assertEqual(ts.T_PYSSIZET, PY_SSIZE_T_MAX)
        ts.T_PYSSIZET = PY_SSIZE_T_MIN
        self.assertEqual(ts.T_PYSSIZET, PY_SSIZE_T_MIN)

    @unittest.skipUnless(hasattr(ts, "T_LONGLONG"), "long long not present")
    def test_longlong(self):
        ts.T_LONGLONG = LLONG_MAX
        self.assertEqual(ts.T_LONGLONG, LLONG_MAX)
        ts.T_LONGLONG = LLONG_MIN
        self.assertEqual(ts.T_LONGLONG, LLONG_MIN)

        ts.T_ULONGLONG = ULLONG_MAX
        self.assertEqual(ts.T_ULONGLONG, ULLONG_MAX)

        ## make sure these will accept a plain int as well as a long
        ts.T_LONGLONG = 3
        self.assertEqual(ts.T_LONGLONG, 3)
        ts.T_ULONGLONG = 4
        self.assertEqual(ts.T_ULONGLONG, 4)

    def test_bad_assignments(self):
        integer_attributes = [
            'T_BOOL',
            'T_BYTE', 'T_UBYTE',
            'T_SHORT', 'T_USHORT',
            'T_INT', 'T_UINT',
            'T_LONG', 'T_ULONG',
            'T_PYSSIZET'
            ]
        if hasattr(ts, 'T_LONGLONG'):
            integer_attributes.extend(['T_LONGLONG', 'T_ULONGLONG'])

        # issue8014: this produced 'bad argument to internal function'
        # internal error
        for nonint in None, 3.2j, "full of eels", {}, []:
            for attr in integer_attributes:
                self.assertRaises(TypeError, setattr, ts, attr, nonint)

    def test_inplace_string(self):
        self.assertEqual(ts.T_STRING_INPLACE, "hi")
        self.assertRaises(TypeError, setattr, ts, "T_STRING_INPLACE", "s")
        self.assertRaises(TypeError, delattr, ts, "T_STRING_INPLACE")


class TestWarnings(unittest.TestCase):

    def test_byte_max(self):
        with support.check_warnings(('', RuntimeWarning)):
            ts.T_BYTE = CHAR_MAX+1

    def test_byte_min(self):
        with support.check_warnings(('', RuntimeWarning)):
            ts.T_BYTE = CHAR_MIN-1

    def test_ubyte_max(self):
        with support.check_warnings(('', RuntimeWarning)):
            ts.T_UBYTE = UCHAR_MAX+1

    def test_short_max(self):
        with support.check_warnings(('', RuntimeWarning)):
            ts.T_SHORT = SHRT_MAX+1

    def test_short_min(self):
        with support.check_warnings(('', RuntimeWarning)):
            ts.T_SHORT = SHRT_MIN-1

    def test_ushort_max(self):
        with support.check_warnings(('', RuntimeWarning)):
            ts.T_USHORT = USHRT_MAX+1


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_structmembers.py
DragonHPC/dragon
"""Tests for the gprof2html script in the Tools directory."""

import os
import sys
import unittest
from unittest import mock
import tempfile

from test.test_tools import skip_if_missing, import_tool

skip_if_missing()

class Gprof2htmlTests(unittest.TestCase):

    def setUp(self):
        self.gprof = import_tool('gprof2html')
        oldargv = sys.argv
        def fixup():
            sys.argv = oldargv
        self.addCleanup(fixup)
        sys.argv = []

    def test_gprof(self):
        # Issue #14508: this used to fail with a NameError.
        with mock.patch.object(self.gprof, 'webbrowser') as wmock, \
                tempfile.TemporaryDirectory() as tmpdir:
            fn = os.path.join(tmpdir, 'abc')
            open(fn, 'w').close()
            sys.argv = ['gprof2html', fn]
            self.gprof.main()
        self.assertTrue(wmock.open.called)


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_tools/test_gprof2html.py
DragonHPC/dragon
"""Tests for the pdeps script in the Tools directory."""

import os
import unittest
import tempfile

from test.test_tools import skip_if_missing, import_tool

skip_if_missing()


class PdepsTests(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        self.pdeps = import_tool('pdeps')

    def test_process_errors(self):
        # Issue #14492: m_import.match(line) can be None.
        with tempfile.TemporaryDirectory() as tmpdir:
            fn = os.path.join(tmpdir, 'foo')
            with open(fn, 'w') as stream:
                stream.write("#!/this/will/fail")
            self.pdeps.process(fn, {})

    def test_inverse_attribute_error(self):
        # Issue #14492: this used to fail with an AttributeError.
        self.pdeps.inverse({'a': []})


if __name__ == '__main__':
    unittest.main()

external/cpython_master/Lib/test/test_tools/test_pdeps.py
DragonHPC/dragon
from test import support

# Skip this test if _tkinter does not exist.
support.import_module('_tkinter')

from tkinter.test import runtktests

def test_main():
    support.run_unittest(
            *runtktests.get_tests(gui=False, packages=['test_ttk']))

if __name__ == '__main__':
    test_main()

external/cpython_master/Lib/test/test_ttk_textonly.py
DragonHPC/dragon
# Test the Unicode versions of normal file functions
# open, os.open, os.stat. os.listdir, os.rename, os.remove, os.mkdir, os.chdir, os.rmdir
import os
import sys
import unittest
import warnings
from unicodedata import normalize
from test import support

filenames = [
    '1_abc',
    '2_ascii',
    '3_Gr\xfc\xdf-Gott',
    '4_\u0393\u03b5\u03b9\u03ac-\u03c3\u03b1\u03c2',
    '5_\u0417\u0434\u0440\u0430\u0432\u0441\u0442\u0432\u0443\u0439\u0442\u0435',
    '6_\u306b\u307d\u3093',
    '7_\u05d4\u05e9\u05e7\u05e6\u05e5\u05e1',
    '8_\u66e8\u66e9\u66eb',
    '9_\u66e8\u05e9\u3093\u0434\u0393\xdf',
    # Specific code points: fn, NFC(fn) and NFKC(fn) all different
    '10_\u1fee\u1ffd',
    ]

# Mac OS X decomposes Unicode names, using Normal Form D.
# http://developer.apple.com/mac/library/qa/qa2001/qa1173.html
# "However, most volume formats do not follow the exact specification for
# these normal forms.  For example, HFS Plus uses a variant of Normal Form D
# in which U+2000 through U+2FFF, U+F900 through U+FAFF, and U+2F800 through
# U+2FAFF are not decomposed."
if sys.platform != 'darwin':
    filenames.extend([
        # Specific code points: NFC(fn), NFD(fn), NFKC(fn) and NFKD(fn) all different
        '11_\u0385\u03d3\u03d4',
        '12_\u00a8\u0301\u03d2\u0301\u03d2\u0308', # == NFD('\u0385\u03d3\u03d4')
        '13_\u0020\u0308\u0301\u038e\u03ab',       # == NFKC('\u0385\u03d3\u03d4')
        '14_\u1e9b\u1fc1\u1fcd\u1fce\u1fcf\u1fdd\u1fde\u1fdf\u1fed',

        # Specific code points: fn, NFC(fn) and NFKC(fn) all different
        '15_\u1fee\u1ffd\ufad1',
        '16_\u2000\u2000\u2000A',
        '17_\u2001\u2001\u2001A',
        '18_\u2003\u2003\u2003A',  # == NFC('\u2001\u2001\u2001A')
        '19_\u0020\u0020\u0020A',  # '\u0020' == ' ' == NFKC('\u2000') ==
                                   #  NFKC('\u2001') == NFKC('\u2003')
    ])


# Is it Unicode-friendly?
if not os.path.supports_unicode_filenames:
    fsencoding = sys.getfilesystemencoding()
    try:
        for name in filenames:
            name.encode(fsencoding)
    except UnicodeEncodeError:
        raise unittest.SkipTest("only NT+ and systems with "
                                "Unicode-friendly filesystem encoding")


class UnicodeFileTests(unittest.TestCase):
    files = set(filenames)
    normal_form = None

    def setUp(self):
        try:
            os.mkdir(support.TESTFN)
        except FileExistsError:
            pass
        self.addCleanup(support.rmtree, support.TESTFN)

        files = set()
        for name in self.files:
            name = os.path.join(support.TESTFN, self.norm(name))
            with open(name, 'wb') as f:
                f.write((name+'\n').encode("utf-8"))
            os.stat(name)
            files.add(name)
        self.files = files

    def norm(self, s):
        if self.normal_form:
            return normalize(self.normal_form, s)
        return s

    def _apply_failure(self, fn, filename,
                       expected_exception=FileNotFoundError,
                       check_filename=True):
        with self.assertRaises(expected_exception) as c:
            fn(filename)
        exc_filename = c.exception.filename
        if check_filename:
            self.assertEqual(exc_filename, filename, "Function '%s(%a) failed "
                             "with bad filename in the exception: %a" %
                             (fn.__name__, filename, exc_filename))

    def test_failures(self):
        # Pass non-existing Unicode filenames all over the place.
        for name in self.files:
            name = "not_" + name
            self._apply_failure(open, name)
            self._apply_failure(os.stat, name)
            self._apply_failure(os.chdir, name)
            self._apply_failure(os.rmdir, name)
            self._apply_failure(os.remove, name)
            self._apply_failure(os.listdir, name)

    if sys.platform == 'win32':
        # Windows is lunatic. Issue #13366.
        _listdir_failure = NotADirectoryError, FileNotFoundError
    else:
        _listdir_failure = NotADirectoryError

    def test_open(self):
        for name in self.files:
            f = open(name, 'wb')
            f.write((name+'\n').encode("utf-8"))
            f.close()
            os.stat(name)
            self._apply_failure(os.listdir, name, self._listdir_failure)

    # Skip the test on darwin, because darwin does normalize the filename to
    # NFD (a variant of Unicode NFD form). Normalize the filename to NFC, NFKC,
    # NFKD in Python is useless, because darwin will normalize it later and so
    # open(), os.stat(), etc. don't raise any exception.
    @unittest.skipIf(sys.platform == 'darwin', 'irrelevant test on Mac OS X')
    def test_normalize(self):
        files = set(self.files)
        others = set()
        for nf in set(['NFC', 'NFD', 'NFKC', 'NFKD']):
            others |= set(normalize(nf, file) for file in files)
        others -= files
        for name in others:
            self._apply_failure(open, name)
            self._apply_failure(os.stat, name)
            self._apply_failure(os.chdir, name)
            self._apply_failure(os.rmdir, name)
            self._apply_failure(os.remove, name)
            self._apply_failure(os.listdir, name)

    # Skip the test on darwin, because darwin uses a normalization different
    # than Python NFD normalization: filenames are different even if we use
    # Python NFD normalization.
    @unittest.skipIf(sys.platform == 'darwin', 'irrelevant test on Mac OS X')
    def test_listdir(self):
        sf0 = set(self.files)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)
            f1 = os.listdir(support.TESTFN.encode(sys.getfilesystemencoding()))
        f2 = os.listdir(support.TESTFN)
        sf2 = set(os.path.join(support.TESTFN, f) for f in f2)
        self.assertEqual(sf0, sf2, "%a != %a" % (sf0, sf2))
        self.assertEqual(len(f1), len(f2))

    def test_rename(self):
        for name in self.files:
            os.rename(name, "tmp")
            os.rename("tmp", name)

    def test_directory(self):
        dirname = os.path.join(support.TESTFN, 'Gr\xfc\xdf-\u66e8\u66e9\u66eb')
        filename = '\xdf-\u66e8\u66e9\u66eb'
        with support.temp_cwd(dirname):
            with open(filename, 'wb') as f:
                f.write((filename + '\n').encode("utf-8"))
            os.access(filename,os.R_OK)
            os.remove(filename)


class UnicodeNFCFileTests(UnicodeFileTests):
    normal_form = 'NFC'


class UnicodeNFDFileTests(UnicodeFileTests):
    normal_form = 'NFD'


class UnicodeNFKCFileTests(UnicodeFileTests):
    normal_form = 'NFKC'


class UnicodeNFKDFileTests(UnicodeFileTests):
    normal_form = 'NFKD'


def test_main():
    support.run_unittest(
        UnicodeFileTests,
        UnicodeNFCFileTests,
        UnicodeNFDFileTests,
        UnicodeNFKCFileTests,
        UnicodeNFKDFileTests,
    )


if __name__ == "__main__":
    test_main()

external/cpython_master/Lib/test/test_unicode_file_functions.py
DragonHPC/dragon
# Tests for extended unpacking, starred expressions.

doctests = """

Unpack tuple

    >>> t = (1, 2, 3)
    >>> a, *b, c = t
    >>> a == 1 and b == [2] and c == 3
    True

Unpack list

    >>> l = [4, 5, 6]
    >>> a, *b = l
    >>> a == 4 and b == [5, 6]
    True

Unpack implied tuple

    >>> *a, = 7, 8, 9
    >>> a == [7, 8, 9]
    True

Unpack string... fun!

    >>> a, *b = 'one'
    >>> a == 'o' and b == ['n', 'e']
    True

Unpack long sequence

    >>> a, b, c, *d, e, f, g = range(10)
    >>> (a, b, c, d, e, f, g) == (0, 1, 2, [3, 4, 5, 6], 7, 8, 9)
    True

Unpack short sequence

    >>> a, *b, c = (1, 2)
    >>> a == 1 and c == 2 and b == []
    True

Unpack generic sequence

    >>> class Seq:
    ...     def __getitem__(self, i):
    ...         if i >= 0 and i < 3: return i
    ...         raise IndexError
    ...
    >>> a, *b = Seq()
    >>> a == 0 and b == [1, 2]
    True

Unpack in for statement

    >>> for a, *b, c in [(1,2,3), (4,5,6,7)]:
    ...     print(a, b, c)
    ...
    1 [2] 3
    4 [5, 6] 7

Unpack in list

    >>> [a, *b, c] = range(5)
    >>> a == 0 and b == [1, 2, 3] and c == 4
    True

Multiple targets

    >>> a, *b, c = *d, e = range(5)
    >>> a == 0 and b == [1, 2, 3] and c == 4 and d == [0, 1, 2, 3] and e == 4
    True

Assignment unpacking

    >>> a, b, *c = range(5)
    >>> a, b, c
    (0, 1, [2, 3, 4])
    >>> *a, b, c = a, b, *c
    >>> a, b, c
    ([0, 1, 2], 3, 4)

Set display element unpacking

    >>> a = [1, 2, 3]
    >>> sorted({1, *a, 0, 4})
    [0, 1, 2, 3, 4]

    >>> {1, *1, 0, 4}
    Traceback (most recent call last):
      ...
    TypeError: 'int' object is not iterable

Dict display element unpacking

    >>> kwds = {'z': 0, 'w': 12}
    >>> sorted({'x': 1, 'y': 2, **kwds}.items())
    [('w', 12), ('x', 1), ('y', 2), ('z', 0)]

    >>> sorted({**{'x': 1}, 'y': 2, **{'z': 3}}.items())
    [('x', 1), ('y', 2), ('z', 3)]

    >>> sorted({**{'x': 1}, 'y': 2, **{'x': 3}}.items())
    [('x', 3), ('y', 2)]

    >>> sorted({**{'x': 1}, **{'x': 3}, 'x': 4}.items())
    [('x', 4)]

    >>> {**{}}
    {}

    >>> a = {}
    >>> {**a}[0] = 1
    >>> a
    {}

    >>> {**1}
    Traceback (most recent call last):
    ...
    TypeError: 'int' object is not a mapping

    >>> {**[]}
    Traceback (most recent call last):
    ...
    TypeError: 'list' object is not a mapping

    >>> len(eval("{" + ", ".join("**{{{}: {}}}".format(i, i)
    ...                          for i in range(1000)) + "}"))
    1000

    >>> {0:1, **{0:2}, 0:3, 0:4}
    {0: 4}

List comprehension element unpacking

    >>> a, b, c = [0, 1, 2], 3, 4
    >>> [*a, b, c]
    [0, 1, 2, 3, 4]

    >>> l = [a, (3, 4), {5}, {6: None}, (i for i in range(7, 10))]
    >>> [*item for item in l]
    Traceback (most recent call last):
    ...
    SyntaxError: iterable unpacking cannot be used in comprehension

    >>> [*[0, 1] for i in range(10)]
    Traceback (most recent call last):
    ...
    SyntaxError: iterable unpacking cannot be used in comprehension

    >>> [*'a' for i in range(10)]
    Traceback (most recent call last):
    ...
    SyntaxError: iterable unpacking cannot be used in comprehension

    >>> [*[] for i in range(10)]
    Traceback (most recent call last):
    ...
    SyntaxError: iterable unpacking cannot be used in comprehension

    >>> {**{} for a in [1]}
    Traceback (most recent call last):
    ...
    SyntaxError: dict unpacking cannot be used in dict comprehension

# Pegen is better here.
# Generator expression in function arguments

#     >>> list(*x for x in (range(5) for i in range(3)))
#     Traceback (most recent call last):
#     ...
#         list(*x for x in (range(5) for i in range(3)))
#                   ^
#     SyntaxError: invalid syntax

    >>> dict(**x for x in [{1:2}])
    Traceback (most recent call last):
    ...
        dict(**x for x in [{1:2}])
                   ^
    SyntaxError: invalid syntax

Iterable argument unpacking

    >>> print(*[1], *[2], 3)
    1 2 3

Make sure that they don't corrupt the passed-in dicts.

    >>> def f(x, y):
    ...     print(x, y)
    ...
    >>> original_dict = {'x': 1}
    >>> f(**original_dict, y=2)
    1 2
    >>> original_dict
    {'x': 1}

Now for some failures

Make sure the raised errors are right for keyword argument unpackings

    >>> from collections.abc import MutableMapping
    >>> class CrazyDict(MutableMapping):
    ...     def __init__(self):
    ...         self.d = {}
    ...
    ...     def __iter__(self):
    ...         for x in self.d.__iter__():
    ...             if x == 'c':
    ...                 self.d['z'] = 10
    ...             yield x
    ...
    ...     def __getitem__(self, k):
    ...         return self.d[k]
    ...
    ...     def __len__(self):
    ...         return len(self.d)
    ...
    ...     def __setitem__(self, k, v):
    ...         self.d[k] = v
    ...
    ...     def __delitem__(self, k):
    ...         del self.d[k]
    ...
    >>> d = CrazyDict()
    >>> d.d = {chr(ord('a') + x): x for x in range(5)}
    >>> e = {**d}
    Traceback (most recent call last):
    ...
    RuntimeError: dictionary changed size during iteration

    >>> d.d = {chr(ord('a') + x): x for x in range(5)}
    >>> def f(**kwargs): print(kwargs)
    >>> f(**d)
    Traceback (most recent call last):
    ...
    RuntimeError: dictionary changed size during iteration

Overridden parameters

    >>> f(x=5, **{'x': 3}, y=2)
    Traceback (most recent call last):
      ...
    TypeError: test.test_unpack_ex.f() got multiple values for keyword argument 'x'

    >>> f(**{'x': 3}, x=5, y=2)
    Traceback (most recent call last):
      ...
    TypeError: test.test_unpack_ex.f() got multiple values for keyword argument 'x'

    >>> f(**{'x': 3}, **{'x': 5}, y=2)
    Traceback (most recent call last):
      ...
    TypeError: test.test_unpack_ex.f() got multiple values for keyword argument 'x'

    >>> f(x=5, **{'x': 3}, **{'x': 2})
    Traceback (most recent call last):
      ...
    TypeError: test.test_unpack_ex.f() got multiple values for keyword argument 'x'

    >>> f(**{1: 3}, **{1: 5})
    Traceback (most recent call last):
      ...
    TypeError: test.test_unpack_ex.f() got multiple values for keyword argument '1'

Unpacking non-sequence

    >>> a, *b = 7
    Traceback (most recent call last):
      ...
    TypeError: cannot unpack non-iterable int object

Unpacking sequence too short

    >>> a, *b, c, d, e = Seq()
    Traceback (most recent call last):
      ...
    ValueError: not enough values to unpack (expected at least 4, got 3)

Unpacking sequence too short and target appears last

    >>> a, b, c, d, *e = Seq()
    Traceback (most recent call last):
      ...
    ValueError: not enough values to unpack (expected at least 4, got 3)

Unpacking a sequence where the test for too long raises a different kind of
error

    >>> class BozoError(Exception):
    ...     pass
    ...
    >>> class BadSeq:
    ...     def __getitem__(self, i):
    ...         if i >= 0 and i < 3:
    ...             return i
    ...         elif i == 3:
    ...             raise BozoError
    ...         else:
    ...             raise IndexError
    ...

Trigger code while not expecting an IndexError (unpack sequence too long, wrong
error)

    >>> a, *b, c, d, e = BadSeq()
    Traceback (most recent call last):
      ...
    test.test_unpack_ex.BozoError

Now some general starred expressions (all fail).

    >>> a, *b, c, *d, e = range(10) # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: multiple starred expressions in assignment

    >>> [*b, *c] = range(10) # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: multiple starred expressions in assignment

    >>> a,*b,*c,*d = range(4) # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: multiple starred expressions in assignment

    >>> *a = range(10) # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: starred assignment target must be in a list or tuple

    >>> *a # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

    >>> *1 # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

    >>> x = *a # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

    >>> (*x),y = 1, 2 # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

    >>> (((*x))),y = 1, 2 # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

    >>> z,(*x),y = 1, 2, 4 # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

    >>> z,(*x) = 1, 2 # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

    >>> ((*x),y) = 1, 2 # doctest:+ELLIPSIS
    Traceback (most recent call last):
      ...
    SyntaxError: can't use starred expression here

Some size constraints (all fail.)

    >>> s = ", ".join("a%d" % i for i in range(1<<8)) + ", *rest = range(1<<8 + 1)"
    >>> compile(s, 'test', 'exec') # doctest:+ELLIPSIS
    Traceback (most recent call last):
     ...
    SyntaxError: too many expressions in star-unpacking assignment

    >>> s = ", ".join("a%d" % i for i in range(1<<8 + 1)) + ", *rest = range(1<<8 + 2)"
    >>> compile(s, 'test', 'exec') # doctest:+ELLIPSIS
    Traceback (most recent call last):
     ...
    SyntaxError: too many expressions in star-unpacking assignment

(there is an additional limit, on the number of expressions after the
'*rest', but it's 1<<24 and testing it takes too much memory.)

"""

__test__ = {'doctests' : doctests}

def test_main(verbose=False):
    from test import support
    from test import test_unpack_ex
    support.run_doctest(test_unpack_ex, verbose)

if __name__ == "__main__":
    test_main(verbose=True)

external/cpython_master/Lib/test/test_unpack_ex.py
DragonHPC/dragon
# Ridiculously simple test of the winsound module for Windows.

import functools
import time
import unittest

from test import support

support.requires('audio')
winsound = support.import_module('winsound')


# Unless we actually have an ear in the room, we have no idea whether a sound
# actually plays, and it's incredibly flaky trying to figure out if a sound
# even *should* play.  Instead of guessing, just call the function and assume
# it either passed or raised the RuntimeError we expect in case of failure.
def sound_func(func):
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        try:
            ret = func(*args, **kwargs)
        except RuntimeError as e:
            if support.verbose:
                print(func.__name__, 'failed:', e)
        else:
            if support.verbose:
                print(func.__name__, 'returned')
            return ret
    return wrapper


safe_Beep = sound_func(winsound.Beep)
safe_MessageBeep = sound_func(winsound.MessageBeep)
safe_PlaySound = sound_func(winsound.PlaySound)


class BeepTest(unittest.TestCase):

    def test_errors(self):
        self.assertRaises(TypeError, winsound.Beep)
        self.assertRaises(ValueError, winsound.Beep, 36, 75)
        self.assertRaises(ValueError, winsound.Beep, 32768, 75)

    def test_extremes(self):
        safe_Beep(37, 75)
        safe_Beep(32767, 75)

    def test_increasingfrequency(self):
        for i in range(100, 2000, 100):
            safe_Beep(i, 75)

    def test_keyword_args(self):
        safe_Beep(duration=75, frequency=2000)


class MessageBeepTest(unittest.TestCase):

    def tearDown(self):
        time.sleep(0.5)

    def test_default(self):
        self.assertRaises(TypeError, winsound.MessageBeep, "bad")
        self.assertRaises(TypeError, winsound.MessageBeep, 42, 42)
        safe_MessageBeep()

    def test_ok(self):
        safe_MessageBeep(winsound.MB_OK)

    def test_asterisk(self):
        safe_MessageBeep(winsound.MB_ICONASTERISK)

    def test_exclamation(self):
        safe_MessageBeep(winsound.MB_ICONEXCLAMATION)

    def test_hand(self):
        safe_MessageBeep(winsound.MB_ICONHAND)

    def test_question(self):
        safe_MessageBeep(winsound.MB_ICONQUESTION)

    def test_keyword_args(self):
        safe_MessageBeep(type=winsound.MB_OK)


class PlaySoundTest(unittest.TestCase):

    def test_errors(self):
        self.assertRaises(TypeError, winsound.PlaySound)
        self.assertRaises(TypeError, winsound.PlaySound, "bad", "bad")
        self.assertRaises(
            RuntimeError,
            winsound.PlaySound,
            "none", winsound.SND_ASYNC | winsound.SND_MEMORY
        )
        self.assertRaises(TypeError, winsound.PlaySound, b"bad", 0)
        self.assertRaises(TypeError, winsound.PlaySound, "bad",
                          winsound.SND_MEMORY)
        self.assertRaises(TypeError, winsound.PlaySound, 1, 0)
        # embedded null character
        self.assertRaises(ValueError, winsound.PlaySound, 'bad\0', 0)

    def test_keyword_args(self):
        safe_PlaySound(flags=winsound.SND_ALIAS, sound="SystemExit")

    def test_snd_memory(self):
        with open(support.findfile('pluck-pcm8.wav',
                                   subdir='audiodata'), 'rb') as f:
            audio_data = f.read()
        safe_PlaySound(audio_data, winsound.SND_MEMORY)
        audio_data = bytearray(audio_data)
        safe_PlaySound(audio_data, winsound.SND_MEMORY)

    def test_snd_filename(self):
        fn = support.findfile('pluck-pcm8.wav', subdir='audiodata')
        safe_PlaySound(fn, winsound.SND_FILENAME | winsound.SND_NODEFAULT)

    def test_aliases(self):
        aliases = [
            "SystemAsterisk",
            "SystemExclamation",
            "SystemExit",
            "SystemHand",
            "SystemQuestion",
        ]
        for alias in aliases:
            with self.subTest(alias=alias):
                safe_PlaySound(alias, winsound.SND_ALIAS)

    def test_alias_fallback(self):
        safe_PlaySound('!"$%&/(#+*', winsound.SND_ALIAS)

    def test_alias_nofallback(self):
        safe_PlaySound('!"$%&/(#+*', winsound.SND_ALIAS | winsound.SND_NODEFAULT)

    def test_stopasync(self):
        safe_PlaySound(
            'SystemQuestion',
            winsound.SND_ALIAS | winsound.SND_ASYNC | winsound.SND_LOOP
        )
        time.sleep(0.5)
        safe_PlaySound('SystemQuestion', winsound.SND_ALIAS | winsound.SND_NOSTOP)
        # Issue 8367: PlaySound(None, winsound.SND_PURGE)
        # does not raise on systems without a sound card.
        winsound.PlaySound(None, winsound.SND_PURGE)


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/test/test_winsound.py
DragonHPC/dragon
"""Loading unittests."""

import os
import re
import sys
import traceback
import types
import functools
import warnings

from fnmatch import fnmatch, fnmatchcase

from . import case, suite, util

__unittest = True

# what about .pyc (etc)
# we would need to avoid loading the same tests multiple times
# from '.py', *and* '.pyc'
VALID_MODULE_NAME = re.compile(r'[_a-z]\w*\.py$', re.IGNORECASE)


class _FailedTest(case.TestCase):
    _testMethodName = None

    def __init__(self, method_name, exception):
        self._exception = exception
        super(_FailedTest, self).__init__(method_name)

    def __getattr__(self, name):
        if name != self._testMethodName:
            return super(_FailedTest, self).__getattr__(name)
        def testFailure():
            raise self._exception
        return testFailure


def _make_failed_import_test(name, suiteClass):
    message = 'Failed to import test module: %s\n%s' % (
        name, traceback.format_exc())
    return _make_failed_test(name, ImportError(message), suiteClass, message)

def _make_failed_load_tests(name, exception, suiteClass):
    message = 'Failed to call load_tests:\n%s' % (traceback.format_exc(),)
    return _make_failed_test(
        name, exception, suiteClass, message)

def _make_failed_test(methodname, exception, suiteClass, message):
    test = _FailedTest(methodname, exception)
    return suiteClass((test,)), message

def _make_skipped_test(methodname, exception, suiteClass):
    @case.skip(str(exception))
    def testSkipped(self):
        pass
    attrs = {methodname: testSkipped}
    TestClass = type("ModuleSkipped", (case.TestCase,), attrs)
    return suiteClass((TestClass(methodname),))

def _jython_aware_splitext(path):
    if path.lower().endswith('$py.class'):
        return path[:-9]
    return os.path.splitext(path)[0]


class TestLoader(object):
    """
    This class is responsible for loading tests according to various criteria
    and returning them wrapped in a TestSuite
    """
    testMethodPrefix = 'test'
    sortTestMethodsUsing = staticmethod(util.three_way_cmp)
    testNamePatterns = None
    suiteClass = suite.TestSuite
    _top_level_dir = None

    def __init__(self):
        super(TestLoader, self).__init__()
        self.errors = []
        # Tracks packages which we have called into via load_tests, to
        # avoid infinite re-entrancy.
        self._loading_packages = set()

    def loadTestsFromTestCase(self, testCaseClass):
        """Return a suite of all test cases contained in testCaseClass"""
        if issubclass(testCaseClass, suite.TestSuite):
            raise TypeError("Test cases should not be derived from "
                            "TestSuite. Maybe you meant to derive from "
                            "TestCase?")
        testCaseNames = self.getTestCaseNames(testCaseClass)
        if not testCaseNames and hasattr(testCaseClass, 'runTest'):
            testCaseNames = ['runTest']
        loaded_suite = self.suiteClass(map(testCaseClass, testCaseNames))
        return loaded_suite

    # XXX After Python 3.5, remove backward compatibility hacks for
    # use_load_tests deprecation via *args and **kws.  See issue 16662.
    def loadTestsFromModule(self, module, *args, pattern=None, **kws):
        """Return a suite of all test cases contained in the given module"""
        # This method used to take an undocumented and unofficial
        # use_load_tests argument.  For backward compatibility, we still
        # accept the argument (which can also be the first position) but we
        # ignore it and issue a deprecation warning if it's present.
        if len(args) > 0 or 'use_load_tests' in kws:
            warnings.warn('use_load_tests is deprecated and ignored',
                          DeprecationWarning)
            kws.pop('use_load_tests', None)
        if len(args) > 1:
            # Complain about the number of arguments, but don't forget the
            # required `module` argument.
            complaint = len(args) + 1
            raise TypeError('loadTestsFromModule() takes 1 positional argument but {} were given'.format(complaint))
        if len(kws) != 0:
            # Since the keyword arguments are unsorted (see PEP 468), just
            # pick the alphabetically sorted first argument to complain about,
            # if multiple were given.  At least the error message will be
            # predictable.
            complaint = sorted(kws)[0]
            raise TypeError("loadTestsFromModule() got an unexpected keyword argument '{}'".format(complaint))
        tests = []
        for name in dir(module):
            obj = getattr(module, name)
            if isinstance(obj, type) and issubclass(obj, case.TestCase):
                tests.append(self.loadTestsFromTestCase(obj))

        load_tests = getattr(module, 'load_tests', None)
        tests = self.suiteClass(tests)
        if load_tests is not None:
            try:
                return load_tests(self, tests, pattern)
            except Exception as e:
                error_case, error_message = _make_failed_load_tests(
                    module.__name__, e, self.suiteClass)
                self.errors.append(error_message)
                return error_case
        return tests

    def loadTestsFromName(self, name, module=None):
        """Return a suite of all test cases given a string specifier.

        The name may resolve either to a module, a test case class, a
        test method within a test case class, or a callable object which
        returns a TestCase or TestSuite instance.

        The method optionally resolves the names relative to a given module.
        """
        parts = name.split('.')
        error_case, error_message = None, None
        if module is None:
            parts_copy = parts[:]
            while parts_copy:
                try:
                    module_name = '.'.join(parts_copy)
                    module = __import__(module_name)
                    break
                except ImportError:
                    next_attribute = parts_copy.pop()
                    # Last error so we can give it to the user if needed.
                    error_case, error_message = _make_failed_import_test(
                        next_attribute, self.suiteClass)
                    if not parts_copy:
                        # Even the top level import failed: report that error.
                        self.errors.append(error_message)
                        return error_case
            parts = parts[1:]
        obj = module
        for part in parts:
            try:
                parent, obj = obj, getattr(obj, part)
            except AttributeError as e:
                # We can't traverse some part of the name.
                if (getattr(obj, '__path__', None) is not None
                    and error_case is not None):
                    # This is a package (no __path__ per importlib docs), and we
                    # encountered an error importing something. We cannot tell
                    # the difference between package.WrongNameTestClass and
                    # package.wrong_module_name so we just report the
                    # ImportError - it is more informative.
                    self.errors.append(error_message)
                    return error_case
                else:
                    # Otherwise, we signal that an AttributeError has occurred.
                    error_case, error_message = _make_failed_test(
                        part, e, self.suiteClass,
                        'Failed to access attribute:\n%s' % (
                            traceback.format_exc(),))
                    self.errors.append(error_message)
                    return error_case

        if isinstance(obj, types.ModuleType):
            return self.loadTestsFromModule(obj)
        elif isinstance(obj, type) and issubclass(obj, case.TestCase):
            return self.loadTestsFromTestCase(obj)
        elif (isinstance(obj, types.FunctionType) and
              isinstance(parent, type) and
              issubclass(parent, case.TestCase)):
            name = parts[-1]
            inst = parent(name)
            # static methods follow a different path
            if not isinstance(getattr(inst, name), types.FunctionType):
                return self.suiteClass([inst])
        elif isinstance(obj, suite.TestSuite):
            return obj
        if callable(obj):
            test = obj()
            if isinstance(test, suite.TestSuite):
                return test
            elif isinstance(test, case.TestCase):
                return self.suiteClass([test])
            else:
                raise TypeError("calling %s returned %s, not a test" %
                                (obj, test))
        else:
            raise TypeError("don't know how to make test from: %s" % obj)

    def loadTestsFromNames(self, names, module=None):
        """Return a suite of all test cases found using the given sequence
        of string specifiers. See 'loadTestsFromName()'.
        """
        suites = [self.loadTestsFromName(name, module) for name in names]
        return self.suiteClass(suites)

    def getTestCaseNames(self, testCaseClass):
        """Return a sorted sequence of method names found within testCaseClass
        """
        def shouldIncludeMethod(attrname):
            if not attrname.startswith(self.testMethodPrefix):
                return False
            testFunc = getattr(testCaseClass, attrname)
            if not callable(testFunc):
                return False
            fullName = f'%s.%s.%s' % (
                testCaseClass.__module__, testCaseClass.__qualname__, attrname
            )
            return self.testNamePatterns is None or \
                any(fnmatchcase(fullName, pattern) for pattern in self.testNamePatterns)
        testFnNames = list(filter(shouldIncludeMethod, dir(testCaseClass)))
        if self.sortTestMethodsUsing:
            testFnNames.sort(key=functools.cmp_to_key(self.sortTestMethodsUsing))
        return testFnNames

    def discover(self, start_dir, pattern='test*.py', top_level_dir=None):
        """Find and return all test modules from the specified start
        directory, recursing into subdirectories to find them and return all
        tests found within them. Only test files that match the pattern will
        be loaded. (Using shell style pattern matching.)

        All test modules must be importable from the top level of the project.
        If the start directory is not the top level directory then the top
        level directory must be specified separately.

        If a test package name (directory with '__init__.py') matches the
        pattern then the package will be checked for a 'load_tests' function. If
        this exists then it will be called with (loader, tests, pattern) unless
        the package has already had load_tests called from the same discovery
        invocation, in which case the package module object is not scanned for
        tests - this ensures that when a package uses discover to further
        discover child tests that infinite recursion does not happen.

        If load_tests exists then discovery does *not* recurse into the package,
        load_tests is responsible for loading all tests in the package.

        The pattern is deliberately not stored as a loader attribute so that
        packages can continue discovery themselves. top_level_dir is stored so
        load_tests does not need to pass this argument in to loader.discover().

        Paths are sorted before being imported to ensure reproducible execution
        order even on filesystems with non-alphabetical ordering like ext3/4.
        """
        set_implicit_top = False
        if top_level_dir is None and self._top_level_dir is not None:
            # make top_level_dir optional if called from load_tests in a package
            top_level_dir = self._top_level_dir
        elif top_level_dir is None:
            set_implicit_top = True
            top_level_dir = start_dir

        top_level_dir = os.path.abspath(top_level_dir)

        if not top_level_dir in sys.path:
            # all test modules must be importable from the top level directory
            # should we *unconditionally* put the start directory in first
            # in sys.path to minimise likelihood of conflicts between installed
            # modules and development versions?
            sys.path.insert(0, top_level_dir)
        self._top_level_dir = top_level_dir

        is_not_importable = False
        is_namespace = False
        tests = []
        if os.path.isdir(os.path.abspath(start_dir)):
            start_dir = os.path.abspath(start_dir)
            if start_dir != top_level_dir:
                is_not_importable = not os.path.isfile(os.path.join(start_dir, '__init__.py'))
        else:
            # support for discovery from dotted module names
            try:
                __import__(start_dir)
            except ImportError:
                is_not_importable = True
            else:
                the_module = sys.modules[start_dir]
                top_part = start_dir.split('.')[0]
                try:
                    start_dir = os.path.abspath(
                       os.path.dirname((the_module.__file__)))
                except AttributeError:
                    # look for namespace packages
                    try:
                        spec = the_module.__spec__
                    except AttributeError:
                        spec = None

                    if spec and spec.loader is None:
                        if spec.submodule_search_locations is not None:
                            is_namespace = True

                            for path in the_module.__path__:
                                if (not set_implicit_top and
                                    not path.startswith(top_level_dir)):
                                    continue
                                self._top_level_dir = \
                                    (path.split(the_module.__name__
                                         .replace(".", os.path.sep))[0])
                                tests.extend(self._find_tests(path,
                                                              pattern,
                                                              namespace=True))
                    elif the_module.__name__ in sys.builtin_module_names:
                        # builtin module
                        raise TypeError('Can not use builtin modules '
                                        'as dotted module names') from None
                    else:
                        raise TypeError(
                            'don\'t know how to discover from {!r}'
                            .format(the_module)) from None

                if set_implicit_top:
                    if not is_namespace:
                        self._top_level_dir = \
                           self._get_directory_containing_module(top_part)
                        sys.path.remove(top_level_dir)
                    else:
                        sys.path.remove(top_level_dir)

        if is_not_importable:
            raise ImportError('Start directory is not importable: %r' % start_dir)

        if not is_namespace:
            tests = list(self._find_tests(start_dir, pattern))
        return self.suiteClass(tests)

    def _get_directory_containing_module(self, module_name):
        module = sys.modules[module_name]
        full_path = os.path.abspath(module.__file__)

        if os.path.basename(full_path).lower().startswith('__init__.py'):
            return os.path.dirname(os.path.dirname(full_path))
        else:
            # here we have been given a module rather than a package - so
            # all we can do is search the *same* directory the module is in
            # should an exception be raised instead
            return os.path.dirname(full_path)

    def _get_name_from_path(self, path):
        if path == self._top_level_dir:
            return '.'
        path = _jython_aware_splitext(os.path.normpath(path))

        _relpath = os.path.relpath(path, self._top_level_dir)
        assert not os.path.isabs(_relpath), "Path must be within the project"
        assert not _relpath.startswith('..'), "Path must be within the project"

        name = _relpath.replace(os.path.sep, '.')
        return name

    def _get_module_from_name(self, name):
        __import__(name)
        return sys.modules[name]

    def _match_path(self, path, full_path, pattern):
        # override this method to use alternative matching strategy
        return fnmatch(path, pattern)

    def _find_tests(self, start_dir, pattern, namespace=False):
        """Used by discovery. Yields test suites it loads."""
        # Handle the __init__ in this package
        name = self._get_name_from_path(start_dir)
        # name is '.' when start_dir == top_level_dir (and top_level_dir is by
        # definition not a package).
        if name != '.' and name not in self._loading_packages:
            # name is in self._loading_packages while we have called into
            # loadTestsFromModule with name.
            tests, should_recurse = self._find_test_path(
                start_dir, pattern, namespace)
            if tests is not None:
                yield tests
            if not should_recurse:
                # Either an error occurred, or load_tests was used by the
                # package.
                return
        # Handle the contents.
        paths = sorted(os.listdir(start_dir))
        for path in paths:
            full_path = os.path.join(start_dir, path)
            tests, should_recurse = self._find_test_path(
                full_path, pattern, namespace)
            if tests is not None:
                yield tests
            if should_recurse:
                # we found a package that didn't use load_tests.
                name = self._get_name_from_path(full_path)
                self._loading_packages.add(name)
                try:
                    yield from self._find_tests(full_path, pattern, namespace)
                finally:
                    self._loading_packages.discard(name)

    def _find_test_path(self, full_path, pattern, namespace=False):
        """Used by discovery.

        Loads tests from a single file, or a directories' __init__.py when
        passed the directory.

        Returns a tuple (None_or_tests_from_file, should_recurse).
        """
        basename = os.path.basename(full_path)
        if os.path.isfile(full_path):
            if not VALID_MODULE_NAME.match(basename):
                # valid Python identifiers only
                return None, False
            if not self._match_path(basename, full_path, pattern):
                return None, False
            # if the test file matches, load it
            name = self._get_name_from_path(full_path)
            try:
                module = self._get_module_from_name(name)
            except case.SkipTest as e:
                return _make_skipped_test(name, e, self.suiteClass), False
            except:
                error_case, error_message = \
                    _make_failed_import_test(name, self.suiteClass)
                self.errors.append(error_message)
                return error_case, False
            else:
                mod_file = os.path.abspath(
                    getattr(module, '__file__', full_path))
                realpath = _jython_aware_splitext(
                    os.path.realpath(mod_file))
                fullpath_noext = _jython_aware_splitext(
                    os.path.realpath(full_path))
                if realpath.lower() != fullpath_noext.lower():
                    module_dir = os.path.dirname(realpath)
                    mod_name = _jython_aware_splitext(
                        os.path.basename(full_path))
                    expected_dir = os.path.dirname(full_path)
                    msg = ("%r module incorrectly imported from %r. Expected "
                           "%r. Is this module globally installed?")
                    raise ImportError(
                        msg % (mod_name, module_dir, expected_dir))
                return self.loadTestsFromModule(module, pattern=pattern), False
        elif os.path.isdir(full_path):
            if (not namespace and
                not os.path.isfile(os.path.join(full_path, '__init__.py'))):
                return None, False

            load_tests = None
            tests = None
            name = self._get_name_from_path(full_path)
            try:
                package = self._get_module_from_name(name)
            except case.SkipTest as e:
                return _make_skipped_test(name, e, self.suiteClass), False
            except:
                error_case, error_message = \
                    _make_failed_import_test(name, self.suiteClass)
                self.errors.append(error_message)
                return error_case, False
            else:
                load_tests = getattr(package, 'load_tests', None)
                # Mark this package as being in load_tests (possibly ;))
                self._loading_packages.add(name)
                try:
                    tests = self.loadTestsFromModule(package, pattern=pattern)
                    if load_tests is not None:
                        # loadTestsFromModule(package) has loaded tests for us.
                        return tests, False
                    return tests, True
                finally:
                    self._loading_packages.discard(name)
        else:
            return None, False


defaultTestLoader = TestLoader()


def _makeLoader(prefix, sortUsing, suiteClass=None, testNamePatterns=None):
    loader = TestLoader()
    loader.sortTestMethodsUsing = sortUsing
    loader.testMethodPrefix = prefix
    loader.testNamePatterns = testNamePatterns
    if suiteClass:
        loader.suiteClass = suiteClass
    return loader

def getTestCaseNames(testCaseClass, prefix, sortUsing=util.three_way_cmp, testNamePatterns=None):
    return _makeLoader(prefix, sortUsing, testNamePatterns=testNamePatterns).getTestCaseNames(testCaseClass)

def makeSuite(testCaseClass, prefix='test', sortUsing=util.three_way_cmp,
              suiteClass=suite.TestSuite):
    return _makeLoader(prefix, sortUsing, suiteClass).loadTestsFromTestCase(
        testCaseClass)

def findTestCases(module, prefix='test', sortUsing=util.three_way_cmp,
                  suiteClass=suite.TestSuite):
    return _makeLoader(prefix, sortUsing, suiteClass).loadTestsFromModule(\
        module)

external/cpython_master/Lib/unittest/loader.py
DragonHPC/dragon
"""Running tests"""

import sys
import time
import warnings

from . import result
from .signals import registerResult

__unittest = True


class _WritelnDecorator(object):
    """Used to decorate file-like objects with a handy 'writeln' method"""
    def __init__(self,stream):
        self.stream = stream

    def __getattr__(self, attr):
        if attr in ('stream', '__getstate__'):
            raise AttributeError(attr)
        return getattr(self.stream,attr)

    def writeln(self, arg=None):
        if arg:
            self.write(arg)
        self.write('\n') # text-mode streams translate to \r\n if needed


class TextTestResult(result.TestResult):
    """A test result class that can print formatted text results to a stream.

    Used by TextTestRunner.
    """
    separator1 = '=' * 70
    separator2 = '-' * 70

    def __init__(self, stream, descriptions, verbosity):
        super(TextTestResult, self).__init__(stream, descriptions, verbosity)
        self.stream = stream
        self.showAll = verbosity > 1
        self.dots = verbosity == 1
        self.descriptions = descriptions

    def getDescription(self, test):
        doc_first_line = test.shortDescription()
        if self.descriptions and doc_first_line:
            return '\n'.join((str(test), doc_first_line))
        else:
            return str(test)

    def startTest(self, test):
        super(TextTestResult, self).startTest(test)
        if self.showAll:
            self.stream.write(self.getDescription(test))
            self.stream.write(" ... ")
            self.stream.flush()

    def addSuccess(self, test):
        super(TextTestResult, self).addSuccess(test)
        if self.showAll:
            self.stream.writeln("ok")
        elif self.dots:
            self.stream.write('.')
            self.stream.flush()

    def addError(self, test, err):
        super(TextTestResult, self).addError(test, err)
        if self.showAll:
            self.stream.writeln("ERROR")
        elif self.dots:
            self.stream.write('E')
            self.stream.flush()

    def addFailure(self, test, err):
        super(TextTestResult, self).addFailure(test, err)
        if self.showAll:
            self.stream.writeln("FAIL")
        elif self.dots:
            self.stream.write('F')
            self.stream.flush()

    def addSkip(self, test, reason):
        super(TextTestResult, self).addSkip(test, reason)
        if self.showAll:
            self.stream.writeln("skipped {0!r}".format(reason))
        elif self.dots:
            self.stream.write("s")
            self.stream.flush()

    def addExpectedFailure(self, test, err):
        super(TextTestResult, self).addExpectedFailure(test, err)
        if self.showAll:
            self.stream.writeln("expected failure")
        elif self.dots:
            self.stream.write("x")
            self.stream.flush()

    def addUnexpectedSuccess(self, test):
        super(TextTestResult, self).addUnexpectedSuccess(test)
        if self.showAll:
            self.stream.writeln("unexpected success")
        elif self.dots:
            self.stream.write("u")
            self.stream.flush()

    def printErrors(self):
        if self.dots or self.showAll:
            self.stream.writeln()
        self.printErrorList('ERROR', self.errors)
        self.printErrorList('FAIL', self.failures)

    def printErrorList(self, flavour, errors):
        for test, err in errors:
            self.stream.writeln(self.separator1)
            self.stream.writeln("%s: %s" % (flavour,self.getDescription(test)))
            self.stream.writeln(self.separator2)
            self.stream.writeln("%s" % err)


class TextTestRunner(object):
    """A test runner class that displays results in textual form.

    It prints out the names of tests as they are run, errors as they
    occur, and a summary of the results at the end of the test run.
    """
    resultclass = TextTestResult

    def __init__(self, stream=None, descriptions=True, verbosity=1,
                 failfast=False, buffer=False, resultclass=None, warnings=None,
                 *, tb_locals=False):
        """Construct a TextTestRunner.

        Subclasses should accept **kwargs to ensure compatibility as the
        interface changes.
        """
        if stream is None:
            stream = sys.stderr
        self.stream = _WritelnDecorator(stream)
        self.descriptions = descriptions
        self.verbosity = verbosity
        self.failfast = failfast
        self.buffer = buffer
        self.tb_locals = tb_locals
        self.warnings = warnings
        if resultclass is not None:
            self.resultclass = resultclass

    def _makeResult(self):
        return self.resultclass(self.stream, self.descriptions, self.verbosity)

    def run(self, test):
        "Run the given test case or test suite."
        result = self._makeResult()
        registerResult(result)
        result.failfast = self.failfast
        result.buffer = self.buffer
        result.tb_locals = self.tb_locals
        with warnings.catch_warnings():
            if self.warnings:
                # if self.warnings is set, use it to filter all the warnings
                warnings.simplefilter(self.warnings)
                # if the filter is 'default' or 'always', special-case the
                # warnings from the deprecated unittest methods to show them
                # no more than once per module, because they can be fairly
                # noisy.  The -Wd and -Wa flags can be used to bypass this
                # only when self.warnings is None.
                if self.warnings in ['default', 'always']:
                    warnings.filterwarnings('module',
                            category=DeprecationWarning,
                            message=r'Please use assert\w+ instead.')
            startTime = time.perf_counter()
            startTestRun = getattr(result, 'startTestRun', None)
            if startTestRun is not None:
                startTestRun()
            try:
                test(result)
            finally:
                stopTestRun = getattr(result, 'stopTestRun', None)
                if stopTestRun is not None:
                    stopTestRun()
            stopTime = time.perf_counter()
        timeTaken = stopTime - startTime
        result.printErrors()
        if hasattr(result, 'separator2'):
            self.stream.writeln(result.separator2)
        run = result.testsRun
        self.stream.writeln("Ran %d test%s in %.3fs" %
                            (run, run != 1 and "s" or "", timeTaken))
        self.stream.writeln()

        expectedFails = unexpectedSuccesses = skipped = 0
        try:
            results = map(len, (result.expectedFailures,
                                result.unexpectedSuccesses,
                                result.skipped))
        except AttributeError:
            pass
        else:
            expectedFails, unexpectedSuccesses, skipped = results

        infos = []
        if not result.wasSuccessful():
            self.stream.write("FAILED")
            failed, errored = len(result.failures), len(result.errors)
            if failed:
                infos.append("failures=%d" % failed)
            if errored:
                infos.append("errors=%d" % errored)
        else:
            self.stream.write("OK")
        if skipped:
            infos.append("skipped=%d" % skipped)
        if expectedFails:
            infos.append("expected failures=%d" % expectedFails)
        if unexpectedSuccesses:
            infos.append("unexpected successes=%d" % unexpectedSuccesses)
        if infos:
            self.stream.writeln(" (%s)" % (", ".join(infos),))
        else:
            self.stream.write("\n")
        return result

external/cpython_master/Lib/unittest/runner.py
DragonHPC/dragon
import datetime
import warnings
import weakref
import unittest
from itertools import product


class Test_Assertions(unittest.TestCase):
    def test_AlmostEqual(self):
        self.assertAlmostEqual(1.00000001, 1.0)
        self.assertNotAlmostEqual(1.0000001, 1.0)
        self.assertRaises(self.failureException,
                          self.assertAlmostEqual, 1.0000001, 1.0)
        self.assertRaises(self.failureException,
                          self.assertNotAlmostEqual, 1.00000001, 1.0)

        self.assertAlmostEqual(1.1, 1.0, places=0)
        self.assertRaises(self.failureException,
                          self.assertAlmostEqual, 1.1, 1.0, places=1)

        self.assertAlmostEqual(0, .1+.1j, places=0)
        self.assertNotAlmostEqual(0, .1+.1j, places=1)
        self.assertRaises(self.failureException,
                          self.assertAlmostEqual, 0, .1+.1j, places=1)
        self.assertRaises(self.failureException,
                          self.assertNotAlmostEqual, 0, .1+.1j, places=0)

        self.assertAlmostEqual(float('inf'), float('inf'))
        self.assertRaises(self.failureException, self.assertNotAlmostEqual,
                          float('inf'), float('inf'))

    def test_AmostEqualWithDelta(self):
        self.assertAlmostEqual(1.1, 1.0, delta=0.5)
        self.assertAlmostEqual(1.0, 1.1, delta=0.5)
        self.assertNotAlmostEqual(1.1, 1.0, delta=0.05)
        self.assertNotAlmostEqual(1.0, 1.1, delta=0.05)

        self.assertAlmostEqual(1.0, 1.0, delta=0.5)
        self.assertRaises(self.failureException, self.assertNotAlmostEqual,
                          1.0, 1.0, delta=0.5)

        self.assertRaises(self.failureException, self.assertAlmostEqual,
                          1.1, 1.0, delta=0.05)
        self.assertRaises(self.failureException, self.assertNotAlmostEqual,
                          1.1, 1.0, delta=0.5)

        self.assertRaises(TypeError, self.assertAlmostEqual,
                          1.1, 1.0, places=2, delta=2)
        self.assertRaises(TypeError, self.assertNotAlmostEqual,
                          1.1, 1.0, places=2, delta=2)

        first = datetime.datetime.now()
        second = first + datetime.timedelta(seconds=10)
        self.assertAlmostEqual(first, second,
                               delta=datetime.timedelta(seconds=20))
        self.assertNotAlmostEqual(first, second,
                                  delta=datetime.timedelta(seconds=5))

    def test_assertRaises(self):
        def _raise(e):
            raise e
        self.assertRaises(KeyError, _raise, KeyError)
        self.assertRaises(KeyError, _raise, KeyError("key"))
        try:
            self.assertRaises(KeyError, lambda: None)
        except self.failureException as e:
            self.assertIn("KeyError not raised", str(e))
        else:
            self.fail("assertRaises() didn't fail")
        try:
            self.assertRaises(KeyError, _raise, ValueError)
        except ValueError:
            pass
        else:
            self.fail("assertRaises() didn't let exception pass through")
        with self.assertRaises(KeyError) as cm:
            try:
                raise KeyError
            except Exception as e:
                exc = e
                raise
        self.assertIs(cm.exception, exc)

        with self.assertRaises(KeyError):
            raise KeyError("key")
        try:
            with self.assertRaises(KeyError):
                pass
        except self.failureException as e:
            self.assertIn("KeyError not raised", str(e))
        else:
            self.fail("assertRaises() didn't fail")
        try:
            with self.assertRaises(KeyError):
                raise ValueError
        except ValueError:
            pass
        else:
            self.fail("assertRaises() didn't let exception pass through")

    def test_assertRaises_frames_survival(self):
        # Issue #9815: assertRaises should avoid keeping local variables
        # in a traceback alive.
        class A:
            pass
        wr = None

        class Foo(unittest.TestCase):

            def foo(self):
                nonlocal wr
                a = A()
                wr = weakref.ref(a)
                try:
                    raise OSError
                except OSError:
                    raise ValueError

            def test_functional(self):
                self.assertRaises(ValueError, self.foo)

            def test_with(self):
                with self.assertRaises(ValueError):
                    self.foo()

        Foo("test_functional").run()
        self.assertIsNone(wr())
        Foo("test_with").run()
        self.assertIsNone(wr())

    def testAssertNotRegex(self):
        self.assertNotRegex('Ala ma kota', r'r+')
        try:
            self.assertNotRegex('Ala ma kota', r'k.t', 'Message')
        except self.failureException as e:
            self.assertIn('Message', e.args[0])
        else:
            self.fail('assertNotRegex should have failed.')


class TestLongMessage(unittest.TestCase):
    """Test that the individual asserts honour longMessage.
    This actually tests all the message behaviour for
    asserts that use longMessage."""

    def setUp(self):
        class TestableTestFalse(unittest.TestCase):
            longMessage = False
            failureException = self.failureException

            def testTest(self):
                pass

        class TestableTestTrue(unittest.TestCase):
            longMessage = True
            failureException = self.failureException

            def testTest(self):
                pass

        self.testableTrue = TestableTestTrue('testTest')
        self.testableFalse = TestableTestFalse('testTest')

    def testDefault(self):
        self.assertTrue(unittest.TestCase.longMessage)

    def test_formatMsg(self):
        self.assertEqual(self.testableFalse._formatMessage(None, "foo"), "foo")
        self.assertEqual(self.testableFalse._formatMessage("foo", "bar"), "foo")

        self.assertEqual(self.testableTrue._formatMessage(None, "foo"), "foo")
        self.assertEqual(self.testableTrue._formatMessage("foo", "bar"), "bar : foo")

        # This blows up if _formatMessage uses string concatenation
        self.testableTrue._formatMessage(object(), 'foo')

    def test_formatMessage_unicode_error(self):
        one = ''.join(chr(i) for i in range(255))
        # this used to cause a UnicodeDecodeError constructing msg
        self.testableTrue._formatMessage(one, '\uFFFD')

    def assertMessages(self, methodName, args, errors):
        """
        Check that methodName(*args) raises the correct error messages.
        errors should be a list of 4 regex that match the error when:
          1) longMessage = False and no msg passed;
          2) longMessage = False and msg passed;
          3) longMessage = True and no msg passed;
          4) longMessage = True and msg passed;
        """
        def getMethod(i):
            useTestableFalse  = i < 2
            if useTestableFalse:
                test = self.testableFalse
            else:
                test = self.testableTrue
            return getattr(test, methodName)

        for i, expected_regex in enumerate(errors):
            testMethod = getMethod(i)
            kwargs = {}
            withMsg = i % 2
            if withMsg:
                kwargs = {"msg": "oops"}

            with self.assertRaisesRegex(self.failureException,
                                        expected_regex=expected_regex):
                testMethod(*args, **kwargs)

    def testAssertTrue(self):
        self.assertMessages('assertTrue', (False,),
                            ["^False is not true$", "^oops$", "^False is not true$",
                             "^False is not true : oops$"])

    def testAssertFalse(self):
        self.assertMessages('assertFalse', (True,),
                            ["^True is not false$", "^oops$", "^True is not false$",
                             "^True is not false : oops$"])

    def testNotEqual(self):
        self.assertMessages('assertNotEqual', (1, 1),
                            ["^1 == 1$", "^oops$", "^1 == 1$",
                             "^1 == 1 : oops$"])

    def testAlmostEqual(self):
        self.assertMessages(
            'assertAlmostEqual', (1, 2),
            [r"^1 != 2 within 7 places \(1 difference\)$", "^oops$",
             r"^1 != 2 within 7 places \(1 difference\)$",
             r"^1 != 2 within 7 places \(1 difference\) : oops$"])

    def testNotAlmostEqual(self):
        self.assertMessages('assertNotAlmostEqual', (1, 1),
                            ["^1 == 1 within 7 places$", "^oops$",
                             "^1 == 1 within 7 places$", "^1 == 1 within 7 places : oops$"])

    def test_baseAssertEqual(self):
        self.assertMessages('_baseAssertEqual', (1, 2),
                            ["^1 != 2$", "^oops$", "^1 != 2$", "^1 != 2 : oops$"])

    def testAssertSequenceEqual(self):
        # Error messages are multiline so not testing on full message
        # assertTupleEqual and assertListEqual delegate to this method
        self.assertMessages('assertSequenceEqual', ([], [None]),
                            [r"\+ \[None\]$", "^oops$", r"\+ \[None\]$",
                             r"\+ \[None\] : oops$"])

    def testAssertSetEqual(self):
        self.assertMessages('assertSetEqual', (set(), set([None])),
                            ["None$", "^oops$", "None$",
                             "None : oops$"])

    def testAssertIn(self):
        self.assertMessages('assertIn', (None, []),
                            [r'^None not found in \[\]$', "^oops$",
                             r'^None not found in \[\]$',
                             r'^None not found in \[\] : oops$'])

    def testAssertNotIn(self):
        self.assertMessages('assertNotIn', (None, [None]),
                            [r'^None unexpectedly found in \[None\]$', "^oops$",
                             r'^None unexpectedly found in \[None\]$',
                             r'^None unexpectedly found in \[None\] : oops$'])

    def testAssertDictEqual(self):
        self.assertMessages('assertDictEqual', ({}, {'key': 'value'}),
                            [r"\+ \{'key': 'value'\}$", "^oops$",
                             r"\+ \{'key': 'value'\}$",
                             r"\+ \{'key': 'value'\} : oops$"])

    def testAssertDictContainsSubset(self):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", DeprecationWarning)

            self.assertMessages('assertDictContainsSubset', ({'key': 'value'}, {}),
                                ["^Missing: 'key'$", "^oops$",
                                 "^Missing: 'key'$",
                                 "^Missing: 'key' : oops$"])

    def testAssertMultiLineEqual(self):
        self.assertMessages('assertMultiLineEqual', ("", "foo"),
                            [r"\+ foo$", "^oops$",
                             r"\+ foo$",
                             r"\+ foo : oops$"])

    def testAssertLess(self):
        self.assertMessages('assertLess', (2, 1),
                            ["^2 not less than 1$", "^oops$",
                             "^2 not less than 1$", "^2 not less than 1 : oops$"])

    def testAssertLessEqual(self):
        self.assertMessages('assertLessEqual', (2, 1),
                            ["^2 not less than or equal to 1$", "^oops$",
                             "^2 not less than or equal to 1$",
                             "^2 not less than or equal to 1 : oops$"])

    def testAssertGreater(self):
        self.assertMessages('assertGreater', (1, 2),
                            ["^1 not greater than 2$", "^oops$",
                             "^1 not greater than 2$",
                             "^1 not greater than 2 : oops$"])

    def testAssertGreaterEqual(self):
        self.assertMessages('assertGreaterEqual', (1, 2),
                            ["^1 not greater than or equal to 2$", "^oops$",
                             "^1 not greater than or equal to 2$",
                             "^1 not greater than or equal to 2 : oops$"])

    def testAssertIsNone(self):
        self.assertMessages('assertIsNone', ('not None',),
                            ["^'not None' is not None$", "^oops$",
                             "^'not None' is not None$",
                             "^'not None' is not None : oops$"])

    def testAssertIsNotNone(self):
        self.assertMessages('assertIsNotNone', (None,),
                            ["^unexpectedly None$", "^oops$",
                             "^unexpectedly None$",
                             "^unexpectedly None : oops$"])

    def testAssertIs(self):
        self.assertMessages('assertIs', (None, 'foo'),
                            ["^None is not 'foo'$", "^oops$",
                             "^None is not 'foo'$",
                             "^None is not 'foo' : oops$"])

    def testAssertIsNot(self):
        self.assertMessages('assertIsNot', (None, None),
                            ["^unexpectedly identical: None$", "^oops$",
                             "^unexpectedly identical: None$",
                             "^unexpectedly identical: None : oops$"])

    def testAssertRegex(self):
        self.assertMessages('assertRegex', ('foo', 'bar'),
                            ["^Regex didn't match:",
                             "^oops$",
                             "^Regex didn't match:",
                             "^Regex didn't match: (.*) : oops$"])

    def testAssertNotRegex(self):
        self.assertMessages('assertNotRegex', ('foo', 'foo'),
                            ["^Regex matched:",
                             "^oops$",
                             "^Regex matched:",
                             "^Regex matched: (.*) : oops$"])


    def assertMessagesCM(self, methodName, args, func, errors):
        """
        Check that the correct error messages are raised while executing:
          with method(*args):
              func()
        *errors* should be a list of 4 regex that match the error when:
          1) longMessage = False and no msg passed;
          2) longMessage = False and msg passed;
          3) longMessage = True and no msg passed;
          4) longMessage = True and msg passed;
        """
        p = product((self.testableFalse, self.testableTrue),
                    ({}, {"msg": "oops"}))
        for (cls, kwargs), err in zip(p, errors):
            method = getattr(cls, methodName)
            with self.assertRaisesRegex(cls.failureException, err):
                with method(*args, **kwargs) as cm:
                    func()

    def testAssertRaises(self):
        self.assertMessagesCM('assertRaises', (TypeError,), lambda: None,
                              ['^TypeError not raised$', '^oops$',
                               '^TypeError not raised$',
                               '^TypeError not raised : oops$'])

    def testAssertRaisesRegex(self):
        # test error not raised
        self.assertMessagesCM('assertRaisesRegex', (TypeError, 'unused regex'),
                              lambda: None,
                              ['^TypeError not raised$', '^oops$',
                               '^TypeError not raised$',
                               '^TypeError not raised : oops$'])
        # test error raised but with wrong message
        def raise_wrong_message():
            raise TypeError('foo')
        self.assertMessagesCM('assertRaisesRegex', (TypeError, 'regex'),
                              raise_wrong_message,
                              ['^"regex" does not match "foo"$', '^oops$',
                               '^"regex" does not match "foo"$',
                               '^"regex" does not match "foo" : oops$'])

    def testAssertWarns(self):
        self.assertMessagesCM('assertWarns', (UserWarning,), lambda: None,
                              ['^UserWarning not triggered$', '^oops$',
                               '^UserWarning not triggered$',
                               '^UserWarning not triggered : oops$'])

    def testAssertWarnsRegex(self):
        # test error not raised
        self.assertMessagesCM('assertWarnsRegex', (UserWarning, 'unused regex'),
                              lambda: None,
                              ['^UserWarning not triggered$', '^oops$',
                               '^UserWarning not triggered$',
                               '^UserWarning not triggered : oops$'])
        # test warning raised but with wrong message
        def raise_wrong_message():
            warnings.warn('foo')
        self.assertMessagesCM('assertWarnsRegex', (UserWarning, 'regex'),
                              raise_wrong_message,
                              ['^"regex" does not match "foo"$', '^oops$',
                               '^"regex" does not match "foo"$',
                               '^"regex" does not match "foo" : oops$'])


if __name__ == "__main__":
    unittest.main()

external/cpython_master/Lib/unittest/test/test_assertions.py
DragonHPC/dragon
"""Python version compatibility support for minidom.

This module contains internal implementation details and
should not be imported; use xml.dom.minidom instead.
"""

# This module should only be imported using "import *".
#
# The following names are defined:
#
#   NodeList      -- lightest possible NodeList implementation
#
#   EmptyNodeList -- lightest possible NodeList that is guaranteed to
#                    remain empty (immutable)
#
#   StringTypes   -- tuple of defined string types
#
#   defproperty   -- function used in conjunction with GetattrMagic;
#                    using these together is needed to make them work
#                    as efficiently as possible in both Python 2.2+
#                    and older versions.  For example:
#
#                        class MyClass(GetattrMagic):
#                            def _get_myattr(self):
#                                return something
#
#                        defproperty(MyClass, "myattr",
#                                    "return some value")
#
#                    For Python 2.2 and newer, this will construct a
#                    property object on the class, which avoids
#                    needing to override __getattr__().  It will only
#                    work for read-only attributes.
#
#                    For older versions of Python, inheriting from
#                    GetattrMagic will use the traditional
#                    __getattr__() hackery to achieve the same effect,
#                    but less efficiently.
#
#                    defproperty() should be used for each version of
#                    the relevant _get_<property>() function.

__all__ = ["NodeList", "EmptyNodeList", "StringTypes", "defproperty"]

import xml.dom

StringTypes = (str,)


class NodeList(list):
    __slots__ = ()

    def item(self, index):
        if 0 <= index < len(self):
            return self[index]

    def _get_length(self):
        return len(self)

    def _set_length(self, value):
        raise xml.dom.NoModificationAllowedErr(
            "attempt to modify read-only attribute 'length'")

    length = property(_get_length, _set_length,
                      doc="The number of nodes in the NodeList.")

    # For backward compatibility
    def __setstate__(self, state):
        if state is None:
            state = []
        self[:] = state


class EmptyNodeList(tuple):
    __slots__ = ()

    def __add__(self, other):
        NL = NodeList()
        NL.extend(other)
        return NL

    def __radd__(self, other):
        NL = NodeList()
        NL.extend(other)
        return NL

    def item(self, index):
        return None

    def _get_length(self):
        return 0

    def _set_length(self, value):
        raise xml.dom.NoModificationAllowedErr(
            "attempt to modify read-only attribute 'length'")

    length = property(_get_length, _set_length,
                      doc="The number of nodes in the NodeList.")


def defproperty(klass, name, doc):
    get = getattr(klass, ("_get_" + name))
    def set(self, value, name=name):
        raise xml.dom.NoModificationAllowedErr(
            "attempt to modify read-only attribute " + repr(name))
    assert not hasattr(klass, "_set_" + name), \
           "expected not to find _set_" + name
    prop = property(get, set, doc=doc)
    setattr(klass, name, prop)

external/cpython_master/Lib/xml/dom/minicompat.py
DragonHPC/dragon
"""Interface to the Expat non-validating XML parser."""
import sys

from pyexpat import *

# provide pyexpat submodules as xml.parsers.expat submodules
sys.modules['xml.parsers.expat.model'] = model
sys.modules['xml.parsers.expat.errors'] = errors

external/cpython_master/Lib/xml/parsers/expat.py
DragonHPC/dragon
#!/usr/bin/env python
"""
This script is used to build "official" universal installers on macOS.

NEW for 3.9.0 and backports:
- 2.7 end-of-life issues:
    - Python 3 installs now update the Current version link
      in /Library/Frameworks/Python.framework/Versions
- fully support running under Python 3 as well as 2.7
- support building on newer macOS systems with SIP
- fully support building on macOS 10.9+
- support 10.6+ on best effort
- support bypassing docs build by supplying a prebuilt
    docs html tarball in the third-party source library,
    in the format and filename conventional of those
    downloadable from python.org:
        python-3.x.y-docs-html.tar.bz2

NEW for 3.7.0:
- support Intel 64-bit-only () and 32-bit-only installer builds
- build and use internal Tcl/Tk 8.6 for 10.6+ builds
- deprecate use of explicit SDK (--sdk-path=) since all but the oldest
  versions of Xcode support implicit setting of an SDK via environment
  variables (SDKROOT and friends, see the xcrun man page for more info).
  The SDK stuff was primarily needed for building universal installers
  for 10.4; so as of 3.7.0, building installers for 10.4 is no longer
  supported with build-installer.
- use generic "gcc" as compiler (CC env var) rather than "gcc-4.2"

TODO:
- test building with SDKROOT and DEVELOPER_DIR xcrun env variables

Usage: see USAGE variable in the script.
"""
import platform, os, sys, getopt, textwrap, shutil, stat, time, pwd, grp
try:
    import urllib2 as urllib_request
except ImportError:
    import urllib.request as urllib_request

STAT_0o755 = ( stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR
             | stat.S_IRGRP |                stat.S_IXGRP
             | stat.S_IROTH |                stat.S_IXOTH )

STAT_0o775 = ( stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR
             | stat.S_IRGRP | stat.S_IWGRP | stat.S_IXGRP
             | stat.S_IROTH |                stat.S_IXOTH )

INCLUDE_TIMESTAMP = 1
VERBOSE = 1

RUNNING_ON_PYTHON2 = sys.version_info.major == 2

if RUNNING_ON_PYTHON2:
    from plistlib import writePlist
else:
    from plistlib import dump
    def writePlist(path, plist):
        with open(plist, 'wb') as fp:
            dump(path, fp)

def shellQuote(value):
    """
    Return the string value in a form that can safely be inserted into
    a shell command.
    """
    return "'%s'"%(value.replace("'", "'\"'\"'"))

def grepValue(fn, variable):
    """
    Return the unquoted value of a variable from a file..
    QUOTED_VALUE='quotes'    -> str('quotes')
    UNQUOTED_VALUE=noquotes  -> str('noquotes')
    """
    variable = variable + '='
    for ln in open(fn, 'r'):
        if ln.startswith(variable):
            value = ln[len(variable):].strip()
            return value.strip("\"'")
    raise RuntimeError("Cannot find variable %s" % variable[:-1])

_cache_getVersion = None

def getVersion():
    global _cache_getVersion
    if _cache_getVersion is None:
        _cache_getVersion = grepValue(
            os.path.join(SRCDIR, 'configure'), 'PACKAGE_VERSION')
    return _cache_getVersion

def getVersionMajorMinor():
    return tuple([int(n) for n in getVersion().split('.', 2)])

_cache_getFullVersion = None

def getFullVersion():
    global _cache_getFullVersion
    if _cache_getFullVersion is not None:
        return _cache_getFullVersion
    fn = os.path.join(SRCDIR, 'Include', 'patchlevel.h')
    for ln in open(fn):
        if 'PY_VERSION' in ln:
            _cache_getFullVersion = ln.split()[-1][1:-1]
            return _cache_getFullVersion
    raise RuntimeError("Cannot find full version??")

FW_PREFIX = ["Library", "Frameworks", "Python.framework"]
FW_VERSION_PREFIX = "--undefined--" # initialized in parseOptions
FW_SSL_DIRECTORY = "--undefined--" # initialized in parseOptions

# The directory we'll use to create the build (will be erased and recreated)
WORKDIR = "/tmp/_py"

# The directory we'll use to store third-party sources. Set this to something
# else if you don't want to re-fetch required libraries every time.
DEPSRC = os.path.join(WORKDIR, 'third-party')
DEPSRC = os.path.expanduser('~/Universal/other-sources')

universal_opts_map = { 'universal2': ('arm64', 'x86_64'),
                       '32-bit': ('i386', 'ppc',),
                       '64-bit': ('x86_64', 'ppc64',),
                       'intel':  ('i386', 'x86_64'),
                       'intel-32':  ('i386',),
                       'intel-64':  ('x86_64',),
                       '3-way':  ('ppc', 'i386', 'x86_64'),
                       'all':    ('i386', 'ppc', 'x86_64', 'ppc64',) }
default_target_map = {
        'universal2': '10.9',
        '64-bit': '10.5',
        '3-way': '10.5',
        'intel': '10.5',
        'intel-32': '10.4',
        'intel-64': '10.5',
        'all': '10.5',
}

UNIVERSALOPTS = tuple(universal_opts_map.keys())

UNIVERSALARCHS = '32-bit'

ARCHLIST = universal_opts_map[UNIVERSALARCHS]

# Source directory (assume we're in Mac/BuildScript)
SRCDIR = os.path.dirname(
        os.path.dirname(
            os.path.dirname(
                os.path.abspath(__file__
        ))))

# $MACOSX_DEPLOYMENT_TARGET -> minimum OS X level
DEPTARGET = '10.5'

def getDeptargetTuple():
    return tuple([int(n) for n in DEPTARGET.split('.')[0:2]])

def getBuildTuple():
    return tuple([int(n) for n in platform.mac_ver()[0].split('.')[0:2]])

def getTargetCompilers():
    target_cc_map = {
        '10.4': ('gcc-4.0', 'g++-4.0'),
        '10.5': ('gcc', 'g++'),
        '10.6': ('gcc', 'g++'),
    }
    return target_cc_map.get(DEPTARGET, ('gcc', 'g++') )

CC, CXX = getTargetCompilers()

PYTHON_3 = getVersionMajorMinor() >= (3, 0)

USAGE = textwrap.dedent("""\
    Usage: build_python [options]

    Options:
    -? or -h:            Show this message
    -b DIR
    --build-dir=DIR:     Create build here (default: %(WORKDIR)r)
    --third-party=DIR:   Store third-party sources here (default: %(DEPSRC)r)
    --sdk-path=DIR:      Location of the SDK (deprecated, use SDKROOT env variable)
    --src-dir=DIR:       Location of the Python sources (default: %(SRCDIR)r)
    --dep-target=10.n    macOS deployment target (default: %(DEPTARGET)r)
    --universal-archs=x  universal architectures (options: %(UNIVERSALOPTS)r, default: %(UNIVERSALARCHS)r)
""")% globals()

# Dict of object file names with shared library names to check after building.
# This is to ensure that we ended up dynamically linking with the shared
# library paths and versions we expected.  For example:
#   EXPECTED_SHARED_LIBS['_tkinter.so'] = [
#                       '/Library/Frameworks/Tcl.framework/Versions/8.5/Tcl',
#                       '/Library/Frameworks/Tk.framework/Versions/8.5/Tk']
EXPECTED_SHARED_LIBS = {}

# Are we building and linking with our own copy of Tcl/TK?
#   For now, do so if deployment target is 10.6+.
def internalTk():
    return getDeptargetTuple() >= (10, 6)

# Do we use 8.6.8 when building our own copy
# of Tcl/Tk or a modern version.
#   We use the old version when buildin on
#   old versions of macOS due to build issues.
def useOldTk():
    return getBuildTuple() < (10, 15)


def tweak_tcl_build(basedir, archList):
    with open("Makefile", "r") as fp:
        contents = fp.readlines()

    # For reasons I don't understand the tcl configure script
    # decides that some stdlib symbols aren't present, before
    # deciding that strtod is broken.
    new_contents = []
    for line in contents:
        if line.startswith("COMPAT_OBJS"):
            # note: the space before strtod.o is intentional,
            # the detection of a broken strtod results in
            # "fixstrod.o" on this line.
            for nm in ("strstr.o", "strtoul.o", " strtod.o"):
                line = line.replace(nm, "")
        new_contents.append(line)

    with open("Makefile", "w") as fp:
        fp.writelines(new_contents)

# List of names of third party software built with this installer.
# The names will be inserted into the rtf version of the License.
THIRD_PARTY_LIBS = []

# Instructions for building libraries that are necessary for building a
# batteries included python.
#   [The recipes are defined here for convenience but instantiated later after
#    command line options have been processed.]
def library_recipes():
    result = []

    LT_10_5 = bool(getDeptargetTuple() < (10, 5))

    # Since Apple removed the header files for the deprecated system
    # OpenSSL as of the Xcode 7 release (for OS X 10.10+), we do not
    # have much choice but to build our own copy here, too.

    result.extend([
          dict(
              name="OpenSSL 1.1.1i",
              url="https://www.openssl.org/source/openssl-1.1.1i.tar.gz",
              checksum='08987c3cf125202e2b0840035efb392c',
              buildrecipe=build_universal_openssl,
              configure=None,
              install=None,
          ),
    ])

    if internalTk():
        if useOldTk():
            tcl_tk_ver='8.6.8'
            tcl_checksum='81656d3367af032e0ae6157eff134f89'

            tk_checksum='5e0faecba458ee1386078fb228d008ba'
            tk_patches = ['tk868_on_10_8_10_9.patch']

        else:
            tcl_tk_ver='8.6.11'
            tcl_checksum='8a4c004f48984a03a7747e9ba06e4da4'

            tk_checksum='c7ee71a2d05bba78dfffd76528dc17c6'
            tk_patches = [ ]


        result.extend([
          dict(
              name="Tcl %s"%(tcl_tk_ver,),
              url="ftp://ftp.tcl.tk/pub/tcl//tcl8_6/tcl%s-src.tar.gz"%(tcl_tk_ver,),
              checksum=tcl_checksum,
              buildDir="unix",
              configure_pre=[
                    '--enable-shared',
                    '--enable-threads',
                    '--libdir=/Library/Frameworks/Python.framework/Versions/%s/lib'%(getVersion(),),
              ],
              useLDFlags=False,
              buildrecipe=tweak_tcl_build,
              install='make TCL_LIBRARY=%(TCL_LIBRARY)s && make install TCL_LIBRARY=%(TCL_LIBRARY)s DESTDIR=%(DESTDIR)s'%{
                  "DESTDIR": shellQuote(os.path.join(WORKDIR, 'libraries')),
                  "TCL_LIBRARY": shellQuote('/Library/Frameworks/Python.framework/Versions/%s/lib/tcl8.6'%(getVersion())),
                  },
              ),
          dict(
              name="Tk %s"%(tcl_tk_ver,),
              url="ftp://ftp.tcl.tk/pub/tcl//tcl8_6/tk%s-src.tar.gz"%(tcl_tk_ver,),
              checksum=tk_checksum,
              patches=tk_patches,
              buildDir="unix",
              configure_pre=[
                    '--enable-aqua',
                    '--enable-shared',
                    '--enable-threads',
                    '--libdir=/Library/Frameworks/Python.framework/Versions/%s/lib'%(getVersion(),),
              ],
              useLDFlags=False,
              install='make TCL_LIBRARY=%(TCL_LIBRARY)s TK_LIBRARY=%(TK_LIBRARY)s && make install TCL_LIBRARY=%(TCL_LIBRARY)s TK_LIBRARY=%(TK_LIBRARY)s DESTDIR=%(DESTDIR)s'%{
                  "DESTDIR": shellQuote(os.path.join(WORKDIR, 'libraries')),
                  "TCL_LIBRARY": shellQuote('/Library/Frameworks/Python.framework/Versions/%s/lib/tcl8.6'%(getVersion())),
                  "TK_LIBRARY": shellQuote('/Library/Frameworks/Python.framework/Versions/%s/lib/tk8.6'%(getVersion())),
                  },
                ),
        ])

    if PYTHON_3:
        result.extend([
          dict(
              name="XZ 5.2.3",
              url="http://tukaani.org/xz/xz-5.2.3.tar.gz",
              checksum='ef68674fb47a8b8e741b34e429d86e9d',
              configure_pre=[
                    '--disable-dependency-tracking',
              ]
              ),
        ])

    result.extend([
          dict(
              name="NCurses 5.9",
              url="http://ftp.gnu.org/pub/gnu/ncurses/ncurses-5.9.tar.gz",
              checksum='8cb9c412e5f2d96bc6f459aa8c6282a1',
              configure_pre=[
                  "--enable-widec",
                  "--without-cxx",
                  "--without-cxx-binding",
                  "--without-ada",
                  "--without-curses-h",
                  "--enable-shared",
                  "--with-shared",
                  "--without-debug",
                  "--without-normal",
                  "--without-tests",
                  "--without-manpages",
                  "--datadir=/usr/share",
                  "--sysconfdir=/etc",
                  "--sharedstatedir=/usr/com",
                  "--with-terminfo-dirs=/usr/share/terminfo",
                  "--with-default-terminfo-dir=/usr/share/terminfo",
                  "--libdir=/Library/Frameworks/Python.framework/Versions/%s/lib"%(getVersion(),),
              ],
              patchscripts=[
                  ("ftp://ftp.invisible-island.net/ncurses//5.9/ncurses-5.9-20120616-patch.sh.bz2",
                   "f54bf02a349f96a7c4f0d00922f3a0d4"),
                   ],
              useLDFlags=False,
              install='make && make install DESTDIR=%s && cd %s/usr/local/lib && ln -fs ../../../Library/Frameworks/Python.framework/Versions/%s/lib/lib* .'%(
                  shellQuote(os.path.join(WORKDIR, 'libraries')),
                  shellQuote(os.path.join(WORKDIR, 'libraries')),
                  getVersion(),
                  ),
          ),
          dict(
              name="SQLite 3.34.0",
              url="https://sqlite.org/2020/sqlite-autoconf-3340000.tar.gz",
              checksum='7f33c9db7b713957fcb9271fe9049fef',
              extra_cflags=('-Os '
                            '-DSQLITE_ENABLE_FTS5 '
                            '-DSQLITE_ENABLE_FTS4 '
                            '-DSQLITE_ENABLE_FTS3_PARENTHESIS '
                            '-DSQLITE_ENABLE_JSON1 '
                            '-DSQLITE_ENABLE_RTREE '
                            '-DSQLITE_TCL=0 '
                 '%s' % ('','-DSQLITE_WITHOUT_ZONEMALLOC ')[LT_10_5]),
              configure_pre=[
                  '--enable-threadsafe',
                  '--enable-shared=no',
                  '--enable-static=yes',
                  '--disable-readline',
                  '--disable-dependency-tracking',
              ]
          ),
        ])

    if getDeptargetTuple() < (10, 5):
        result.extend([
          dict(
              name="Bzip2 1.0.6",
              url="http://bzip.org/1.0.6/bzip2-1.0.6.tar.gz",
              checksum='00b516f4704d4a7cb50a1d97e6e8e15b',
              configure=None,
              install='make install CC=%s CXX=%s, PREFIX=%s/usr/local/ CFLAGS="-arch %s"'%(
                  CC, CXX,
                  shellQuote(os.path.join(WORKDIR, 'libraries')),
                  ' -arch '.join(ARCHLIST),
              ),
          ),
          dict(
              name="ZLib 1.2.3",
              url="http://www.gzip.org/zlib/zlib-1.2.3.tar.gz",
              checksum='debc62758716a169df9f62e6ab2bc634',
              configure=None,
              install='make install CC=%s CXX=%s, prefix=%s/usr/local/ CFLAGS="-arch %s"'%(
                  CC, CXX,
                  shellQuote(os.path.join(WORKDIR, 'libraries')),
                  ' -arch '.join(ARCHLIST),
              ),
          ),
          dict(
              # Note that GNU readline is GPL'd software
              name="GNU Readline 6.1.2",
              url="http://ftp.gnu.org/pub/gnu/readline/readline-6.1.tar.gz" ,
              checksum='fc2f7e714fe792db1ce6ddc4c9fb4ef3',
              patchlevel='0',
              patches=[
                  # The readline maintainers don't do actual micro releases, but
                  # just ship a set of patches.
                  ('http://ftp.gnu.org/pub/gnu/readline/readline-6.1-patches/readline61-001',
                   'c642f2e84d820884b0bf9fd176bc6c3f'),
                  ('http://ftp.gnu.org/pub/gnu/readline/readline-6.1-patches/readline61-002',
                   '1a76781a1ea734e831588285db7ec9b1'),
              ]
          ),
        ])

    if not PYTHON_3:
        result.extend([
          dict(
              name="Sleepycat DB 4.7.25",
              url="http://download.oracle.com/berkeley-db/db-4.7.25.tar.gz",
              checksum='ec2b87e833779681a0c3a814aa71359e',
              buildDir="build_unix",
              configure="../dist/configure",
              configure_pre=[
                  '--includedir=/usr/local/include/db4',
              ]
          ),
        ])

    return result


# Instructions for building packages inside the .mpkg.
def pkg_recipes():
    unselected_for_python3 = ('selected', 'unselected')[PYTHON_3]
    result = [
        dict(
            name="PythonFramework",
            long_name="Python Framework",
            source="/Library/Frameworks/Python.framework",
            readme="""\
                This package installs Python.framework, that is the python
                interpreter and the standard library.
            """,
            postflight="scripts/postflight.framework",
            selected='selected',
        ),
        dict(
            name="PythonApplications",
            long_name="GUI Applications",
            source="/Applications/Python %(VER)s",
            readme="""\
                This package installs IDLE (an interactive Python IDE),
                Python Launcher and Build Applet (create application bundles
                from python scripts).

                It also installs a number of examples and demos.
                """,
            required=False,
            selected='selected',
        ),
        dict(
            name="PythonUnixTools",
            long_name="UNIX command-line tools",
            source="/usr/local/bin",
            readme="""\
                This package installs the unix tools in /usr/local/bin for
                compatibility with older releases of Python. This package
                is not necessary to use Python.
                """,
            required=False,
            selected='selected',
        ),
        dict(
            name="PythonDocumentation",
            long_name="Python Documentation",
            topdir="/Library/Frameworks/Python.framework/Versions/%(VER)s/Resources/English.lproj/Documentation",
            source="/pydocs",
            readme="""\
                This package installs the python documentation at a location
                that is useable for pydoc and IDLE.
                """,
            postflight="scripts/postflight.documentation",
            required=False,
            selected='selected',
        ),
        dict(
            name="PythonProfileChanges",
            long_name="Shell profile updater",
            readme="""\
                This packages updates your shell profile to make sure that
                the Python tools are found by your shell in preference of
                the system provided Python tools.

                If you don't install this package you'll have to add
                "/Library/Frameworks/Python.framework/Versions/%(VER)s/bin"
                to your PATH by hand.
                """,
            postflight="scripts/postflight.patch-profile",
            topdir="/Library/Frameworks/Python.framework",
            source="/empty-dir",
            required=False,
            selected='selected',
        ),
        dict(
            name="PythonInstallPip",
            long_name="Install or upgrade pip",
            readme="""\
                This package installs (or upgrades from an earlier version)
                pip, a tool for installing and managing Python packages.
                """,
            postflight="scripts/postflight.ensurepip",
            topdir="/Library/Frameworks/Python.framework",
            source="/empty-dir",
            required=False,
            selected='selected',
        ),
    ]

    return result

def fatal(msg):
    """
    A fatal error, bail out.
    """
    sys.stderr.write('FATAL: ')
    sys.stderr.write(msg)
    sys.stderr.write('\n')
    sys.exit(1)

def fileContents(fn):
    """
    Return the contents of the named file
    """
    return open(fn, 'r').read()

def runCommand(commandline):
    """
    Run a command and raise RuntimeError if it fails. Output is suppressed
    unless the command fails.
    """
    fd = os.popen(commandline, 'r')
    data = fd.read()
    xit = fd.close()
    if xit is not None:
        sys.stdout.write(data)
        raise RuntimeError("command failed: %s"%(commandline,))

    if VERBOSE:
        sys.stdout.write(data); sys.stdout.flush()

def captureCommand(commandline):
    fd = os.popen(commandline, 'r')
    data = fd.read()
    xit = fd.close()
    if xit is not None:
        sys.stdout.write(data)
        raise RuntimeError("command failed: %s"%(commandline,))

    return data

def getTclTkVersion(configfile, versionline):
    """
    search Tcl or Tk configuration file for version line
    """
    try:
        f = open(configfile, "r")
    except OSError:
        fatal("Framework configuration file not found: %s" % configfile)

    for l in f:
        if l.startswith(versionline):
            f.close()
            return l

    fatal("Version variable %s not found in framework configuration file: %s"
            % (versionline, configfile))

def checkEnvironment():
    """
    Check that we're running on a supported system.
    """

    if sys.version_info[0:2] < (2, 7):
        fatal("This script must be run with Python 2.7 (or later)")

    if platform.system() != 'Darwin':
        fatal("This script should be run on a macOS 10.5 (or later) system")

    if int(platform.release().split('.')[0]) < 8:
        fatal("This script should be run on a macOS 10.5 (or later) system")

    # Because we only support dynamic load of only one major/minor version of
    # Tcl/Tk, if we are not using building and using our own private copy of
    # Tcl/Tk, ensure:
    # 1. there is a user-installed framework (usually ActiveTcl) in (or linked
    #       in) SDKROOT/Library/Frameworks.  As of Python 3.7.0, we no longer
    #       enforce that the version of the user-installed framework also
    #       exists in the system-supplied Tcl/Tk frameworks.  Time to support
    #       Tcl/Tk 8.6 even if Apple does not.
    if not internalTk():
        frameworks = {}
        for framework in ['Tcl', 'Tk']:
            fwpth = 'Library/Frameworks/%s.framework/Versions/Current' % framework
            libfw = os.path.join('/', fwpth)
            usrfw = os.path.join(os.getenv('HOME'), fwpth)
            frameworks[framework] = os.readlink(libfw)
            if not os.path.exists(libfw):
                fatal("Please install a link to a current %s %s as %s so "
                        "the user can override the system framework."
                        % (framework, frameworks[framework], libfw))
            if os.path.exists(usrfw):
                fatal("Please rename %s to avoid possible dynamic load issues."
                        % usrfw)

        if frameworks['Tcl'] != frameworks['Tk']:
            fatal("The Tcl and Tk frameworks are not the same version.")

        print(" -- Building with external Tcl/Tk %s frameworks"
                    % frameworks['Tk'])

        # add files to check after build
        EXPECTED_SHARED_LIBS['_tkinter.so'] = [
                "/Library/Frameworks/Tcl.framework/Versions/%s/Tcl"
                    % frameworks['Tcl'],
                "/Library/Frameworks/Tk.framework/Versions/%s/Tk"
                    % frameworks['Tk'],
                ]
    else:
        print(" -- Building private copy of Tcl/Tk")
    print("")

    # Remove inherited environment variables which might influence build
    environ_var_prefixes = ['CPATH', 'C_INCLUDE_', 'DYLD_', 'LANG', 'LC_',
                            'LD_', 'LIBRARY_', 'PATH', 'PYTHON']
    for ev in list(os.environ):
        for prefix in environ_var_prefixes:
            if ev.startswith(prefix) :
                print("INFO: deleting environment variable %s=%s" % (
                                                    ev, os.environ[ev]))
                del os.environ[ev]

    base_path = '/bin:/sbin:/usr/bin:/usr/sbin'
    if 'SDK_TOOLS_BIN' in os.environ:
        base_path = os.environ['SDK_TOOLS_BIN'] + ':' + base_path
    # Xcode 2.5 on OS X 10.4 does not include SetFile in its usr/bin;
    # add its fixed location here if it exists
    OLD_DEVELOPER_TOOLS = '/Developer/Tools'
    if os.path.isdir(OLD_DEVELOPER_TOOLS):
        base_path = base_path + ':' + OLD_DEVELOPER_TOOLS
    os.environ['PATH'] = base_path
    print("Setting default PATH: %s"%(os.environ['PATH']))

def parseOptions(args=None):
    """
    Parse arguments and update global settings.
    """
    global WORKDIR, DEPSRC, SRCDIR, DEPTARGET
    global UNIVERSALOPTS, UNIVERSALARCHS, ARCHLIST, CC, CXX
    global FW_VERSION_PREFIX
    global FW_SSL_DIRECTORY

    if args is None:
        args = sys.argv[1:]

    try:
        options, args = getopt.getopt(args, '?hb',
                [ 'build-dir=', 'third-party=', 'sdk-path=' , 'src-dir=',
                  'dep-target=', 'universal-archs=', 'help' ])
    except getopt.GetoptError:
        print(sys.exc_info()[1])
        sys.exit(1)

    if args:
        print("Additional arguments")
        sys.exit(1)

    deptarget = None
    for k, v in options:
        if k in ('-h', '-?', '--help'):
            print(USAGE)
            sys.exit(0)

        elif k in ('-d', '--build-dir'):
            WORKDIR=v

        elif k in ('--third-party',):
            DEPSRC=v

        elif k in ('--sdk-path',):
            print(" WARNING: --sdk-path is no longer supported")

        elif k in ('--src-dir',):
            SRCDIR=v

        elif k in ('--dep-target', ):
            DEPTARGET=v
            deptarget=v

        elif k in ('--universal-archs', ):
            if v in UNIVERSALOPTS:
                UNIVERSALARCHS = v
                ARCHLIST = universal_opts_map[UNIVERSALARCHS]
                if deptarget is None:
                    # Select alternate default deployment
                    # target
                    DEPTARGET = default_target_map.get(v, '10.5')
            else:
                raise NotImplementedError(v)

        else:
            raise NotImplementedError(k)

    SRCDIR=os.path.abspath(SRCDIR)
    WORKDIR=os.path.abspath(WORKDIR)
    DEPSRC=os.path.abspath(DEPSRC)

    CC, CXX = getTargetCompilers()

    FW_VERSION_PREFIX = FW_PREFIX[:] + ["Versions", getVersion()]
    FW_SSL_DIRECTORY = FW_VERSION_PREFIX[:] + ["etc", "openssl"]

    print("-- Settings:")
    print("   * Source directory:    %s" % SRCDIR)
    print("   * Build directory:     %s" % WORKDIR)
    print("   * Third-party source:  %s" % DEPSRC)
    print("   * Deployment target:   %s" % DEPTARGET)
    print("   * Universal archs:     %s" % str(ARCHLIST))
    print("   * C compiler:          %s" % CC)
    print("   * C++ compiler:        %s" % CXX)
    print("")
    print(" -- Building a Python %s framework at patch level %s"
                % (getVersion(), getFullVersion()))
    print("")

def extractArchive(builddir, archiveName):
    """
    Extract a source archive into 'builddir'. Returns the path of the
    extracted archive.

    XXX: This function assumes that archives contain a toplevel directory
    that is has the same name as the basename of the archive. This is
    safe enough for almost anything we use.  Unfortunately, it does not
    work for current Tcl and Tk source releases where the basename of
    the archive ends with "-src" but the uncompressed directory does not.
    For now, just special case Tcl and Tk tar.gz downloads.
    """
    curdir = os.getcwd()
    try:
        os.chdir(builddir)
        if archiveName.endswith('.tar.gz'):
            retval = os.path.basename(archiveName[:-7])
            if ((retval.startswith('tcl') or retval.startswith('tk'))
                    and retval.endswith('-src')):
                retval = retval[:-4]
            if os.path.exists(retval):
                shutil.rmtree(retval)
            fp = os.popen("tar zxf %s 2>&1"%(shellQuote(archiveName),), 'r')

        elif archiveName.endswith('.tar.bz2'):
            retval = os.path.basename(archiveName[:-8])
            if os.path.exists(retval):
                shutil.rmtree(retval)
            fp = os.popen("tar jxf %s 2>&1"%(shellQuote(archiveName),), 'r')

        elif archiveName.endswith('.tar'):
            retval = os.path.basename(archiveName[:-4])
            if os.path.exists(retval):
                shutil.rmtree(retval)
            fp = os.popen("tar xf %s 2>&1"%(shellQuote(archiveName),), 'r')

        elif archiveName.endswith('.zip'):
            retval = os.path.basename(archiveName[:-4])
            if os.path.exists(retval):
                shutil.rmtree(retval)
            fp = os.popen("unzip %s 2>&1"%(shellQuote(archiveName),), 'r')

        data = fp.read()
        xit = fp.close()
        if xit is not None:
            sys.stdout.write(data)
            raise RuntimeError("Cannot extract %s"%(archiveName,))

        return os.path.join(builddir, retval)

    finally:
        os.chdir(curdir)

def downloadURL(url, fname):
    """
    Download the contents of the url into the file.
    """
    fpIn = urllib_request.urlopen(url)
    fpOut = open(fname, 'wb')
    block = fpIn.read(10240)
    try:
        while block:
            fpOut.write(block)
            block = fpIn.read(10240)
        fpIn.close()
        fpOut.close()
    except:
        try:
            os.unlink(fname)
        except OSError:
            pass

def verifyThirdPartyFile(url, checksum, fname):
    """
    Download file from url to filename fname if it does not already exist.
    Abort if file contents does not match supplied md5 checksum.
    """
    name = os.path.basename(fname)
    if os.path.exists(fname):
        print("Using local copy of %s"%(name,))
    else:
        print("Did not find local copy of %s"%(name,))
        print("Downloading %s"%(name,))
        downloadURL(url, fname)
        print("Archive for %s stored as %s"%(name, fname))
    if os.system(
            'MD5=$(openssl md5 %s) ; test "${MD5##*= }" = "%s"'
                % (shellQuote(fname), checksum) ):
        fatal('MD5 checksum mismatch for file %s' % fname)

def build_universal_openssl(basedir, archList):
    """
    Special case build recipe for universal build of openssl.

    The upstream OpenSSL build system does not directly support
    OS X universal builds.  We need to build each architecture
    separately then lipo them together into fat libraries.
    """

    # OpenSSL fails to build with Xcode 2.5 (on OS X 10.4).
    # If we are building on a 10.4.x or earlier system,
    # unilaterally disable assembly code building to avoid the problem.
    no_asm = int(platform.release().split(".")[0]) < 9

    def build_openssl_arch(archbase, arch):
        "Build one architecture of openssl"
        arch_opts = {
            "i386": ["darwin-i386-cc"],
            "x86_64": ["darwin64-x86_64-cc", "enable-ec_nistp_64_gcc_128"],
            "arm64": ["darwin64-arm64-cc"],
            "ppc": ["darwin-ppc-cc"],
            "ppc64": ["darwin64-ppc-cc"],
        }

        # Somewhere between OpenSSL 1.1.0j and 1.1.1c, changes cause the
        # "enable-ec_nistp_64_gcc_128" option to get compile errors when
        # building on our 10.6 gcc-4.2 environment.  There have been other
        # reports of projects running into this when using older compilers.
        # So, for now, do not try to use "enable-ec_nistp_64_gcc_128" when
        # building for 10.6.
        if getDeptargetTuple() == (10, 6):
            arch_opts['x86_64'].remove('enable-ec_nistp_64_gcc_128')

        configure_opts = [
            "no-idea",
            "no-mdc2",
            "no-rc5",
            "no-zlib",
            "no-ssl3",
            # "enable-unit-test",
            "shared",
            "--prefix=%s"%os.path.join("/", *FW_VERSION_PREFIX),
            "--openssldir=%s"%os.path.join("/", *FW_SSL_DIRECTORY),
        ]
        if no_asm:
            configure_opts.append("no-asm")
        runCommand(" ".join(["perl", "Configure"]
                        + arch_opts[arch] + configure_opts))
        runCommand("make depend")
        runCommand("make all")
        runCommand("make install_sw DESTDIR=%s"%shellQuote(archbase))
        # runCommand("make test")
        return

    srcdir = os.getcwd()
    universalbase = os.path.join(srcdir, "..",
                        os.path.basename(srcdir) + "-universal")
    os.mkdir(universalbase)
    archbasefws = []
    for arch in archList:
        # fresh copy of the source tree
        archsrc = os.path.join(universalbase, arch, "src")
        shutil.copytree(srcdir, archsrc, symlinks=True)
        # install base for this arch
        archbase = os.path.join(universalbase, arch, "root")
        os.mkdir(archbase)
        # Python framework base within install_prefix:
        # the build will install into this framework..
        # This is to ensure that the resulting shared libs have
        # the desired real install paths built into them.
        archbasefw = os.path.join(archbase, *FW_VERSION_PREFIX)

        # build one architecture
        os.chdir(archsrc)
        build_openssl_arch(archbase, arch)
        os.chdir(srcdir)
        archbasefws.append(archbasefw)

    # copy arch-independent files from last build into the basedir framework
    basefw = os.path.join(basedir, *FW_VERSION_PREFIX)
    shutil.copytree(
            os.path.join(archbasefw, "include", "openssl"),
            os.path.join(basefw, "include", "openssl")
            )

    shlib_version_number = grepValue(os.path.join(archsrc, "Makefile"),
            "SHLIB_VERSION_NUMBER")
    #   e.g. -> "1.0.0"
    libcrypto = "libcrypto.dylib"
    libcrypto_versioned = libcrypto.replace(".", "."+shlib_version_number+".")
    #   e.g. -> "libcrypto.1.0.0.dylib"
    libssl = "libssl.dylib"
    libssl_versioned = libssl.replace(".", "."+shlib_version_number+".")
    #   e.g. -> "libssl.1.0.0.dylib"

    try:
        os.mkdir(os.path.join(basefw, "lib"))
    except OSError:
        pass

    # merge the individual arch-dependent shared libs into a fat shared lib
    archbasefws.insert(0, basefw)
    for (lib_unversioned, lib_versioned) in [
                (libcrypto, libcrypto_versioned),
                (libssl, libssl_versioned)
            ]:
        runCommand("lipo -create -output " +
                    " ".join(shellQuote(
                            os.path.join(fw, "lib", lib_versioned))
                                    for fw in archbasefws))
        # and create an unversioned symlink of it
        os.symlink(lib_versioned, os.path.join(basefw, "lib", lib_unversioned))

    # Create links in the temp include and lib dirs that will be injected
    # into the Python build so that setup.py can find them while building
    # and the versioned links so that the setup.py post-build import test
    # does not fail.
    relative_path = os.path.join("..", "..", "..", *FW_VERSION_PREFIX)
    for fn in [
            ["include", "openssl"],
            ["lib", libcrypto],
            ["lib", libssl],
            ["lib", libcrypto_versioned],
            ["lib", libssl_versioned],
        ]:
        os.symlink(
            os.path.join(relative_path, *fn),
            os.path.join(basedir, "usr", "local", *fn)
        )

    return

def buildRecipe(recipe, basedir, archList):
    """
    Build software using a recipe. This function does the
    'configure;make;make install' dance for C software, with a possibility
    to customize this process, basically a poor-mans DarwinPorts.
    """
    curdir = os.getcwd()

    name = recipe['name']
    THIRD_PARTY_LIBS.append(name)
    url = recipe['url']
    configure = recipe.get('configure', './configure')
    buildrecipe = recipe.get('buildrecipe', None)
    install = recipe.get('install', 'make && make install DESTDIR=%s'%(
        shellQuote(basedir)))

    archiveName = os.path.split(url)[-1]
    sourceArchive = os.path.join(DEPSRC, archiveName)

    if not os.path.exists(DEPSRC):
        os.mkdir(DEPSRC)

    verifyThirdPartyFile(url, recipe['checksum'], sourceArchive)
    print("Extracting archive for %s"%(name,))
    buildDir=os.path.join(WORKDIR, '_bld')
    if not os.path.exists(buildDir):
        os.mkdir(buildDir)

    workDir = extractArchive(buildDir, sourceArchive)
    os.chdir(workDir)

    for patch in recipe.get('patches', ()):
        if isinstance(patch, tuple):
            url, checksum = patch
            fn = os.path.join(DEPSRC, os.path.basename(url))
            verifyThirdPartyFile(url, checksum, fn)
        else:
            # patch is a file in the source directory
            fn = os.path.join(curdir, patch)
        runCommand('patch -p%s < %s'%(recipe.get('patchlevel', 1),
            shellQuote(fn),))

    for patchscript in recipe.get('patchscripts', ()):
        if isinstance(patchscript, tuple):
            url, checksum = patchscript
            fn = os.path.join(DEPSRC, os.path.basename(url))
            verifyThirdPartyFile(url, checksum, fn)
        else:
            # patch is a file in the source directory
            fn = os.path.join(curdir, patchscript)
        if fn.endswith('.bz2'):
            runCommand('bunzip2 -fk %s' % shellQuote(fn))
            fn = fn[:-4]
        runCommand('sh %s' % shellQuote(fn))
        os.unlink(fn)

    if 'buildDir' in recipe:
        os.chdir(recipe['buildDir'])

    if configure is not None:
        configure_args = [
            "--prefix=/usr/local",
            "--enable-static",
            "--disable-shared",
            #"CPP=gcc -arch %s -E"%(' -arch '.join(archList,),),
        ]

        if 'configure_pre' in recipe:
            args = list(recipe['configure_pre'])
            if '--disable-static' in args:
                configure_args.remove('--enable-static')
            if '--enable-shared' in args:
                configure_args.remove('--disable-shared')
            configure_args.extend(args)

        if recipe.get('useLDFlags', 1):
            configure_args.extend([
                "CFLAGS=%s-mmacosx-version-min=%s -arch %s "
                            "-I%s/usr/local/include"%(
                        recipe.get('extra_cflags', ''),
                        DEPTARGET,
                        ' -arch '.join(archList),
                        shellQuote(basedir)[1:-1],),
                "LDFLAGS=-mmacosx-version-min=%s -L%s/usr/local/lib -arch %s"%(
                    DEPTARGET,
                    shellQuote(basedir)[1:-1],
                    ' -arch '.join(archList)),
            ])
        else:
            configure_args.extend([
                "CFLAGS=%s-mmacosx-version-min=%s -arch %s "
                            "-I%s/usr/local/include"%(
                        recipe.get('extra_cflags', ''),
                        DEPTARGET,
                        ' -arch '.join(archList),
                        shellQuote(basedir)[1:-1],),
            ])

        if 'configure_post' in recipe:
            configure_args = configure_args + list(recipe['configure_post'])

        configure_args.insert(0, configure)
        configure_args = [ shellQuote(a) for a in configure_args ]

        print("Running configure for %s"%(name,))
        runCommand(' '.join(configure_args) + ' 2>&1')

    if buildrecipe is not None:
        # call special-case build recipe, e.g. for openssl
        buildrecipe(basedir, archList)

    if install is not None:
        print("Running install for %s"%(name,))
        runCommand('{ ' + install + ' ;} 2>&1')

    print("Done %s"%(name,))
    print("")

    os.chdir(curdir)

def buildLibraries():
    """
    Build our dependencies into $WORKDIR/libraries/usr/local
    """
    print("")
    print("Building required libraries")
    print("")
    universal = os.path.join(WORKDIR, 'libraries')
    os.mkdir(universal)
    os.makedirs(os.path.join(universal, 'usr', 'local', 'lib'))
    os.makedirs(os.path.join(universal, 'usr', 'local', 'include'))

    for recipe in library_recipes():
        buildRecipe(recipe, universal, ARCHLIST)



def buildPythonDocs():
    # This stores the documentation as Resources/English.lproj/Documentation
    # inside the framework. pydoc and IDLE will pick it up there.
    print("Install python documentation")
    rootDir = os.path.join(WORKDIR, '_root')
    buildDir = os.path.join('../../Doc')
    docdir = os.path.join(rootDir, 'pydocs')
    curDir = os.getcwd()
    os.chdir(buildDir)
    runCommand('make clean')

    # Search third-party source directory for a pre-built version of the docs.
    #   Use the naming convention of the docs.python.org html downloads:
    #       python-3.9.0b1-docs-html.tar.bz2
    doctarfiles = [ f for f in os.listdir(DEPSRC)
        if f.startswith('python-'+getFullVersion())
        if f.endswith('-docs-html.tar.bz2') ]
    if doctarfiles:
        doctarfile = doctarfiles[0]
        if not os.path.exists('build'):
            os.mkdir('build')
        # if build directory existed, it was emptied by make clean, above
        os.chdir('build')
        # Extract the first archive found for this version into build
        runCommand('tar xjf %s'%shellQuote(os.path.join(DEPSRC, doctarfile)))
        # see if tar extracted a directory ending in -docs-html
        archivefiles = [ f for f in os.listdir('.')
            if f.endswith('-docs-html')
            if os.path.isdir(f) ]
        if archivefiles:
            archivefile = archivefiles[0]
            # make it our 'Docs/build/html' directory
            print(' -- using pre-built python documentation from %s'%archivefile)
            os.rename(archivefile, 'html')
        os.chdir(buildDir)

    htmlDir = os.path.join('build', 'html')
    if not os.path.exists(htmlDir):
        # Create virtual environment for docs builds with blurb and sphinx
        runCommand('make venv')
        runCommand('make html PYTHON=venv/bin/python')
    os.rename(htmlDir, docdir)
    os.chdir(curDir)


def buildPython():
    print("Building a universal python for %s architectures" % UNIVERSALARCHS)

    buildDir = os.path.join(WORKDIR, '_bld', 'python')
    rootDir = os.path.join(WORKDIR, '_root')

    if os.path.exists(buildDir):
        shutil.rmtree(buildDir)
    if os.path.exists(rootDir):
        shutil.rmtree(rootDir)
    os.makedirs(buildDir)
    os.makedirs(rootDir)
    os.makedirs(os.path.join(rootDir, 'empty-dir'))
    curdir = os.getcwd()
    os.chdir(buildDir)

    # Extract the version from the configure file, needed to calculate
    # several paths.
    version = getVersion()

    # Since the extra libs are not in their installed framework location
    # during the build, augment the library path so that the interpreter
    # will find them during its extension import sanity checks.

    print("Running configure...")
    runCommand("%s -C --enable-framework --enable-universalsdk=/ "
               "--with-universal-archs=%s "
               "%s "
               "%s "
               "%s "
               "%s "
               "%s "
               "LDFLAGS='-g -L%s/libraries/usr/local/lib' "
               "CFLAGS='-g -I%s/libraries/usr/local/include' 2>&1"%(
        shellQuote(os.path.join(SRCDIR, 'configure')),
        UNIVERSALARCHS,
        (' ', '--with-computed-gotos ')[PYTHON_3],
        (' ', '--without-ensurepip ')[PYTHON_3],
        (' ', "--with-openssl='%s/libraries/usr/local'"%(
                            shellQuote(WORKDIR)[1:-1],))[PYTHON_3],
        (' ', "--with-tcltk-includes='-I%s/libraries/usr/local/include'"%(
                            shellQuote(WORKDIR)[1:-1],))[internalTk()],
        (' ', "--with-tcltk-libs='-L%s/libraries/usr/local/lib -ltcl8.6 -ltk8.6'"%(
                            shellQuote(WORKDIR)[1:-1],))[internalTk()],
        shellQuote(WORKDIR)[1:-1],
        shellQuote(WORKDIR)[1:-1]))

    # As of macOS 10.11 with SYSTEM INTEGRITY PROTECTION, DYLD_*
    # environment variables are no longer automatically inherited
    # by child processes from their parents. We used to just set
    # DYLD_LIBRARY_PATH, pointing to the third-party libs,
    # in build-installer.py's process environment and it was
    # passed through the make utility into the environment of
    # setup.py. Instead, we now append DYLD_LIBRARY_PATH to
    # the existing RUNSHARED configuration value when we call
    # make for extension module builds.

    runshared_for_make = "".join([
            " RUNSHARED=",
            "'",
            grepValue("Makefile", "RUNSHARED"),
            ' DYLD_LIBRARY_PATH=',
            os.path.join(WORKDIR, 'libraries', 'usr', 'local', 'lib'),
            "'" ])

    # Look for environment value BUILDINSTALLER_BUILDPYTHON_MAKE_EXTRAS
    # and, if defined, append its value to the make command.  This allows
    # us to pass in version control tags, like GITTAG, to a build from a
    # tarball rather than from a vcs checkout, thus eliminating the need
    # to have a working copy of the vcs program on the build machine.
    #
    # A typical use might be:
    #      export BUILDINSTALLER_BUILDPYTHON_MAKE_EXTRAS=" \
    #                         GITVERSION='echo 123456789a' \
    #                         GITTAG='echo v3.6.0' \
    #                         GITBRANCH='echo 3.6'"

    make_extras = os.getenv("BUILDINSTALLER_BUILDPYTHON_MAKE_EXTRAS")
    if make_extras:
        make_cmd = "make " + make_extras + runshared_for_make
    else:
        make_cmd = "make" + runshared_for_make
    print("Running " + make_cmd)
    runCommand(make_cmd)

    make_cmd = "make install DESTDIR=%s %s"%(
        shellQuote(rootDir),
        runshared_for_make)
    print("Running " + make_cmd)
    runCommand(make_cmd)

    make_cmd = "make frameworkinstallextras DESTDIR=%s %s"%(
        shellQuote(rootDir),
        runshared_for_make)
    print("Running " + make_cmd)
    runCommand(make_cmd)

    print("Copying required shared libraries")
    if os.path.exists(os.path.join(WORKDIR, 'libraries', 'Library')):
        build_lib_dir = os.path.join(
                WORKDIR, 'libraries', 'Library', 'Frameworks',
                'Python.framework', 'Versions', getVersion(), 'lib')
        fw_lib_dir = os.path.join(
                WORKDIR, '_root', 'Library', 'Frameworks',
                'Python.framework', 'Versions', getVersion(), 'lib')
        if internalTk():
            # move Tcl and Tk pkgconfig files
            runCommand("mv %s/pkgconfig/* %s/pkgconfig"%(
                        shellQuote(build_lib_dir),
                        shellQuote(fw_lib_dir) ))
            runCommand("rm -r %s/pkgconfig"%(
                        shellQuote(build_lib_dir), ))
        runCommand("mv %s/* %s"%(
                    shellQuote(build_lib_dir),
                    shellQuote(fw_lib_dir) ))

    frmDir = os.path.join(rootDir, 'Library', 'Frameworks', 'Python.framework')
    frmDirVersioned = os.path.join(frmDir, 'Versions', version)
    path_to_lib = os.path.join(frmDirVersioned, 'lib', 'python%s'%(version,))
    # create directory for OpenSSL certificates
    sslDir = os.path.join(frmDirVersioned, 'etc', 'openssl')
    os.makedirs(sslDir)

    print("Fix file modes")
    gid = grp.getgrnam('admin').gr_gid

    shared_lib_error = False
    for dirpath, dirnames, filenames in os.walk(frmDir):
        for dn in dirnames:
            os.chmod(os.path.join(dirpath, dn), STAT_0o775)
            os.chown(os.path.join(dirpath, dn), -1, gid)

        for fn in filenames:
            if os.path.islink(fn):
                continue

            # "chmod g+w $fn"
            p = os.path.join(dirpath, fn)
            st = os.stat(p)
            os.chmod(p, stat.S_IMODE(st.st_mode) | stat.S_IWGRP)
            os.chown(p, -1, gid)

            if fn in EXPECTED_SHARED_LIBS:
                # check to see that this file was linked with the
                # expected library path and version
                data = captureCommand("otool -L %s" % shellQuote(p))
                for sl in EXPECTED_SHARED_LIBS[fn]:
                    if ("\t%s " % sl) not in data:
                        print("Expected shared lib %s was not linked with %s"
                                % (sl, p))
                        shared_lib_error = True

    if shared_lib_error:
        fatal("Unexpected shared library errors.")

    if PYTHON_3:
        LDVERSION=None
        VERSION=None
        ABIFLAGS=None

        fp = open(os.path.join(buildDir, 'Makefile'), 'r')
        for ln in fp:
            if ln.startswith('VERSION='):
                VERSION=ln.split()[1]
            if ln.startswith('ABIFLAGS='):
                ABIFLAGS=ln.split()
                ABIFLAGS=ABIFLAGS[1] if len(ABIFLAGS) > 1 else ''
            if ln.startswith('LDVERSION='):
                LDVERSION=ln.split()[1]
        fp.close()

        LDVERSION = LDVERSION.replace('$(VERSION)', VERSION)
        LDVERSION = LDVERSION.replace('$(ABIFLAGS)', ABIFLAGS)
        config_suffix = '-' + LDVERSION
        if getVersionMajorMinor() >= (3, 6):
            config_suffix = config_suffix + '-darwin'
    else:
        config_suffix = ''      # Python 2.x

    # We added some directories to the search path during the configure
    # phase. Remove those because those directories won't be there on
    # the end-users system. Also remove the directories from _sysconfigdata.py
    # (added in 3.3) if it exists.

    include_path = '-I%s/libraries/usr/local/include' % (WORKDIR,)
    lib_path = '-L%s/libraries/usr/local/lib' % (WORKDIR,)

    # fix Makefile
    path = os.path.join(path_to_lib, 'config' + config_suffix, 'Makefile')
    fp = open(path, 'r')
    data = fp.read()
    fp.close()

    for p in (include_path, lib_path):
        data = data.replace(" " + p, '')
        data = data.replace(p + " ", '')

    fp = open(path, 'w')
    fp.write(data)
    fp.close()

    # fix _sysconfigdata
    #
    # TODO: make this more robust!  test_sysconfig_module of
    # distutils.tests.test_sysconfig.SysconfigTestCase tests that
    # the output from get_config_var in both sysconfig and
    # distutils.sysconfig is exactly the same for both CFLAGS and
    # LDFLAGS.  The fixing up is now complicated by the pretty
    # printing in _sysconfigdata.py.  Also, we are using the
    # pprint from the Python running the installer build which
    # may not cosmetically format the same as the pprint in the Python
    # being built (and which is used to originally generate
    # _sysconfigdata.py).

    import pprint
    if getVersionMajorMinor() >= (3, 6):
        # XXX this is extra-fragile
        path = os.path.join(path_to_lib,
            '_sysconfigdata_%s_darwin_darwin.py' % (ABIFLAGS,))
    else:
        path = os.path.join(path_to_lib, '_sysconfigdata.py')
    fp = open(path, 'r')
    data = fp.read()
    fp.close()
    # create build_time_vars dict
    if RUNNING_ON_PYTHON2:
        exec(data)
    else:
        g_dict = {}
        l_dict = {}
        exec(data, g_dict, l_dict)
        build_time_vars = l_dict['build_time_vars']
    vars = {}
    for k, v in build_time_vars.items():
        if type(v) == type(''):
            for p in (include_path, lib_path):
                v = v.replace(' ' + p, '')
                v = v.replace(p + ' ', '')
        vars[k] = v

    fp = open(path, 'w')
    # duplicated from sysconfig._generate_posix_vars()
    fp.write('# system configuration generated and used by'
                ' the sysconfig module\n')
    fp.write('build_time_vars = ')
    pprint.pprint(vars, stream=fp)
    fp.close()

    # Add symlinks in /usr/local/bin, using relative links
    usr_local_bin = os.path.join(rootDir, 'usr', 'local', 'bin')
    to_framework = os.path.join('..', '..', '..', 'Library', 'Frameworks',
            'Python.framework', 'Versions', version, 'bin')
    if os.path.exists(usr_local_bin):
        shutil.rmtree(usr_local_bin)
    os.makedirs(usr_local_bin)
    for fn in os.listdir(
                os.path.join(frmDir, 'Versions', version, 'bin')):
        os.symlink(os.path.join(to_framework, fn),
                   os.path.join(usr_local_bin, fn))

    os.chdir(curdir)

def patchFile(inPath, outPath):
    data = fileContents(inPath)
    data = data.replace('$FULL_VERSION', getFullVersion())
    data = data.replace('$VERSION', getVersion())
    data = data.replace('$MACOSX_DEPLOYMENT_TARGET', ''.join((DEPTARGET, ' or later')))
    data = data.replace('$ARCHITECTURES', ", ".join(universal_opts_map[UNIVERSALARCHS]))
    data = data.replace('$INSTALL_SIZE', installSize())
    data = data.replace('$THIRD_PARTY_LIBS', "\\\n".join(THIRD_PARTY_LIBS))

    # This one is not handy as a template variable
    data = data.replace('$PYTHONFRAMEWORKINSTALLDIR', '/Library/Frameworks/Python.framework')
    fp = open(outPath, 'w')
    fp.write(data)
    fp.close()

def patchScript(inPath, outPath):
    major, minor = getVersionMajorMinor()
    data = fileContents(inPath)
    data = data.replace('@PYMAJOR@', str(major))
    data = data.replace('@PYVER@', getVersion())
    fp = open(outPath, 'w')
    fp.write(data)
    fp.close()
    os.chmod(outPath, STAT_0o755)



def packageFromRecipe(targetDir, recipe):
    curdir = os.getcwd()
    try:
        # The major version (such as 2.5) is included in the package name
        # because having two version of python installed at the same time is
        # common.
        pkgname = '%s-%s'%(recipe['name'], getVersion())
        srcdir  = recipe.get('source')
        pkgroot = recipe.get('topdir', srcdir)
        postflight = recipe.get('postflight')
        readme = textwrap.dedent(recipe['readme'])
        isRequired = recipe.get('required', True)

        print("- building package %s"%(pkgname,))

        # Substitute some variables
        textvars = dict(
            VER=getVersion(),
            FULLVER=getFullVersion(),
        )
        readme = readme % textvars

        if pkgroot is not None:
            pkgroot = pkgroot % textvars
        else:
            pkgroot = '/'

        if srcdir is not None:
            srcdir = os.path.join(WORKDIR, '_root', srcdir[1:])
            srcdir = srcdir % textvars

        if postflight is not None:
            postflight = os.path.abspath(postflight)

        packageContents = os.path.join(targetDir, pkgname + '.pkg', 'Contents')
        os.makedirs(packageContents)

        if srcdir is not None:
            os.chdir(srcdir)
            runCommand("pax -wf %s . 2>&1"%(shellQuote(os.path.join(packageContents, 'Archive.pax')),))
            runCommand("gzip -9 %s 2>&1"%(shellQuote(os.path.join(packageContents, 'Archive.pax')),))
            runCommand("mkbom . %s 2>&1"%(shellQuote(os.path.join(packageContents, 'Archive.bom')),))

        fn = os.path.join(packageContents, 'PkgInfo')
        fp = open(fn, 'w')
        fp.write('pmkrpkg1')
        fp.close()

        rsrcDir = os.path.join(packageContents, "Resources")
        os.mkdir(rsrcDir)
        fp = open(os.path.join(rsrcDir, 'ReadMe.txt'), 'w')
        fp.write(readme)
        fp.close()

        if postflight is not None:
            patchScript(postflight, os.path.join(rsrcDir, 'postflight'))

        vers = getFullVersion()
        major, minor = getVersionMajorMinor()
        pl = dict(
                CFBundleGetInfoString="Python.%s %s"%(pkgname, vers,),
                CFBundleIdentifier='org.python.Python.%s'%(pkgname,),
                CFBundleName='Python.%s'%(pkgname,),
                CFBundleShortVersionString=vers,
                IFMajorVersion=major,
                IFMinorVersion=minor,
                IFPkgFormatVersion=0.10000000149011612,
                IFPkgFlagAllowBackRev=False,
                IFPkgFlagAuthorizationAction="RootAuthorization",
                IFPkgFlagDefaultLocation=pkgroot,
                IFPkgFlagFollowLinks=True,
                IFPkgFlagInstallFat=True,
                IFPkgFlagIsRequired=isRequired,
                IFPkgFlagOverwritePermissions=False,
                IFPkgFlagRelocatable=False,
                IFPkgFlagRestartAction="NoRestart",
                IFPkgFlagRootVolumeOnly=True,
                IFPkgFlagUpdateInstalledLangauges=False,
            )
        writePlist(pl, os.path.join(packageContents, 'Info.plist'))

        pl = dict(
                    IFPkgDescriptionDescription=readme,
                    IFPkgDescriptionTitle=recipe.get('long_name', "Python.%s"%(pkgname,)),
                    IFPkgDescriptionVersion=vers,
                )
        writePlist(pl, os.path.join(packageContents, 'Resources', 'Description.plist'))

    finally:
        os.chdir(curdir)


def makeMpkgPlist(path):

    vers = getFullVersion()
    major, minor = getVersionMajorMinor()

    pl = dict(
            CFBundleGetInfoString="Python %s"%(vers,),
            CFBundleIdentifier='org.python.Python',
            CFBundleName='Python',
            CFBundleShortVersionString=vers,
            IFMajorVersion=major,
            IFMinorVersion=minor,
            IFPkgFlagComponentDirectory="Contents/Packages",
            IFPkgFlagPackageList=[
                dict(
                    IFPkgFlagPackageLocation='%s-%s.pkg'%(item['name'], getVersion()),
                    IFPkgFlagPackageSelection=item.get('selected', 'selected'),
                )
                for item in pkg_recipes()
            ],
            IFPkgFormatVersion=0.10000000149011612,
            IFPkgFlagBackgroundScaling="proportional",
            IFPkgFlagBackgroundAlignment="left",
            IFPkgFlagAuthorizationAction="RootAuthorization",
        )

    writePlist(pl, path)


def buildInstaller():

    # Zap all compiled files
    for dirpath, _, filenames in os.walk(os.path.join(WORKDIR, '_root')):
        for fn in filenames:
            if fn.endswith('.pyc') or fn.endswith('.pyo'):
                os.unlink(os.path.join(dirpath, fn))

    outdir = os.path.join(WORKDIR, 'installer')
    if os.path.exists(outdir):
        shutil.rmtree(outdir)
    os.mkdir(outdir)

    pkgroot = os.path.join(outdir, 'Python.mpkg', 'Contents')
    pkgcontents = os.path.join(pkgroot, 'Packages')
    os.makedirs(pkgcontents)
    for recipe in pkg_recipes():
        packageFromRecipe(pkgcontents, recipe)

    rsrcDir = os.path.join(pkgroot, 'Resources')

    fn = os.path.join(pkgroot, 'PkgInfo')
    fp = open(fn, 'w')
    fp.write('pmkrpkg1')
    fp.close()

    os.mkdir(rsrcDir)

    makeMpkgPlist(os.path.join(pkgroot, 'Info.plist'))
    pl = dict(
                IFPkgDescriptionTitle="Python",
                IFPkgDescriptionVersion=getVersion(),
            )

    writePlist(pl, os.path.join(pkgroot, 'Resources', 'Description.plist'))
    for fn in os.listdir('resources'):
        if fn == '.svn': continue
        if fn.endswith('.jpg'):
            shutil.copy(os.path.join('resources', fn), os.path.join(rsrcDir, fn))
        else:
            patchFile(os.path.join('resources', fn), os.path.join(rsrcDir, fn))


def installSize(clear=False, _saved=[]):
    if clear:
        del _saved[:]
    if not _saved:
        data = captureCommand("du -ks %s"%(
                    shellQuote(os.path.join(WORKDIR, '_root'))))
        _saved.append("%d"%((0.5 + (int(data.split()[0]) / 1024.0)),))
    return _saved[0]


def buildDMG():
    """
    Create DMG containing the rootDir.
    """
    outdir = os.path.join(WORKDIR, 'diskimage')
    if os.path.exists(outdir):
        shutil.rmtree(outdir)

    # We used to use the deployment target as the last characters of the
    # installer file name. With the introduction of weaklinked installer
    # variants, we may have two variants with the same file name, i.e.
    # both ending in '10.9'.  To avoid this, we now use the major/minor
    # version numbers of the macOS version we are building on, i.e.
    # '10.9' as before for 10.9+ variant, '11.0' for universal2 11.0-.
    # it's not ideal but should cause the least disruption to packaging
    # workflows.
    build_system_version = '.'.join(platform.mac_ver()[0].split('.')[0:2])
    imagepath = os.path.join(outdir,
                    'python-%s-macosx%s'%(getFullVersion(),build_system_version))
    if INCLUDE_TIMESTAMP:
        imagepath = imagepath + '-%04d-%02d-%02d'%(time.localtime()[:3])
    imagepath = imagepath + '.dmg'

    os.mkdir(outdir)

    # Try to mitigate race condition in certain versions of macOS, e.g. 10.9,
    # when hdiutil create fails with  "Resource busy".  For now, just retry
    # the create a few times and hope that it eventually works.

    volname='Python %s'%(getFullVersion())
    cmd = ("hdiutil create -format UDRW -volname %s -srcfolder %s -size 100m %s"%(
            shellQuote(volname),
            shellQuote(os.path.join(WORKDIR, 'installer')),
            shellQuote(imagepath + ".tmp.dmg" )))
    for i in range(5):
        fd = os.popen(cmd, 'r')
        data = fd.read()
        xit = fd.close()
        if not xit:
            break
        sys.stdout.write(data)
        print(" -- retrying hdiutil create")
        time.sleep(5)
    else:
        raise RuntimeError("command failed: %s"%(cmd,))

    if not os.path.exists(os.path.join(WORKDIR, "mnt")):
        os.mkdir(os.path.join(WORKDIR, "mnt"))
    runCommand("hdiutil attach %s -mountroot %s"%(
        shellQuote(imagepath + ".tmp.dmg"), shellQuote(os.path.join(WORKDIR, "mnt"))))

    # Custom icon for the DMG, shown when the DMG is mounted.
    shutil.copy("../Icons/Disk Image.icns",
            os.path.join(WORKDIR, "mnt", volname, ".VolumeIcon.icns"))
    runCommand("SetFile -a C %s/"%(
            shellQuote(os.path.join(WORKDIR, "mnt", volname)),))

    runCommand("hdiutil detach %s"%(shellQuote(os.path.join(WORKDIR, "mnt", volname))))

    setIcon(imagepath + ".tmp.dmg", "../Icons/Disk Image.icns")
    runCommand("hdiutil convert %s -format UDZO -o %s"%(
            shellQuote(imagepath + ".tmp.dmg"), shellQuote(imagepath)))
    setIcon(imagepath, "../Icons/Disk Image.icns")

    os.unlink(imagepath + ".tmp.dmg")

    return imagepath


def setIcon(filePath, icnsPath):
    """
    Set the custom icon for the specified file or directory.
    """

    dirPath = os.path.normpath(os.path.dirname(__file__))
    toolPath = os.path.join(dirPath, "seticon.app/Contents/MacOS/seticon")
    if not os.path.exists(toolPath) or os.stat(toolPath).st_mtime < os.stat(dirPath + '/seticon.m').st_mtime:
        # NOTE: The tool is created inside an .app bundle, otherwise it won't work due
        # to connections to the window server.
        appPath = os.path.join(dirPath, "seticon.app/Contents/MacOS")
        if not os.path.exists(appPath):
            os.makedirs(appPath)
        runCommand("cc -o %s %s/seticon.m -framework Cocoa"%(
            shellQuote(toolPath), shellQuote(dirPath)))

    runCommand("%s %s %s"%(shellQuote(os.path.abspath(toolPath)), shellQuote(icnsPath),
        shellQuote(filePath)))

def main():
    # First parse options and check if we can perform our work
    parseOptions()
    checkEnvironment()

    os.environ['MACOSX_DEPLOYMENT_TARGET'] = DEPTARGET
    os.environ['CC'] = CC
    os.environ['CXX'] = CXX

    if os.path.exists(WORKDIR):
        shutil.rmtree(WORKDIR)
    os.mkdir(WORKDIR)

    os.environ['LC_ALL'] = 'C'

    # Then build third-party libraries such as sleepycat DB4.
    buildLibraries()

    # Now build python itself
    buildPython()

    # And then build the documentation
    # Remove the Deployment Target from the shell
    # environment, it's no longer needed and
    # an unexpected build target can cause problems
    # when Sphinx and its dependencies need to
    # be (re-)installed.
    del os.environ['MACOSX_DEPLOYMENT_TARGET']
    buildPythonDocs()


    # Prepare the applications folder
    folder = os.path.join(WORKDIR, "_root", "Applications", "Python %s"%(
        getVersion(),))
    fn = os.path.join(folder, "License.rtf")
    patchFile("resources/License.rtf",  fn)
    fn = os.path.join(folder, "ReadMe.rtf")
    patchFile("resources/ReadMe.rtf",  fn)
    fn = os.path.join(folder, "Update Shell Profile.command")
    patchScript("scripts/postflight.patch-profile",  fn)
    fn = os.path.join(folder, "Install Certificates.command")
    patchScript("resources/install_certificates.command",  fn)
    os.chmod(folder, STAT_0o755)
    setIcon(folder, "../Icons/Python Folder.icns")

    # Create the installer
    buildInstaller()

    # And copy the readme into the directory containing the installer
    patchFile('resources/ReadMe.rtf',
                os.path.join(WORKDIR, 'installer', 'ReadMe.rtf'))

    # Ditto for the license file.
    patchFile('resources/License.rtf',
                os.path.join(WORKDIR, 'installer', 'License.rtf'))

    fp = open(os.path.join(WORKDIR, 'installer', 'Build.txt'), 'w')
    fp.write("# BUILD INFO\n")
    fp.write("# Date: %s\n" % time.ctime())
    fp.write("# By: %s\n" % pwd.getpwuid(os.getuid()).pw_gecos)
    fp.close()

    # And copy it to a DMG
    buildDMG()

if __name__ == "__main__":
    main()

external/cpython_master/Mac/BuildScript/build-installer.py
DragonHPC/dragon
# An absurd workaround for the lack of arithmetic in MS's resource compiler.
# After building Python, run this, then paste the output into the appropriate
# part of PC\python_nt.rc.
# Example output:
#
# * For 2.3a0,
# * PY_MICRO_VERSION = 0
# * PY_RELEASE_LEVEL = 'alpha' = 0xA
# * PY_RELEASE_SERIAL = 1
# *
# * and 0*1000 + 10*10 + 1 = 101.
# */
# #define FIELD3 101

import sys

major, minor, micro, level, serial = sys.version_info
levelnum = {'alpha': 0xA,
            'beta': 0xB,
            'candidate': 0xC,
            'final': 0xF,
           }[level]
string = sys.version.split()[0] # like '2.3a0'

print(" * For %s," % string)
print(" * PY_MICRO_VERSION = %d" % micro)
print(" * PY_RELEASE_LEVEL = %r = %s" % (level, hex(levelnum)))
print(" * PY_RELEASE_SERIAL = %d" % serial)
print(" *")

field3 = micro * 1000 + levelnum * 10 + serial

print(" * and %d*1000 + %d*10 + %d = %d" % (micro, levelnum, serial, field3))
print(" */")
print("#define FIELD3", field3)

external/cpython_master/PCbuild/field3.py
DragonHPC/dragon
#!/usr/bin/env python3
#
# Argument Clinic
# Copyright 2012-2013 by Larry Hastings.
# Licensed to the PSF under a contributor agreement.
#

import abc
import ast
import collections
import contextlib
import copy
import cpp
import functools
import hashlib
import inspect
import io
import itertools
import os
import pprint
import re
import shlex
import string
import sys
import tempfile
import textwrap
import traceback
import types

from types import *
NoneType = type(None)

# TODO:
#
# soon:
#
# * allow mixing any two of {positional-only, positional-or-keyword,
#   keyword-only}
#       * dict constructor uses positional-only and keyword-only
#       * max and min use positional only with an optional group
#         and keyword-only
#

version = '1'

NoneType = type(None)

class Unspecified:
    def __repr__(self):
        return '<Unspecified>'

unspecified = Unspecified()


class Null:
    def __repr__(self):
        return '<Null>'

NULL = Null()


class Unknown:
    def __repr__(self):
        return '<Unknown>'

unknown = Unknown()

sig_end_marker = '--'


_text_accumulator_nt = collections.namedtuple("_text_accumulator", "text append output")

def _text_accumulator():
    text = []
    def output():
        s = ''.join(text)
        text.clear()
        return s
    return _text_accumulator_nt(text, text.append, output)


text_accumulator_nt = collections.namedtuple("text_accumulator", "text append")

def text_accumulator():
    """
    Creates a simple text accumulator / joiner.

    Returns a pair of callables:
        append, output
    "append" appends a string to the accumulator.
    "output" returns the contents of the accumulator
       joined together (''.join(accumulator)) and
       empties the accumulator.
    """
    text, append, output = _text_accumulator()
    return text_accumulator_nt(append, output)


def warn_or_fail(fail=False, *args, filename=None, line_number=None):
    joined = " ".join([str(a) for a in args])
    add, output = text_accumulator()
    if fail:
        add("Error")
    else:
        add("Warning")
    if clinic:
        if filename is None:
            filename = clinic.filename
        if getattr(clinic, 'block_parser', None) and (line_number is None):
            line_number = clinic.block_parser.line_number
    if filename is not None:
        add(' in file "' + filename + '"')
    if line_number is not None:
        add(" on line " + str(line_number))
    add(':\n')
    add(joined)
    print(output())
    if fail:
        sys.exit(-1)


def warn(*args, filename=None, line_number=None):
    return warn_or_fail(False, *args, filename=filename, line_number=line_number)

def fail(*args, filename=None, line_number=None):
    return warn_or_fail(True, *args, filename=filename, line_number=line_number)


def quoted_for_c_string(s):
    for old, new in (
        ('\\', '\\\\'), # must be first!
        ('"', '\\"'),
        ("'", "\\'"),
        ):
        s = s.replace(old, new)
    return s

def c_repr(s):
    return '"' + s + '"'


is_legal_c_identifier = re.compile('^[A-Za-z_][A-Za-z0-9_]*$').match

def is_legal_py_identifier(s):
    return all(is_legal_c_identifier(field) for field in s.split('.'))

# identifiers that are okay in Python but aren't a good idea in C.
# so if they're used Argument Clinic will add "_value" to the end
# of the name in C.
c_keywords = set("""
asm auto break case char const continue default do double
else enum extern float for goto if inline int long
register return short signed sizeof static struct switch
typedef typeof union unsigned void volatile while
""".strip().split())

def ensure_legal_c_identifier(s):
    # for now, just complain if what we're given isn't legal
    if not is_legal_c_identifier(s):
        fail("Illegal C identifier: {}".format(s))
    # but if we picked a C keyword, pick something else
    if s in c_keywords:
        return s + "_value"
    return s

def rstrip_lines(s):
    text, add, output = _text_accumulator()
    for line in s.split('\n'):
        add(line.rstrip())
        add('\n')
    text.pop()
    return output()

def format_escape(s):
    # double up curly-braces, this string will be used
    # as part of a format_map() template later
    s = s.replace('{', '{{')
    s = s.replace('}', '}}')
    return s

def linear_format(s, **kwargs):
    """
    Perform str.format-like substitution, except:
      * The strings substituted must be on lines by
        themselves.  (This line is the "source line".)
      * If the substitution text is empty, the source line
        is removed in the output.
      * If the field is not recognized, the original line
        is passed unmodified through to the output.
      * If the substitution text is not empty:
          * Each line of the substituted text is indented
            by the indent of the source line.
          * A newline will be added to the end.
    """

    add, output = text_accumulator()
    for line in s.split('\n'):
        indent, curly, trailing = line.partition('{')
        if not curly:
            add(line)
            add('\n')
            continue

        name, curly, trailing = trailing.partition('}')
        if not curly or name not in kwargs:
            add(line)
            add('\n')
            continue

        if trailing:
            fail("Text found after {" + name + "} block marker!  It must be on a line by itself.")
        if indent.strip():
            fail("Non-whitespace characters found before {" + name + "} block marker!  It must be on a line by itself.")

        value = kwargs[name]
        if not value:
            continue

        value = textwrap.indent(rstrip_lines(value), indent)
        add(value)
        add('\n')

    return output()[:-1]

def indent_all_lines(s, prefix):
    """
    Returns 's', with 'prefix' prepended to all lines.

    If the last line is empty, prefix is not prepended
    to it.  (If s is blank, returns s unchanged.)

    (textwrap.indent only adds to non-blank lines.)
    """
    split = s.split('\n')
    last = split.pop()
    final = []
    for line in split:
        final.append(prefix)
        final.append(line)
        final.append('\n')
    if last:
        final.append(prefix)
        final.append(last)
    return ''.join(final)

def suffix_all_lines(s, suffix):
    """
    Returns 's', with 'suffix' appended to all lines.

    If the last line is empty, suffix is not appended
    to it.  (If s is blank, returns s unchanged.)
    """
    split = s.split('\n')
    last = split.pop()
    final = []
    for line in split:
        final.append(line)
        final.append(suffix)
        final.append('\n')
    if last:
        final.append(last)
        final.append(suffix)
    return ''.join(final)


def version_splitter(s):
    """Splits a version string into a tuple of integers.

    The following ASCII characters are allowed, and employ
    the following conversions:
        a -> -3
        b -> -2
        c -> -1
    (This permits Python-style version strings such as "1.4b3".)
    """
    version = []
    accumulator = []
    def flush():
        if not accumulator:
            raise ValueError('Unsupported version string: ' + repr(s))
        version.append(int(''.join(accumulator)))
        accumulator.clear()

    for c in s:
        if c.isdigit():
            accumulator.append(c)
        elif c == '.':
            flush()
        elif c in 'abc':
            flush()
            version.append('abc'.index(c) - 3)
        else:
            raise ValueError('Illegal character ' + repr(c) + ' in version string ' + repr(s))
    flush()
    return tuple(version)

def version_comparitor(version1, version2):
    iterator = itertools.zip_longest(version_splitter(version1), version_splitter(version2), fillvalue=0)
    for i, (a, b) in enumerate(iterator):
        if a < b:
            return -1
        if a > b:
            return 1
    return 0


class CRenderData:
    def __init__(self):

        # The C statements to declare variables.
        # Should be full lines with \n eol characters.
        self.declarations = []

        # The C statements required to initialize the variables before the parse call.
        # Should be full lines with \n eol characters.
        self.initializers = []

        # The C statements needed to dynamically modify the values
        # parsed by the parse call, before calling the impl.
        self.modifications = []

        # The entries for the "keywords" array for PyArg_ParseTuple.
        # Should be individual strings representing the names.
        self.keywords = []

        # The "format units" for PyArg_ParseTuple.
        # Should be individual strings that will get
        self.format_units = []

        # The varargs arguments for PyArg_ParseTuple.
        self.parse_arguments = []

        # The parameter declarations for the impl function.
        self.impl_parameters = []

        # The arguments to the impl function at the time it's called.
        self.impl_arguments = []

        # For return converters: the name of the variable that
        # should receive the value returned by the impl.
        self.return_value = "return_value"

        # For return converters: the code to convert the return
        # value from the parse function.  This is also where
        # you should check the _return_value for errors, and
        # "goto exit" if there are any.
        self.return_conversion = []

        # The C statements required to clean up after the impl call.
        self.cleanup = []


class FormatCounterFormatter(string.Formatter):
    """
    This counts how many instances of each formatter
    "replacement string" appear in the format string.

    e.g. after evaluating "string {a}, {b}, {c}, {a}"
         the counts dict would now look like
         {'a': 2, 'b': 1, 'c': 1}
    """
    def __init__(self):
        self.counts = collections.Counter()

    def get_value(self, key, args, kwargs):
        self.counts[key] += 1
        return ''

class Language(metaclass=abc.ABCMeta):

    start_line = ""
    body_prefix = ""
    stop_line = ""
    checksum_line = ""

    def __init__(self, filename):
        pass

    @abc.abstractmethod
    def render(self, clinic, signatures):
        pass

    def parse_line(self, line):
        pass

    def validate(self):
        def assert_only_one(attr, *additional_fields):
            """
            Ensures that the string found at getattr(self, attr)
            contains exactly one formatter replacement string for
            each valid field.  The list of valid fields is
            ['dsl_name'] extended by additional_fields.

            e.g.
                self.fmt = "{dsl_name} {a} {b}"

                # this passes
                self.assert_only_one('fmt', 'a', 'b')

                # this fails, the format string has a {b} in it
                self.assert_only_one('fmt', 'a')

                # this fails, the format string doesn't have a {c} in it
                self.assert_only_one('fmt', 'a', 'b', 'c')

                # this fails, the format string has two {a}s in it,
                # it must contain exactly one
                self.fmt2 = '{dsl_name} {a} {a}'
                self.assert_only_one('fmt2', 'a')

            """
            fields = ['dsl_name']
            fields.extend(additional_fields)
            line = getattr(self, attr)
            fcf = FormatCounterFormatter()
            fcf.format(line)
            def local_fail(should_be_there_but_isnt):
                if should_be_there_but_isnt:
                    fail("{} {} must contain {{{}}} exactly once!".format(
                        self.__class__.__name__, attr, name))
                else:
                    fail("{} {} must not contain {{{}}}!".format(
                        self.__class__.__name__, attr, name))

            for name, count in fcf.counts.items():
                if name in fields:
                    if count > 1:
                        local_fail(True)
                else:
                    local_fail(False)
            for name in fields:
                if fcf.counts.get(name) != 1:
                    local_fail(True)

        assert_only_one('start_line')
        assert_only_one('stop_line')

        field = "arguments" if "{arguments}" in self.checksum_line else "checksum"
        assert_only_one('checksum_line', field)



class PythonLanguage(Language):

    language      = 'Python'
    start_line    = "#/*[{dsl_name} input]"
    body_prefix   = "#"
    stop_line     = "#[{dsl_name} start generated code]*/"
    checksum_line = "#/*[{dsl_name} end generated code: {arguments}]*/"


def permute_left_option_groups(l):
    """
    Given [1, 2, 3], should yield:
       ()
       (3,)
       (2, 3)
       (1, 2, 3)
    """
    yield tuple()
    accumulator = []
    for group in reversed(l):
        accumulator = list(group) + accumulator
        yield tuple(accumulator)


def permute_right_option_groups(l):
    """
    Given [1, 2, 3], should yield:
      ()
      (1,)
      (1, 2)
      (1, 2, 3)
    """
    yield tuple()
    accumulator = []
    for group in l:
        accumulator.extend(group)
        yield tuple(accumulator)


def permute_optional_groups(left, required, right):
    """
    Generator function that computes the set of acceptable
    argument lists for the provided iterables of
    argument groups.  (Actually it generates a tuple of tuples.)

    Algorithm: prefer left options over right options.

    If required is empty, left must also be empty.
    """
    required = tuple(required)
    result = []

    if not required:
        assert not left

    accumulator = []
    counts = set()
    for r in permute_right_option_groups(right):
        for l in permute_left_option_groups(left):
            t = l + required + r
            if len(t) in counts:
                continue
            counts.add(len(t))
            accumulator.append(t)

    accumulator.sort(key=len)
    return tuple(accumulator)


def strip_leading_and_trailing_blank_lines(s):
    lines = s.rstrip().split('\n')
    while lines:
        line = lines[0]
        if line.strip():
            break
        del lines[0]
    return '\n'.join(lines)

@functools.lru_cache()
def normalize_snippet(s, *, indent=0):
    """
    Reformats s:
        * removes leading and trailing blank lines
        * ensures that it does not end with a newline
        * dedents so the first nonwhite character on any line is at column "indent"
    """
    s = strip_leading_and_trailing_blank_lines(s)
    s = textwrap.dedent(s)
    if indent:
        s = textwrap.indent(s, ' ' * indent)
    return s


def wrap_declarations(text, length=78):
    """
    A simple-minded text wrapper for C function declarations.

    It views a declaration line as looking like this:
        xxxxxxxx(xxxxxxxxx,xxxxxxxxx)
    If called with length=30, it would wrap that line into
        xxxxxxxx(xxxxxxxxx,
                 xxxxxxxxx)
    (If the declaration has zero or one parameters, this
    function won't wrap it.)

    If this doesn't work properly, it's probably better to
    start from scratch with a more sophisticated algorithm,
    rather than try and improve/debug this dumb little function.
    """
    lines = []
    for line in text.split('\n'):
        prefix, _, after_l_paren = line.partition('(')
        if not after_l_paren:
            lines.append(line)
            continue
        parameters, _, after_r_paren = after_l_paren.partition(')')
        if not _:
            lines.append(line)
            continue
        if ',' not in parameters:
            lines.append(line)
            continue
        parameters = [x.strip() + ", " for x in parameters.split(',')]
        prefix += "("
        if len(prefix) < length:
            spaces = " " * len(prefix)
        else:
            spaces = " " * 4

        while parameters:
            line = prefix
            first = True
            while parameters:
                if (not first and
                    (len(line) + len(parameters[0]) > length)):
                    break
                line += parameters.pop(0)
                first = False
            if not parameters:
                line = line.rstrip(", ") + ")" + after_r_paren
            lines.append(line.rstrip())
            prefix = spaces
    return "\n".join(lines)


class CLanguage(Language):

    body_prefix   = "#"
    language      = 'C'
    start_line    = "/*[{dsl_name} input]"
    body_prefix   = ""
    stop_line     = "[{dsl_name} start generated code]*/"
    checksum_line = "/*[{dsl_name} end generated code: {arguments}]*/"

    def __init__(self, filename):
        super().__init__(filename)
        self.cpp = cpp.Monitor(filename)
        self.cpp.fail = fail

    def parse_line(self, line):
        self.cpp.writeline(line)

    def render(self, clinic, signatures):
        function = None
        for o in signatures:
            if isinstance(o, Function):
                if function:
                    fail("You may specify at most one function per block.\nFound a block containing at least two:\n\t" + repr(function) + " and " + repr(o))
                function = o
        return self.render_function(clinic, function)

    def docstring_for_c_string(self, f):
        if re.search(r'[^\x00-\x7F]', f.docstring):
            warn("Non-ascii character appear in docstring.")

        text, add, output = _text_accumulator()
        # turn docstring into a properly quoted C string
        for line in f.docstring.split('\n'):
            add('"')
            add(quoted_for_c_string(line))
            add('\\n"\n')

        if text[-2] == sig_end_marker:
            # If we only have a signature, add the blank line that the
            # __text_signature__ getter expects to be there.
            add('"\\n"')
        else:
            text.pop()
            add('"')
        return ''.join(text)

    def output_templates(self, f):
        parameters = list(f.parameters.values())
        assert parameters
        assert isinstance(parameters[0].converter, self_converter)
        del parameters[0]
        converters = [p.converter for p in parameters]

        has_option_groups = parameters and (parameters[0].group or parameters[-1].group)
        default_return_converter = (not f.return_converter or
            f.return_converter.type == 'PyObject *')

        new_or_init = f.kind in (METHOD_NEW, METHOD_INIT)

        pos_only = min_pos = max_pos = min_kw_only = 0
        for i, p in enumerate(parameters, 1):
            if p.is_keyword_only():
                assert not p.is_positional_only()
                if not p.is_optional():
                    min_kw_only = i - max_pos
            else:
                max_pos = i
                if p.is_positional_only():
                    pos_only = i
                if not p.is_optional():
                    min_pos = i

        requires_defining_class = any(
            isinstance(p.converter, defining_class_converter)
            for p in parameters)

        meth_o = (len(parameters) == 1 and
              parameters[0].is_positional_only() and
              not converters[0].is_optional() and
              not requires_defining_class and
              not new_or_init)

        # we have to set these things before we're done:
        #
        # docstring_prototype
        # docstring_definition
        # impl_prototype
        # methoddef_define
        # parser_prototype
        # parser_definition
        # impl_definition
        # cpp_if
        # cpp_endif
        # methoddef_ifndef

        return_value_declaration = "PyObject *return_value = NULL;"

        methoddef_define = normalize_snippet("""
            #define {methoddef_name}    \\
                {{"{name}", {methoddef_cast}{c_basename}, {methoddef_flags}, {c_basename}__doc__}},
            """)
        if new_or_init and not f.docstring:
            docstring_prototype = docstring_definition = ''
        else:
            docstring_prototype = normalize_snippet("""
                PyDoc_VAR({c_basename}__doc__);
                """)
            docstring_definition = normalize_snippet("""
                PyDoc_STRVAR({c_basename}__doc__,
                {docstring});
                """)
        impl_definition = normalize_snippet("""
            static {impl_return_type}
            {c_basename}_impl({impl_parameters})
            """)
        impl_prototype = parser_prototype = parser_definition = None

        parser_prototype_keyword = normalize_snippet("""
            static PyObject *
            {c_basename}({self_type}{self_name}, PyObject *args, PyObject *kwargs)
            """)

        parser_prototype_varargs = normalize_snippet("""
            static PyObject *
            {c_basename}({self_type}{self_name}, PyObject *args)
            """)

        parser_prototype_fastcall = normalize_snippet("""
            static PyObject *
            {c_basename}({self_type}{self_name}, PyObject *const *args, Py_ssize_t nargs)
            """)

        parser_prototype_fastcall_keywords = normalize_snippet("""
            static PyObject *
            {c_basename}({self_type}{self_name}, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
            """)

        parser_prototype_def_class = normalize_snippet("""
            static PyObject *
            {c_basename}({self_type}{self_name}, PyTypeObject *{defining_class_name}, PyObject *const *args, Py_ssize_t nargs, PyObject *kwnames)
        """)

        # parser_body_fields remembers the fields passed in to the
        # previous call to parser_body. this is used for an awful hack.
        parser_body_fields = ()
        parser_body_declarations = ''
        def parser_body(prototype, *fields, declarations=''):
            nonlocal parser_body_fields, parser_body_declarations
            add, output = text_accumulator()
            add(prototype)
            parser_body_fields = fields
            parser_body_declarations = declarations

            fields = list(fields)
            fields.insert(0, normalize_snippet("""
                {{
                    {return_value_declaration}
                    {parser_declarations}
                    {declarations}
                    {initializers}
                """) + "\n")
            # just imagine--your code is here in the middle
            fields.append(normalize_snippet("""
                    {modifications}
                    {return_value} = {c_basename}_impl({impl_arguments});
                    {return_conversion}

                {exit_label}
                    {cleanup}
                    return return_value;
                }}
                """))
            for field in fields:
                add('\n')
                add(field)
            return linear_format(output(), parser_declarations=declarations)

        if not parameters:
            # no parameters, METH_NOARGS

            flags = "METH_NOARGS"

            parser_prototype = normalize_snippet("""
                static PyObject *
                {c_basename}({self_type}{self_name}, PyObject *Py_UNUSED(ignored))
                """)
            parser_definition = parser_prototype

            if default_return_converter:
                parser_definition = parser_prototype + '\n' + normalize_snippet("""
                    {{
                        return {c_basename}_impl({impl_arguments});
                    }}
                    """)
            else:
                parser_definition = parser_body(parser_prototype)

        elif meth_o:
            flags = "METH_O"

            if (isinstance(converters[0], object_converter) and
                converters[0].format_unit == 'O'):
                meth_o_prototype = normalize_snippet("""
                    static PyObject *
                    {c_basename}({impl_parameters})
                    """)

                if default_return_converter:
                    # maps perfectly to METH_O, doesn't need a return converter.
                    # so we skip making a parse function
                    # and call directly into the impl function.
                    impl_prototype = parser_prototype = parser_definition = ''
                    impl_definition = meth_o_prototype
                else:
                    # SLIGHT HACK
                    # use impl_parameters for the parser here!
                    parser_prototype = meth_o_prototype
                    parser_definition = parser_body(parser_prototype)

            else:
                argname = 'arg'
                if parameters[0].name == argname:
                    argname += '_'
                parser_prototype = normalize_snippet("""
                    static PyObject *
                    {c_basename}({self_type}{self_name}, PyObject *%s)
                    """ % argname)

                displayname = parameters[0].get_displayname(0)
                parsearg = converters[0].parse_arg(argname, displayname)
                if parsearg is None:
                    parsearg = """
                        if (!PyArg_Parse(%s, "{format_units}:{name}", {parse_arguments})) {{
                            goto exit;
                        }}
                        """ % argname
                parser_definition = parser_body(parser_prototype,
                                                normalize_snippet(parsearg, indent=4))

        elif has_option_groups:
            # positional parameters with option groups
            # (we have to generate lots of PyArg_ParseTuple calls
            #  in a big switch statement)

            flags = "METH_VARARGS"
            parser_prototype = parser_prototype_varargs

            parser_definition = parser_body(parser_prototype, '    {option_group_parsing}')

        elif not requires_defining_class and pos_only == len(parameters):
            if not new_or_init:
                # positional-only, but no option groups
                # we only need one call to _PyArg_ParseStack

                flags = "METH_FASTCALL"
                parser_prototype = parser_prototype_fastcall
                nargs = 'nargs'
                argname_fmt = 'args[%d]'
            else:
                # positional-only, but no option groups
                # we only need one call to PyArg_ParseTuple

                flags = "METH_VARARGS"
                parser_prototype = parser_prototype_varargs
                nargs = 'PyTuple_GET_SIZE(args)'
                argname_fmt = 'PyTuple_GET_ITEM(args, %d)'

            parser_code = [normalize_snippet("""
                if (!_PyArg_CheckPositional("{name}", %s, %d, %d)) {{
                    goto exit;
                }}
                """ % (nargs, min_pos, max_pos), indent=4)]
            has_optional = False
            for i, p in enumerate(parameters):
                displayname = p.get_displayname(i+1)
                parsearg = p.converter.parse_arg(argname_fmt % i, displayname)
                if parsearg is None:
                    #print('Cannot convert %s %r for %s' % (p.converter.__class__.__name__, p.converter.format_unit, p.converter.name), file=sys.stderr)
                    parser_code = None
                    break
                if has_optional or p.is_optional():
                    has_optional = True
                    parser_code.append(normalize_snippet("""
                        if (%s < %d) {{
                            goto skip_optional;
                        }}
                        """, indent=4) % (nargs, i + 1))
                parser_code.append(normalize_snippet(parsearg, indent=4))

            if parser_code is not None:
                if has_optional:
                    parser_code.append("skip_optional:")
            else:
                if not new_or_init:
                    parser_code = [normalize_snippet("""
                        if (!_PyArg_ParseStack(args, nargs, "{format_units}:{name}",
                            {parse_arguments})) {{
                            goto exit;
                        }}
                        """, indent=4)]
                else:
                    parser_code = [normalize_snippet("""
                        if (!PyArg_ParseTuple(args, "{format_units}:{name}",
                            {parse_arguments})) {{
                            goto exit;
                        }}
                        """, indent=4)]
            parser_definition = parser_body(parser_prototype, *parser_code)

        else:
            has_optional_kw = (max(pos_only, min_pos) + min_kw_only < len(converters))
            if not new_or_init:
                flags = "METH_FASTCALL|METH_KEYWORDS"
                parser_prototype = parser_prototype_fastcall_keywords
                argname_fmt = 'args[%d]'
                declarations = normalize_snippet("""
                    static const char * const _keywords[] = {{{keywords} NULL}};
                    static _PyArg_Parser _parser = {{NULL, _keywords, "{name}", 0}};
                    PyObject *argsbuf[%s];
                    """ % len(converters))
                if has_optional_kw:
                    declarations += "\nPy_ssize_t noptargs = nargs + (kwnames ? PyTuple_GET_SIZE(kwnames) : 0) - %d;" % (min_pos + min_kw_only)
                parser_code = [normalize_snippet("""
                    args = _PyArg_UnpackKeywords(args, nargs, NULL, kwnames, &_parser, %d, %d, %d, argsbuf);
                    if (!args) {{
                        goto exit;
                    }}
                    """ % (min_pos, max_pos, min_kw_only), indent=4)]
            else:
                # positional-or-keyword arguments
                flags = "METH_VARARGS|METH_KEYWORDS"
                parser_prototype = parser_prototype_keyword
                argname_fmt = 'fastargs[%d]'
                declarations = normalize_snippet("""
                    static const char * const _keywords[] = {{{keywords} NULL}};
                    static _PyArg_Parser _parser = {{NULL, _keywords, "{name}", 0}};
                    PyObject *argsbuf[%s];
                    PyObject * const *fastargs;
                    Py_ssize_t nargs = PyTuple_GET_SIZE(args);
                    """ % len(converters))
                if has_optional_kw:
                    declarations += "\nPy_ssize_t noptargs = nargs + (kwargs ? PyDict_GET_SIZE(kwargs) : 0) - %d;" % (min_pos + min_kw_only)
                parser_code = [normalize_snippet("""
                    fastargs = _PyArg_UnpackKeywords(_PyTuple_CAST(args)->ob_item, nargs, kwargs, NULL, &_parser, %d, %d, %d, argsbuf);
                    if (!fastargs) {{
                        goto exit;
                    }}
                    """ % (min_pos, max_pos, min_kw_only), indent=4)]
            if requires_defining_class:
                flags = 'METH_METHOD|' + flags
                parser_prototype = parser_prototype_def_class

            add_label = None
            for i, p in enumerate(parameters):
                displayname = p.get_displayname(i+1)
                parsearg = p.converter.parse_arg(argname_fmt % i, displayname)
                if parsearg is None:
                    #print('Cannot convert %s %r for %s' % (p.converter.__class__.__name__, p.converter.format_unit, p.converter.name), file=sys.stderr)
                    parser_code = None
                    break
                if add_label and (i == pos_only or i == max_pos):
                    parser_code.append("%s:" % add_label)
                    add_label = None
                if not p.is_optional():
                    parser_code.append(normalize_snippet(parsearg, indent=4))
                elif i < pos_only:
                    add_label = 'skip_optional_posonly'
                    parser_code.append(normalize_snippet("""
                        if (nargs < %d) {{
                            goto %s;
                        }}
                        """ % (i + 1, add_label), indent=4))
                    if has_optional_kw:
                        parser_code.append(normalize_snippet("""
                            noptargs--;
                            """, indent=4))
                    parser_code.append(normalize_snippet(parsearg, indent=4))
                else:
                    if i < max_pos:
                        label = 'skip_optional_pos'
                        first_opt = max(min_pos, pos_only)
                    else:
                        label = 'skip_optional_kwonly'
                        first_opt = max_pos + min_kw_only
                    if i == first_opt:
                        add_label = label
                        parser_code.append(normalize_snippet("""
                            if (!noptargs) {{
                                goto %s;
                            }}
                            """ % add_label, indent=4))
                    if i + 1 == len(parameters):
                        parser_code.append(normalize_snippet(parsearg, indent=4))
                    else:
                        add_label = label
                        parser_code.append(normalize_snippet("""
                            if (%s) {{
                            """ % (argname_fmt % i), indent=4))
                        parser_code.append(normalize_snippet(parsearg, indent=8))
                        parser_code.append(normalize_snippet("""
                                if (!--noptargs) {{
                                    goto %s;
                                }}
                            }}
                            """ % add_label, indent=4))

            if parser_code is not None:
                if add_label:
                    parser_code.append("%s:" % add_label)
            else:
                declarations = (
                    'static const char * const _keywords[] = {{{keywords} NULL}};\n'
                    'static _PyArg_Parser _parser = {{"{format_units}:{name}", _keywords, 0}};')
                if not new_or_init:
                    parser_code = [normalize_snippet("""
                        if (!_PyArg_ParseStackAndKeywords(args, nargs, kwnames, &_parser{parse_arguments_comma}
                            {parse_arguments})) {{
                            goto exit;
                        }}
                        """, indent=4)]
                else:
                    parser_code = [normalize_snippet("""
                        if (!_PyArg_ParseTupleAndKeywordsFast(args, kwargs, &_parser,
                            {parse_arguments})) {{
                            goto exit;
                        }}
                        """, indent=4)]
            parser_definition = parser_body(parser_prototype, *parser_code,
                                            declarations=declarations)


        if new_or_init:
            methoddef_define = ''

            if f.kind == METHOD_NEW:
                parser_prototype = parser_prototype_keyword
            else:
                return_value_declaration = "int return_value = -1;"
                parser_prototype = normalize_snippet("""
                    static int
                    {c_basename}({self_type}{self_name}, PyObject *args, PyObject *kwargs)
                    """)

            fields = list(parser_body_fields)
            parses_positional = 'METH_NOARGS' not in flags
            parses_keywords = 'METH_KEYWORDS' in flags
            if parses_keywords:
                assert parses_positional

            if requires_defining_class:
                raise ValueError("Slot methods cannot access their defining class.")

            if not parses_keywords:
                fields.insert(0, normalize_snippet("""
                    if ({self_type_check}!_PyArg_NoKeywords("{name}", kwargs)) {{
                        goto exit;
                    }}
                    """, indent=4))
                if not parses_positional:
                    fields.insert(0, normalize_snippet("""
                        if ({self_type_check}!_PyArg_NoPositional("{name}", args)) {{
                            goto exit;
                        }}
                        """, indent=4))

            parser_definition = parser_body(parser_prototype, *fields,
                                            declarations=parser_body_declarations)


        if flags in ('METH_NOARGS', 'METH_O', 'METH_VARARGS'):
            methoddef_cast = "(PyCFunction)"
        else:
            methoddef_cast = "(PyCFunction)(void(*)(void))"

        if f.methoddef_flags:
            flags += '|' + f.methoddef_flags

        methoddef_define = methoddef_define.replace('{methoddef_flags}', flags)
        methoddef_define = methoddef_define.replace('{methoddef_cast}', methoddef_cast)

        methoddef_ifndef = ''
        conditional = self.cpp.condition()
        if not conditional:
            cpp_if = cpp_endif = ''
        else:
            cpp_if = "#if " + conditional
            cpp_endif = "#endif /* " + conditional + " */"

            if methoddef_define and f.full_name not in clinic.ifndef_symbols:
                clinic.ifndef_symbols.add(f.full_name)
                methoddef_ifndef = normalize_snippet("""
                    #ifndef {methoddef_name}
                        #define {methoddef_name}
                    #endif /* !defined({methoddef_name}) */
                    """)


        # add ';' to the end of parser_prototype and impl_prototype
        # (they mustn't be None, but they could be an empty string.)
        assert parser_prototype is not None
        if parser_prototype:
            assert not parser_prototype.endswith(';')
            parser_prototype += ';'

        if impl_prototype is None:
            impl_prototype = impl_definition
        if impl_prototype:
            impl_prototype += ";"

        parser_definition = parser_definition.replace("{return_value_declaration}", return_value_declaration)

        d = {
            "docstring_prototype" : docstring_prototype,
            "docstring_definition" : docstring_definition,
            "impl_prototype" : impl_prototype,
            "methoddef_define" : methoddef_define,
            "parser_prototype" : parser_prototype,
            "parser_definition" : parser_definition,
            "impl_definition" : impl_definition,
            "cpp_if" : cpp_if,
            "cpp_endif" : cpp_endif,
            "methoddef_ifndef" : methoddef_ifndef,
        }

        # make sure we didn't forget to assign something,
        # and wrap each non-empty value in \n's
        d2 = {}
        for name, value in d.items():
            assert value is not None, "got a None value for template " + repr(name)
            if value:
                value = '\n' + value + '\n'
            d2[name] = value
        return d2

    @staticmethod
    def group_to_variable_name(group):
        adjective = "left_" if group < 0 else "right_"
        return "group_" + adjective + str(abs(group))

    def render_option_group_parsing(self, f, template_dict):
        # positional only, grouped, optional arguments!
        # can be optional on the left or right.
        # here's an example:
        #
        # [ [ [ A1 A2 ] B1 B2 B3 ] C1 C2 ] D1 D2 D3 [ E1 E2 E3 [ F1 F2 F3 ] ]
        #
        # Here group D are required, and all other groups are optional.
        # (Group D's "group" is actually None.)
        # We can figure out which sets of arguments we have based on
        # how many arguments are in the tuple.
        #
        # Note that you need to count up on both sides.  For example,
        # you could have groups C+D, or C+D+E, or C+D+E+F.
        #
        # What if the number of arguments leads us to an ambiguous result?
        # Clinic prefers groups on the left.  So in the above example,
        # five arguments would map to B+C, not C+D.

        add, output = text_accumulator()
        parameters = list(f.parameters.values())
        if isinstance(parameters[0].converter, self_converter):
            del parameters[0]

        groups = []
        group = None
        left = []
        right = []
        required = []
        last = unspecified

        for p in parameters:
            group_id = p.group
            if group_id != last:
                last = group_id
                group = []
                if group_id < 0:
                    left.append(group)
                elif group_id == 0:
                    group = required
                else:
                    right.append(group)
            group.append(p)

        count_min = sys.maxsize
        count_max = -1

        add("switch (PyTuple_GET_SIZE(args)) {\n")
        for subset in permute_optional_groups(left, required, right):
            count = len(subset)
            count_min = min(count_min, count)
            count_max = max(count_max, count)

            if count == 0:
                add("""    case 0:
        break;
""")
                continue

            group_ids = {p.group for p in subset}  # eliminate duplicates
            d = {}
            d['count'] = count
            d['name'] = f.name
            d['format_units'] = "".join(p.converter.format_unit for p in subset)

            parse_arguments = []
            for p in subset:
                p.converter.parse_argument(parse_arguments)
            d['parse_arguments'] = ", ".join(parse_arguments)

            group_ids.discard(0)
            lines = [self.group_to_variable_name(g) + " = 1;" for g in group_ids]
            lines = "\n".join(lines)

            s = """\
    case {count}:
        if (!PyArg_ParseTuple(args, "{format_units}:{name}", {parse_arguments})) {{
            goto exit;
        }}
        {group_booleans}
        break;
"""
            s = linear_format(s, group_booleans=lines)
            s = s.format_map(d)
            add(s)

        add("    default:\n")
        s = '        PyErr_SetString(PyExc_TypeError, "{} requires {} to {} arguments");\n'
        add(s.format(f.full_name, count_min, count_max))
        add('        goto exit;\n')
        add("}")
        template_dict['option_group_parsing'] = format_escape(output())

    def render_function(self, clinic, f):
        if not f:
            return ""

        add, output = text_accumulator()
        data = CRenderData()

        assert f.parameters, "We should always have a 'self' at this point!"
        parameters = f.render_parameters
        converters = [p.converter for p in parameters]

        templates = self.output_templates(f)

        f_self = parameters[0]
        selfless = parameters[1:]
        assert isinstance(f_self.converter, self_converter), "No self parameter in " + repr(f.full_name) + "!"

        last_group = 0
        first_optional = len(selfless)
        positional = selfless and selfless[-1].is_positional_only()
        new_or_init = f.kind in (METHOD_NEW, METHOD_INIT)
        default_return_converter = (not f.return_converter or
            f.return_converter.type == 'PyObject *')
        has_option_groups = False

        # offset i by -1 because first_optional needs to ignore self
        for i, p in enumerate(parameters, -1):
            c = p.converter

            if (i != -1) and (p.default is not unspecified):
                first_optional = min(first_optional, i)

            # insert group variable
            group = p.group
            if last_group != group:
                last_group = group
                if group:
                    group_name = self.group_to_variable_name(group)
                    data.impl_arguments.append(group_name)
                    data.declarations.append("int " + group_name + " = 0;")
                    data.impl_parameters.append("int " + group_name)
                    has_option_groups = True

            c.render(p, data)

        if has_option_groups and (not positional):
            fail("You cannot use optional groups ('[' and ']')\nunless all parameters are positional-only ('/').")

        # HACK
        # when we're METH_O, but have a custom return converter,
        # we use "impl_parameters" for the parsing function
        # because that works better.  but that means we must
        # suppress actually declaring the impl's parameters
        # as variables in the parsing function.  but since it's
        # METH_O, we have exactly one anyway, so we know exactly
        # where it is.
        if ("METH_O" in templates['methoddef_define'] and
            '{impl_parameters}' in templates['parser_prototype']):
            data.declarations.pop(0)

        template_dict = {}

        full_name = f.full_name
        template_dict['full_name'] = full_name

        if new_or_init:
            name = f.cls.name
        else:
            name = f.name

        template_dict['name'] = name

        if f.c_basename:
            c_basename = f.c_basename
        else:
            fields = full_name.split(".")
            if fields[-1] == '__new__':
                fields.pop()
            c_basename = "_".join(fields)

        template_dict['c_basename'] = c_basename

        methoddef_name = "{}_METHODDEF".format(c_basename.upper())
        template_dict['methoddef_name'] = methoddef_name

        template_dict['docstring'] = self.docstring_for_c_string(f)

        template_dict['self_name'] = template_dict['self_type'] = template_dict['self_type_check'] = ''
        for converter in converters:
            converter.set_template_dict(template_dict)

        f.return_converter.render(f, data)
        template_dict['impl_return_type'] = f.return_converter.type

        template_dict['declarations'] = format_escape("\n".join(data.declarations))
        template_dict['initializers'] = "\n\n".join(data.initializers)
        template_dict['modifications'] = '\n\n'.join(data.modifications)
        template_dict['keywords'] = ' '.join('"' + k + '",' for k in data.keywords)
        template_dict['format_units'] = ''.join(data.format_units)
        template_dict['parse_arguments'] = ', '.join(data.parse_arguments)
        if data.parse_arguments:
            template_dict['parse_arguments_comma'] = ',';
        else:
            template_dict['parse_arguments_comma'] = '';
        template_dict['impl_parameters'] = ", ".join(data.impl_parameters)
        template_dict['impl_arguments'] = ", ".join(data.impl_arguments)
        template_dict['return_conversion'] = format_escape("".join(data.return_conversion).rstrip())
        template_dict['cleanup'] = format_escape("".join(data.cleanup))
        template_dict['return_value'] = data.return_value

        # used by unpack tuple code generator
        ignore_self = -1 if isinstance(converters[0], self_converter) else 0
        unpack_min = first_optional
        unpack_max = len(selfless)
        template_dict['unpack_min'] = str(unpack_min)
        template_dict['unpack_max'] = str(unpack_max)

        if has_option_groups:
            self.render_option_group_parsing(f, template_dict)

        # buffers, not destination
        for name, destination in clinic.destination_buffers.items():
            template = templates[name]
            if has_option_groups:
                template = linear_format(template,
                        option_group_parsing=template_dict['option_group_parsing'])
            template = linear_format(template,
                declarations=template_dict['declarations'],
                return_conversion=template_dict['return_conversion'],
                initializers=template_dict['initializers'],
                modifications=template_dict['modifications'],
                cleanup=template_dict['cleanup'],
                )

            # Only generate the "exit:" label
            # if we have any gotos
            need_exit_label = "goto exit;" in template
            template = linear_format(template,
                exit_label="exit:" if need_exit_label else ''
                )

            s = template.format_map(template_dict)

            # mild hack:
            # reflow long impl declarations
            if name in {"impl_prototype", "impl_definition"}:
                s = wrap_declarations(s)

            if clinic.line_prefix:
                s = indent_all_lines(s, clinic.line_prefix)
            if clinic.line_suffix:
                s = suffix_all_lines(s, clinic.line_suffix)

            destination.append(s)

        return clinic.get_destination('block').dump()




@contextlib.contextmanager
def OverrideStdioWith(stdout):
    saved_stdout = sys.stdout
    sys.stdout = stdout
    try:
        yield
    finally:
        assert sys.stdout is stdout
        sys.stdout = saved_stdout


def create_regex(before, after, word=True, whole_line=True):
    """Create an re object for matching marker lines."""
    group_re = r"\w+" if word else ".+"
    pattern = r'{}({}){}'
    if whole_line:
        pattern = '^' + pattern + '$'
    pattern = pattern.format(re.escape(before), group_re, re.escape(after))
    return re.compile(pattern)


class Block:
    r"""
    Represents a single block of text embedded in
    another file.  If dsl_name is None, the block represents
    verbatim text, raw original text from the file, in
    which case "input" will be the only non-false member.
    If dsl_name is not None, the block represents a Clinic
    block.

    input is always str, with embedded \n characters.
    input represents the original text from the file;
    if it's a Clinic block, it is the original text with
    the body_prefix and redundant leading whitespace removed.

    dsl_name is either str or None.  If str, it's the text
    found on the start line of the block between the square
    brackets.

    signatures is either list or None.  If it's a list,
    it may only contain clinic.Module, clinic.Class, and
    clinic.Function objects.  At the moment it should
    contain at most one of each.

    output is either str or None.  If str, it's the output
    from this block, with embedded '\n' characters.

    indent is either str or None.  It's the leading whitespace
    that was found on every line of input.  (If body_prefix is
    not empty, this is the indent *after* removing the
    body_prefix.)

    preindent is either str or None.  It's the whitespace that
    was found in front of every line of input *before* the
    "body_prefix" (see the Language object).  If body_prefix
    is empty, preindent must always be empty too.

    To illustrate indent and preindent: Assume that '_'
    represents whitespace.  If the block processed was in a
    Python file, and looked like this:
      ____#/*[python]
      ____#__for a in range(20):
      ____#____print(a)
      ____#[python]*/
    "preindent" would be "____" and "indent" would be "__".

    """
    def __init__(self, input, dsl_name=None, signatures=None, output=None, indent='', preindent=''):
        assert isinstance(input, str)
        self.input = input
        self.dsl_name = dsl_name
        self.signatures = signatures or []
        self.output = output
        self.indent = indent
        self.preindent = preindent

    def __repr__(self):
        dsl_name = self.dsl_name or "text"
        def summarize(s):
            s = repr(s)
            if len(s) > 30:
                return s[:26] + "..." + s[0]
            return s
        return "".join((
            "<Block ", dsl_name, " input=", summarize(self.input), " output=", summarize(self.output), ">"))


class BlockParser:
    """
    Block-oriented parser for Argument Clinic.
    Iterator, yields Block objects.
    """

    def __init__(self, input, language, *, verify=True):
        """
        "input" should be a str object
        with embedded \n characters.

        "language" should be a Language object.
        """
        language.validate()

        self.input = collections.deque(reversed(input.splitlines(keepends=True)))
        self.block_start_line_number = self.line_number = 0

        self.language = language
        before, _, after = language.start_line.partition('{dsl_name}')
        assert _ == '{dsl_name}'
        self.find_start_re = create_regex(before, after, whole_line=False)
        self.start_re = create_regex(before, after)
        self.verify = verify
        self.last_checksum_re = None
        self.last_dsl_name = None
        self.dsl_name = None
        self.first_block = True

    def __iter__(self):
        return self

    def __next__(self):
        while True:
            if not self.input:
                raise StopIteration

            if self.dsl_name:
                return_value = self.parse_clinic_block(self.dsl_name)
                self.dsl_name = None
                self.first_block = False
                return return_value
            block = self.parse_verbatim_block()
            if self.first_block and not block.input:
                continue
            self.first_block = False
            return block


    def is_start_line(self, line):
        match = self.start_re.match(line.lstrip())
        return match.group(1) if match else None

    def _line(self, lookahead=False):
        self.line_number += 1
        line = self.input.pop()
        if not lookahead:
            self.language.parse_line(line)
        return line

    def parse_verbatim_block(self):
        add, output = text_accumulator()
        self.block_start_line_number = self.line_number

        while self.input:
            line = self._line()
            dsl_name = self.is_start_line(line)
            if dsl_name:
                self.dsl_name = dsl_name
                break
            add(line)

        return Block(output())

    def parse_clinic_block(self, dsl_name):
        input_add, input_output = text_accumulator()
        self.block_start_line_number = self.line_number + 1
        stop_line = self.language.stop_line.format(dsl_name=dsl_name)
        body_prefix = self.language.body_prefix.format(dsl_name=dsl_name)

        def is_stop_line(line):
            # make sure to recognize stop line even if it
            # doesn't end with EOL (it could be the very end of the file)
            if not line.startswith(stop_line):
                return False
            remainder = line[len(stop_line):]
            return (not remainder) or remainder.isspace()

        # consume body of program
        while self.input:
            line = self._line()
            if is_stop_line(line) or self.is_start_line(line):
                break
            if body_prefix:
                line = line.lstrip()
                assert line.startswith(body_prefix)
                line = line[len(body_prefix):]
            input_add(line)

        # consume output and checksum line, if present.
        if self.last_dsl_name == dsl_name:
            checksum_re = self.last_checksum_re
        else:
            before, _, after = self.language.checksum_line.format(dsl_name=dsl_name, arguments='{arguments}').partition('{arguments}')
            assert _ == '{arguments}'
            checksum_re = create_regex(before, after, word=False)
            self.last_dsl_name = dsl_name
            self.last_checksum_re = checksum_re

        # scan forward for checksum line
        output_add, output_output = text_accumulator()
        arguments = None
        while self.input:
            line = self._line(lookahead=True)
            match = checksum_re.match(line.lstrip())
            arguments = match.group(1) if match else None
            if arguments:
                break
            output_add(line)
            if self.is_start_line(line):
                break

        output = output_output()
        if arguments:
            d = {}
            for field in shlex.split(arguments):
                name, equals, value = field.partition('=')
                if not equals:
                    fail("Mangled Argument Clinic marker line: {!r}".format(line))
                d[name.strip()] = value.strip()

            if self.verify:
                if 'input' in d:
                    checksum = d['output']
                    input_checksum = d['input']
                else:
                    checksum = d['checksum']
                    input_checksum = None

                computed = compute_checksum(output, len(checksum))
                if checksum != computed:
                    fail("Checksum mismatch!\nExpected: {}\nComputed: {}\n"
                         "Suggested fix: remove all generated code including "
                         "the end marker,\n"
                         "or use the '-f' option."
                        .format(checksum, computed))
        else:
            # put back output
            output_lines = output.splitlines(keepends=True)
            self.line_number -= len(output_lines)
            self.input.extend(reversed(output_lines))
            output = None

        return Block(input_output(), dsl_name, output=output)


class BlockPrinter:

    def __init__(self, language, f=None):
        self.language = language
        self.f = f or io.StringIO()

    def print_block(self, block):
        input = block.input
        output = block.output
        dsl_name = block.dsl_name
        write = self.f.write

        assert not ((dsl_name == None) ^ (output == None)), "you must specify dsl_name and output together, dsl_name " + repr(dsl_name)

        if not dsl_name:
            write(input)
            return

        write(self.language.start_line.format(dsl_name=dsl_name))
        write("\n")

        body_prefix = self.language.body_prefix.format(dsl_name=dsl_name)
        if not body_prefix:
            write(input)
        else:
            for line in input.split('\n'):
                write(body_prefix)
                write(line)
                write("\n")

        write(self.language.stop_line.format(dsl_name=dsl_name))
        write("\n")

        input = ''.join(block.input)
        output = ''.join(block.output)
        if output:
            if not output.endswith('\n'):
                output += '\n'
            write(output)

        arguments="output={} input={}".format(compute_checksum(output, 16), compute_checksum(input, 16))
        write(self.language.checksum_line.format(dsl_name=dsl_name, arguments=arguments))
        write("\n")

    def write(self, text):
        self.f.write(text)


class BufferSeries:
    """
    Behaves like a "defaultlist".
    When you ask for an index that doesn't exist yet,
    the object grows the list until that item exists.
    So o[n] will always work.

    Supports negative indices for actual items.
    e.g. o[-1] is an element immediately preceding o[0].
    """

    def __init__(self):
        self._start = 0
        self._array = []
        self._constructor = _text_accumulator

    def __getitem__(self, i):
        i -= self._start
        if i < 0:
            self._start += i
            prefix = [self._constructor() for x in range(-i)]
            self._array = prefix + self._array
            i = 0
        while i >= len(self._array):
            self._array.append(self._constructor())
        return self._array[i]

    def clear(self):
        for ta in self._array:
            ta._text.clear()

    def dump(self):
        texts = [ta.output() for ta in self._array]
        return "".join(texts)


class Destination:
    def __init__(self, name, type, clinic, *args):
        self.name = name
        self.type = type
        self.clinic = clinic
        valid_types = ('buffer', 'file', 'suppress')
        if type not in valid_types:
            fail("Invalid destination type " + repr(type) + " for " + name + " , must be " + ', '.join(valid_types))
        extra_arguments = 1 if type == "file" else 0
        if len(args) < extra_arguments:
            fail("Not enough arguments for destination " + name + " new " + type)
        if len(args) > extra_arguments:
            fail("Too many arguments for destination " + name + " new " + type)
        if type =='file':
            d = {}
            filename = clinic.filename
            d['path'] = filename
            dirname, basename = os.path.split(filename)
            if not dirname:
                dirname = '.'
            d['dirname'] = dirname
            d['basename'] = basename
            d['basename_root'], d['basename_extension'] = os.path.splitext(filename)
            self.filename = args[0].format_map(d)

        self.buffers = BufferSeries()

    def __repr__(self):
        if self.type == 'file':
            file_repr = " " + repr(self.filename)
        else:
            file_repr = ''
        return "".join(("<Destination ", self.name, " ", self.type, file_repr, ">"))

    def clear(self):
        if self.type != 'buffer':
            fail("Can't clear destination" + self.name + " , it's not of type buffer")
        self.buffers.clear()

    def dump(self):
        return self.buffers.dump()


# maps strings to Language objects.
# "languages" maps the name of the language ("C", "Python").
# "extensions" maps the file extension ("c", "py").
languages = { 'C': CLanguage, 'Python': PythonLanguage }
extensions = { name: CLanguage for name in "c cc cpp cxx h hh hpp hxx".split() }
extensions['py'] = PythonLanguage


# maps strings to callables.
# these callables must be of the form:
#   def foo(name, default, *, ...)
# The callable may have any number of keyword-only parameters.
# The callable must return a CConverter object.
# The callable should not call builtins.print.
converters = {}

# maps strings to callables.
# these callables follow the same rules as those for "converters" above.
# note however that they will never be called with keyword-only parameters.
legacy_converters = {}


# maps strings to callables.
# these callables must be of the form:
#   def foo(*, ...)
# The callable may have any number of keyword-only parameters.
# The callable must return a CConverter object.
# The callable should not call builtins.print.
return_converters = {}


def write_file(filename, new_contents):
    try:
        with open(filename, 'r', encoding="utf-8") as fp:
            old_contents = fp.read()

        if old_contents == new_contents:
            # no change: avoid modifying the file modification time
            return
    except FileNotFoundError:
        pass

    # Atomic write using a temporary file and os.replace()
    filename_new = f"{filename}.new"
    with open(filename_new, "w", encoding="utf-8") as fp:
        fp.write(new_contents)

    try:
        os.replace(filename_new, filename)
    except:
        os.unlink(filename_new)
        raise


clinic = None
class Clinic:

    presets_text = """
preset block
everything block
methoddef_ifndef buffer 1
docstring_prototype suppress
parser_prototype suppress
cpp_if suppress
cpp_endif suppress

preset original
everything block
methoddef_ifndef buffer 1
docstring_prototype suppress
parser_prototype suppress
cpp_if suppress
cpp_endif suppress

preset file
everything file
methoddef_ifndef file 1
docstring_prototype suppress
parser_prototype suppress
impl_definition block

preset buffer
everything buffer
methoddef_ifndef buffer 1
impl_definition block
docstring_prototype suppress
impl_prototype suppress
parser_prototype suppress

preset partial-buffer
everything buffer
methoddef_ifndef buffer 1
docstring_prototype block
impl_prototype suppress
methoddef_define block
parser_prototype block
impl_definition block

"""

    def __init__(self, language, printer=None, *, verify=True, filename=None):
        # maps strings to Parser objects.
        # (instantiated from the "parsers" global.)
        self.parsers = {}
        self.language = language
        if printer:
            fail("Custom printers are broken right now")
        self.printer = printer or BlockPrinter(language)
        self.verify = verify
        self.filename = filename
        self.modules = collections.OrderedDict()
        self.classes = collections.OrderedDict()
        self.functions = []

        self.line_prefix = self.line_suffix = ''

        self.destinations = {}
        self.add_destination("block", "buffer")
        self.add_destination("suppress", "suppress")
        self.add_destination("buffer", "buffer")
        if filename:
            self.add_destination("file", "file", "{dirname}/clinic/{basename}.h")

        d = self.get_destination_buffer
        self.destination_buffers = collections.OrderedDict((
            ('cpp_if', d('file')),
            ('docstring_prototype', d('suppress')),
            ('docstring_definition', d('file')),
            ('methoddef_define', d('file')),
            ('impl_prototype', d('file')),
            ('parser_prototype', d('suppress')),
            ('parser_definition', d('file')),
            ('cpp_endif', d('file')),
            ('methoddef_ifndef', d('file', 1)),
            ('impl_definition', d('block')),
        ))

        self.destination_buffers_stack = []
        self.ifndef_symbols = set()

        self.presets = {}
        preset = None
        for line in self.presets_text.strip().split('\n'):
            line = line.strip()
            if not line:
                continue
            name, value, *options = line.split()
            if name == 'preset':
                self.presets[value] = preset = collections.OrderedDict()
                continue

            if len(options):
                index = int(options[0])
            else:
                index = 0
            buffer = self.get_destination_buffer(value, index)

            if name == 'everything':
                for name in self.destination_buffers:
                    preset[name] = buffer
                continue

            assert name in self.destination_buffers
            preset[name] = buffer

        global clinic
        clinic = self

    def add_destination(self, name, type, *args):
        if name in self.destinations:
            fail("Destination already exists: " + repr(name))
        self.destinations[name] = Destination(name, type, self, *args)

    def get_destination(self, name):
        d = self.destinations.get(name)
        if not d:
            fail("Destination does not exist: " + repr(name))
        return d

    def get_destination_buffer(self, name, item=0):
        d = self.get_destination(name)
        return d.buffers[item]

    def parse(self, input):
        printer = self.printer
        self.block_parser = BlockParser(input, self.language, verify=self.verify)
        for block in self.block_parser:
            dsl_name = block.dsl_name
            if dsl_name:
                if dsl_name not in self.parsers:
                    assert dsl_name in parsers, "No parser to handle {!r} block.".format(dsl_name)
                    self.parsers[dsl_name] = parsers[dsl_name](self)
                parser = self.parsers[dsl_name]
                try:
                    parser.parse(block)
                except Exception:
                    fail('Exception raised during parsing:\n' +
                         traceback.format_exc().rstrip())
            printer.print_block(block)

        second_pass_replacements = {}

        # these are destinations not buffers
        for name, destination in self.destinations.items():
            if destination.type == 'suppress':
                continue
            output = destination.dump()

            if output:

                block = Block("", dsl_name="clinic", output=output)

                if destination.type == 'buffer':
                    block.input = "dump " + name + "\n"
                    warn("Destination buffer " + repr(name) + " not empty at end of file, emptying.")
                    printer.write("\n")
                    printer.print_block(block)
                    continue

                if destination.type == 'file':
                    try:
                        dirname = os.path.dirname(destination.filename)
                        try:
                            os.makedirs(dirname)
                        except FileExistsError:
                            if not os.path.isdir(dirname):
                                fail("Can't write to destination {}, "
                                     "can't make directory {}!".format(
                                        destination.filename, dirname))
                        if self.verify:
                            with open(destination.filename, "rt") as f:
                                parser_2 = BlockParser(f.read(), language=self.language)
                                blocks = list(parser_2)
                                if (len(blocks) != 1) or (blocks[0].input != 'preserve\n'):
                                    fail("Modified destination file " + repr(destination.filename) + ", not overwriting!")
                    except FileNotFoundError:
                        pass

                    block.input = 'preserve\n'
                    printer_2 = BlockPrinter(self.language)
                    printer_2.print_block(block)
                    write_file(destination.filename, printer_2.f.getvalue())
                    continue
        text = printer.f.getvalue()

        if second_pass_replacements:
            printer_2 = BlockPrinter(self.language)
            parser_2 = BlockParser(text, self.language)
            changed = False
            for block in parser_2:
                if block.dsl_name:
                    for id, replacement in second_pass_replacements.items():
                        if id in block.output:
                            changed = True
                            block.output = block.output.replace(id, replacement)
                printer_2.print_block(block)
            if changed:
                text = printer_2.f.getvalue()

        return text


    def _module_and_class(self, fields):
        """
        fields should be an iterable of field names.
        returns a tuple of (module, class).
        the module object could actually be self (a clinic object).
        this function is only ever used to find the parent of where
        a new class/module should go.
        """
        in_classes = False
        parent = module = self
        cls = None
        so_far = []

        for field in fields:
            so_far.append(field)
            if not in_classes:
                child = parent.modules.get(field)
                if child:
                    parent = module = child
                    continue
                in_classes = True
            if not hasattr(parent, 'classes'):
                return module, cls
            child = parent.classes.get(field)
            if not child:
                fail('Parent class or module ' + '.'.join(so_far) + " does not exist.")
            cls = parent = child

        return module, cls


def parse_file(filename, *, verify=True, output=None):
    if not output:
        output = filename

    extension = os.path.splitext(filename)[1][1:]
    if not extension:
        fail("Can't extract file type for file " + repr(filename))

    try:
        language = extensions[extension](filename)
    except KeyError:
        fail("Can't identify file type for file " + repr(filename))

    with open(filename, 'r', encoding="utf-8") as f:
        raw = f.read()

    # exit quickly if there are no clinic markers in the file
    find_start_re = BlockParser("", language).find_start_re
    if not find_start_re.search(raw):
        return

    clinic = Clinic(language, verify=verify, filename=filename)
    cooked = clinic.parse(raw)

    write_file(output, cooked)


def compute_checksum(input, length=None):
    input = input or ''
    s = hashlib.sha1(input.encode('utf-8')).hexdigest()
    if length:
        s = s[:length]
    return s




class PythonParser:
    def __init__(self, clinic):
        pass

    def parse(self, block):
        s = io.StringIO()
        with OverrideStdioWith(s):
            exec(block.input)
        block.output = s.getvalue()


class Module:
    def __init__(self, name, module=None):
        self.name = name
        self.module = self.parent = module

        self.modules = collections.OrderedDict()
        self.classes = collections.OrderedDict()
        self.functions = []

    def __repr__(self):
        return "<clinic.Module " + repr(self.name) + " at " + str(id(self)) + ">"

class Class:
    def __init__(self, name, module=None, cls=None, typedef=None, type_object=None):
        self.name = name
        self.module = module
        self.cls = cls
        self.typedef = typedef
        self.type_object = type_object
        self.parent = cls or module

        self.classes = collections.OrderedDict()
        self.functions = []

    def __repr__(self):
        return "<clinic.Class " + repr(self.name) + " at " + str(id(self)) + ">"

unsupported_special_methods = set("""

__abs__
__add__
__and__
__bytes__
__call__
__complex__
__delitem__
__divmod__
__eq__
__float__
__floordiv__
__ge__
__getattr__
__getattribute__
__getitem__
__gt__
__hash__
__iadd__
__iand__
__ifloordiv__
__ilshift__
__imatmul__
__imod__
__imul__
__index__
__int__
__invert__
__ior__
__ipow__
__irshift__
__isub__
__iter__
__itruediv__
__ixor__
__le__
__len__
__lshift__
__lt__
__matmul__
__mod__
__mul__
__neg__
__new__
__next__
__or__
__pos__
__pow__
__radd__
__rand__
__rdivmod__
__repr__
__rfloordiv__
__rlshift__
__rmatmul__
__rmod__
__rmul__
__ror__
__rpow__
__rrshift__
__rshift__
__rsub__
__rtruediv__
__rxor__
__setattr__
__setitem__
__str__
__sub__
__truediv__
__xor__

""".strip().split())


INVALID, CALLABLE, STATIC_METHOD, CLASS_METHOD, METHOD_INIT, METHOD_NEW = """
INVALID, CALLABLE, STATIC_METHOD, CLASS_METHOD, METHOD_INIT, METHOD_NEW
""".replace(",", "").strip().split()

class Function:
    """
    Mutable duck type for inspect.Function.

    docstring - a str containing
        * embedded line breaks
        * text outdented to the left margin
        * no trailing whitespace.
        It will always be true that
            (not docstring) or ((not docstring[0].isspace()) and (docstring.rstrip() == docstring))
    """

    def __init__(self, parameters=None, *, name,
                 module, cls=None, c_basename=None,
                 full_name=None,
                 return_converter, return_annotation=inspect.Signature.empty,
                 docstring=None, kind=CALLABLE, coexist=False,
                 docstring_only=False):
        self.parameters = parameters or collections.OrderedDict()
        self.return_annotation = return_annotation
        self.name = name
        self.full_name = full_name
        self.module = module
        self.cls = cls
        self.parent = cls or module
        self.c_basename = c_basename
        self.return_converter = return_converter
        self.docstring = docstring or ''
        self.kind = kind
        self.coexist = coexist
        self.self_converter = None
        # docstring_only means "don't generate a machine-readable
        # signature, just a normal docstring".  it's True for
        # functions with optional groups because we can't represent
        # those accurately with inspect.Signature in 3.4.
        self.docstring_only = docstring_only

        self.rendered_parameters = None

    __render_parameters__ = None
    @property
    def render_parameters(self):
        if not self.__render_parameters__:
            self.__render_parameters__ = l = []
            for p in self.parameters.values():
                p = p.copy()
                p.converter.pre_render()
                l.append(p)
        return self.__render_parameters__

    @property
    def methoddef_flags(self):
        if self.kind in (METHOD_INIT, METHOD_NEW):
            return None
        flags = []
        if self.kind == CLASS_METHOD:
            flags.append('METH_CLASS')
        elif self.kind == STATIC_METHOD:
            flags.append('METH_STATIC')
        else:
            assert self.kind == CALLABLE, "unknown kind: " + repr(self.kind)
        if self.coexist:
            flags.append('METH_COEXIST')
        return '|'.join(flags)

    def __repr__(self):
        return '<clinic.Function ' + self.name + '>'

    def copy(self, **overrides):
        kwargs = {
            'name': self.name, 'module': self.module, 'parameters': self.parameters,
            'cls': self.cls, 'c_basename': self.c_basename,
            'full_name': self.full_name,
            'return_converter': self.return_converter, 'return_annotation': self.return_annotation,
            'docstring': self.docstring, 'kind': self.kind, 'coexist': self.coexist,
            'docstring_only': self.docstring_only,
            }
        kwargs.update(overrides)
        f = Function(**kwargs)

        parameters = collections.OrderedDict()
        for name, value in f.parameters.items():
            value = value.copy(function=f)
            parameters[name] = value
        f.parameters = parameters
        return f


class Parameter:
    """
    Mutable duck type of inspect.Parameter.
    """

    def __init__(self, name, kind, *, default=inspect.Parameter.empty,
                 function, converter, annotation=inspect.Parameter.empty,
                 docstring=None, group=0):
        self.name = name
        self.kind = kind
        self.default = default
        self.function = function
        self.converter = converter
        self.annotation = annotation
        self.docstring = docstring or ''
        self.group = group

    def __repr__(self):
        return '<clinic.Parameter ' + self.name + '>'

    def is_keyword_only(self):
        return self.kind == inspect.Parameter.KEYWORD_ONLY

    def is_positional_only(self):
        return self.kind == inspect.Parameter.POSITIONAL_ONLY

    def is_optional(self):
        return (self.default is not unspecified)

    def copy(self, **overrides):
        kwargs = {
            'name': self.name, 'kind': self.kind, 'default':self.default,
                 'function': self.function, 'converter': self.converter, 'annotation': self.annotation,
                 'docstring': self.docstring, 'group': self.group,
            }
        kwargs.update(overrides)
        if 'converter' not in overrides:
            converter = copy.copy(self.converter)
            converter.function = kwargs['function']
            kwargs['converter'] = converter
        return Parameter(**kwargs)

    def get_displayname(self, i):
        if i == 0:
            return '"argument"'
        if not self.is_positional_only():
            return '''"argument '{}'"'''.format(self.name)
        else:
            return '"argument {}"'.format(i)


class LandMine:
    # try to access any
    def __init__(self, message):
        self.__message__ = message

    def __repr__(self):
        return '<LandMine ' + repr(self.__message__) + ">"

    def __getattribute__(self, name):
        if name in ('__repr__', '__message__'):
            return super().__getattribute__(name)
        # raise RuntimeError(repr(name))
        fail("Stepped on a land mine, trying to access attribute " + repr(name) + ":\n" + self.__message__)


def add_c_converter(f, name=None):
    if not name:
        name = f.__name__
        if not name.endswith('_converter'):
            return f
        name = name[:-len('_converter')]
    converters[name] = f
    return f

def add_default_legacy_c_converter(cls):
    # automatically add converter for default format unit
    # (but without stomping on the existing one if it's already
    # set, in case you subclass)
    if ((cls.format_unit not in ('O&', '')) and
        (cls.format_unit not in legacy_converters)):
        legacy_converters[cls.format_unit] = cls
    return cls

def add_legacy_c_converter(format_unit, **kwargs):
    """
    Adds a legacy converter.
    """
    def closure(f):
        if not kwargs:
            added_f = f
        else:
            added_f = functools.partial(f, **kwargs)
        if format_unit:
            legacy_converters[format_unit] = added_f
        return f
    return closure

class CConverterAutoRegister(type):
    def __init__(cls, name, bases, classdict):
        add_c_converter(cls)
        add_default_legacy_c_converter(cls)

class CConverter(metaclass=CConverterAutoRegister):
    """
    For the init function, self, name, function, and default
    must be keyword-or-positional parameters.  All other
    parameters must be keyword-only.
    """

    # The C name to use for this variable.
    name = None

    # The Python name to use for this variable.
    py_name = None

    # The C type to use for this variable.
    # 'type' should be a Python string specifying the type, e.g. "int".
    # If this is a pointer type, the type string should end with ' *'.
    type = None

    # The Python default value for this parameter, as a Python value.
    # Or the magic value "unspecified" if there is no default.
    # Or the magic value "unknown" if this value is a cannot be evaluated
    # at Argument-Clinic-preprocessing time (but is presumed to be valid
    # at runtime).
    default = unspecified

    # If not None, default must be isinstance() of this type.
    # (You can also specify a tuple of types.)
    default_type = None

    # "default" converted into a C value, as a string.
    # Or None if there is no default.
    c_default = None

    # "default" converted into a Python value, as a string.
    # Or None if there is no default.
    py_default = None

    # The default value used to initialize the C variable when
    # there is no default, but not specifying a default may
    # result in an "uninitialized variable" warning.  This can
    # easily happen when using option groups--although
    # properly-written code won't actually use the variable,
    # the variable does get passed in to the _impl.  (Ah, if
    # only dataflow analysis could inline the static function!)
    #
    # This value is specified as a string.
    # Every non-abstract subclass should supply a valid value.
    c_ignored_default = 'NULL'

    # The C converter *function* to be used, if any.
    # (If this is not None, format_unit must be 'O&'.)
    converter = None

    # Should Argument Clinic add a '&' before the name of
    # the variable when passing it into the _impl function?
    impl_by_reference = False

    # Should Argument Clinic add a '&' before the name of
    # the variable when passing it into PyArg_ParseTuple (AndKeywords)?
    parse_by_reference = True

    #############################################################
    #############################################################
    ## You shouldn't need to read anything below this point to ##
    ## write your own converter functions.                     ##
    #############################################################
    #############################################################

    # The "format unit" to specify for this variable when
    # parsing arguments using PyArg_ParseTuple (AndKeywords).
    # Custom converters should always use the default value of 'O&'.
    format_unit = 'O&'

    # What encoding do we want for this variable?  Only used
    # by format units starting with 'e'.
    encoding = None

    # Should this object be required to be a subclass of a specific type?
    # If not None, should be a string representing a pointer to a
    # PyTypeObject (e.g. "&PyUnicode_Type").
    # Only used by the 'O!' format unit (and the "object" converter).
    subclass_of = None

    # Do we want an adjacent '_length' variable for this variable?
    # Only used by format units ending with '#'.
    length = False

    # Should we show this parameter in the generated
    # __text_signature__? This is *almost* always True.
    # (It's only False for __new__, __init__, and METH_STATIC functions.)
    show_in_signature = True

    # Overrides the name used in a text signature.
    # The name used for a "self" parameter must be one of
    # self, type, or module; however users can set their own.
    # This lets the self_converter overrule the user-settable
    # name, *just* for the text signature.
    # Only set by self_converter.
    signature_name = None

    # keep in sync with self_converter.__init__!
    def __init__(self, name, py_name, function, default=unspecified, *, c_default=None, py_default=None, annotation=unspecified, **kwargs):
        self.name = ensure_legal_c_identifier(name)
        self.py_name = py_name

        if default is not unspecified:
            if self.default_type and not isinstance(default, (self.default_type, Unknown)):
                if isinstance(self.default_type, type):
                    types_str = self.default_type.__name__
                else:
                    types_str = ', '.join((cls.__name__ for cls in self.default_type))
                fail("{}: default value {!r} for field {} is not of type {}".format(
                    self.__class__.__name__, default, name, types_str))
            self.default = default

        if c_default:
            self.c_default = c_default
        if py_default:
            self.py_default = py_default

        if annotation != unspecified:
            fail("The 'annotation' parameter is not currently permitted.")

        # this is deliberate, to prevent you from caching information
        # about the function in the init.
        # (that breaks if we get cloned.)
        # so after this change we will noisily fail.
        self.function = LandMine("Don't access members of self.function inside converter_init!")
        self.converter_init(**kwargs)
        self.function = function

    def converter_init(self):
        pass

    def is_optional(self):
        return (self.default is not unspecified)

    def _render_self(self, parameter, data):
        self.parameter = parameter
        name = self.name

        # impl_arguments
        s = ("&" if self.impl_by_reference else "") + name
        data.impl_arguments.append(s)
        if self.length:
            data.impl_arguments.append(self.length_name())

        # impl_parameters
        data.impl_parameters.append(self.simple_declaration(by_reference=self.impl_by_reference))
        if self.length:
            data.impl_parameters.append("Py_ssize_clean_t " + self.length_name())

    def _render_non_self(self, parameter, data):
        self.parameter = parameter
        name = self.name

        # declarations
        d = self.declaration()
        data.declarations.append(d)

        # initializers
        initializers = self.initialize()
        if initializers:
            data.initializers.append('/* initializers for ' + name + ' */\n' + initializers.rstrip())

        # modifications
        modifications = self.modify()
        if modifications:
            data.modifications.append('/* modifications for ' + name + ' */\n' + modifications.rstrip())

        # keywords
        if parameter.is_positional_only():
            data.keywords.append('')
        else:
            data.keywords.append(parameter.name)

        # format_units
        if self.is_optional() and '|' not in data.format_units:
            data.format_units.append('|')
        if parameter.is_keyword_only() and '$' not in data.format_units:
            data.format_units.append('$')
        data.format_units.append(self.format_unit)

        # parse_arguments
        self.parse_argument(data.parse_arguments)

        # cleanup
        cleanup = self.cleanup()
        if cleanup:
            data.cleanup.append('/* Cleanup for ' + name + ' */\n' + cleanup.rstrip() + "\n")

    def render(self, parameter, data):
        """
        parameter is a clinic.Parameter instance.
        data is a CRenderData instance.
        """
        self._render_self(parameter, data)
        self._render_non_self(parameter, data)

    def length_name(self):
        """Computes the name of the associated "length" variable."""
        if not self.length:
            return None
        return self.name + "_length"

    # Why is this one broken out separately?
    # For "positional-only" function parsing,
    # which generates a bunch of PyArg_ParseTuple calls.
    def parse_argument(self, list):
        assert not (self.converter and self.encoding)
        if self.format_unit == 'O&':
            assert self.converter
            list.append(self.converter)

        if self.encoding:
            list.append(c_repr(self.encoding))
        elif self.subclass_of:
            list.append(self.subclass_of)

        s = ("&" if self.parse_by_reference else "") + self.name
        list.append(s)

        if self.length:
            list.append("&" + self.length_name())

    #
    # All the functions after here are intended as extension points.
    #

    def simple_declaration(self, by_reference=False):
        """
        Computes the basic declaration of the variable.
        Used in computing the prototype declaration and the
        variable declaration.
        """
        prototype = [self.type]
        if by_reference or not self.type.endswith('*'):
            prototype.append(" ")
        if by_reference:
            prototype.append('*')
        prototype.append(self.name)
        return "".join(prototype)

    def declaration(self):
        """
        The C statement to declare this variable.
        """
        declaration = [self.simple_declaration()]
        default = self.c_default
        if not default and self.parameter.group:
            default = self.c_ignored_default
        if default:
            declaration.append(" = ")
            declaration.append(default)
        declaration.append(";")
        if self.length:
            declaration.append('\nPy_ssize_clean_t ')
            declaration.append(self.length_name())
            declaration.append(';')
        return "".join(declaration)

    def initialize(self):
        """
        The C statements required to set up this variable before parsing.
        Returns a string containing this code indented at column 0.
        If no initialization is necessary, returns an empty string.
        """
        return ""

    def modify(self):
        """
        The C statements required to modify this variable after parsing.
        Returns a string containing this code indented at column 0.
        If no initialization is necessary, returns an empty string.
        """
        return ""

    def cleanup(self):
        """
        The C statements required to clean up after this variable.
        Returns a string containing this code indented at column 0.
        If no cleanup is necessary, returns an empty string.
        """
        return ""

    def pre_render(self):
        """
        A second initialization function, like converter_init,
        called just before rendering.
        You are permitted to examine self.function here.
        """
        pass

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'O&':
            return """
                if (!{converter}({argname}, &{paramname})) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name,
                           converter=self.converter)
        if self.format_unit == 'O!':
            cast = '(%s)' % self.type if self.type != 'PyObject *' else ''
            if self.subclass_of in type_checks:
                typecheck, typename = type_checks[self.subclass_of]
                return """
                    if (!{typecheck}({argname})) {{{{
                        _PyArg_BadArgument("{{name}}", {displayname}, "{typename}", {argname});
                        goto exit;
                    }}}}
                    {paramname} = {cast}{argname};
                    """.format(argname=argname, paramname=self.name,
                               displayname=displayname, typecheck=typecheck,
                               typename=typename, cast=cast)
            return """
                if (!PyObject_TypeCheck({argname}, {subclass_of})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, ({subclass_of})->tp_name, {argname});
                    goto exit;
                }}}}
                {paramname} = {cast}{argname};
                """.format(argname=argname, paramname=self.name,
                           subclass_of=self.subclass_of, cast=cast,
                           displayname=displayname)
        if self.format_unit == 'O':
            cast = '(%s)' % self.type if self.type != 'PyObject *' else ''
            return """
                {paramname} = {cast}{argname};
                """.format(argname=argname, paramname=self.name, cast=cast)
        return None

    def set_template_dict(self, template_dict):
        pass


type_checks = {
    '&PyLong_Type': ('PyLong_Check', 'int'),
    '&PyTuple_Type': ('PyTuple_Check', 'tuple'),
    '&PyList_Type': ('PyList_Check', 'list'),
    '&PySet_Type': ('PySet_Check', 'set'),
    '&PyFrozenSet_Type': ('PyFrozenSet_Check', 'frozenset'),
    '&PyDict_Type': ('PyDict_Check', 'dict'),
    '&PyUnicode_Type': ('PyUnicode_Check', 'str'),
    '&PyBytes_Type': ('PyBytes_Check', 'bytes'),
    '&PyByteArray_Type': ('PyByteArray_Check', 'bytearray'),
}


class bool_converter(CConverter):
    type = 'int'
    default_type = bool
    format_unit = 'p'
    c_ignored_default = '0'

    def converter_init(self, *, accept={object}):
        if accept == {int}:
            self.format_unit = 'i'
        elif accept != {object}:
            fail("bool_converter: illegal 'accept' argument " + repr(accept))
        if self.default is not unspecified:
            self.default = bool(self.default)
            self.c_default = str(int(self.default))

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'i':
            # XXX PyFloat_Check can be removed after the end of the
            # deprecation in _PyLong_FromNbIndexOrNbInt.
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {paramname} = _PyLong_AsInt({argname});
                if ({paramname} == -1 && PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        elif self.format_unit == 'p':
            return """
                {paramname} = PyObject_IsTrue({argname});
                if ({paramname} < 0) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

class defining_class_converter(CConverter):
    """
    A special-case converter:
    this is the default converter used for the defining class.
    """
    type = 'PyTypeObject *'
    format_unit = ''
    show_in_signature = False

    def converter_init(self, *, type=None):
        self.specified_type = type

    def render(self, parameter, data):
        self._render_self(parameter, data)

    def set_template_dict(self, template_dict):
        template_dict['defining_class_name'] = self.name


class char_converter(CConverter):
    type = 'char'
    default_type = (bytes, bytearray)
    format_unit = 'c'
    c_ignored_default = "'\0'"

    def converter_init(self):
        if isinstance(self.default, self.default_type):
            if len(self.default) != 1:
                fail("char_converter: illegal default value " + repr(self.default))

            self.c_default = repr(bytes(self.default))[1:]
            if self.c_default == '"\'"':
                self.c_default = r"'\''"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'c':
            return """
                if (PyBytes_Check({argname}) && PyBytes_GET_SIZE({argname}) == 1) {{{{
                    {paramname} = PyBytes_AS_STRING({argname})[0];
                }}}}
                else if (PyByteArray_Check({argname}) && PyByteArray_GET_SIZE({argname}) == 1) {{{{
                    {paramname} = PyByteArray_AS_STRING({argname})[0];
                }}}}
                else {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "a byte string of length 1", {argname});
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        return super().parse_arg(argname, displayname)


@add_legacy_c_converter('B', bitwise=True)
class unsigned_char_converter(CConverter):
    type = 'unsigned char'
    default_type = int
    format_unit = 'b'
    c_ignored_default = "'\0'"

    def converter_init(self, *, bitwise=False):
        if bitwise:
            self.format_unit = 'B'

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'b':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {{{{
                    long ival = PyLong_AsLong({argname});
                    if (ival == -1 && PyErr_Occurred()) {{{{
                        goto exit;
                    }}}}
                    else if (ival < 0) {{{{
                        PyErr_SetString(PyExc_OverflowError,
                                        "unsigned byte integer is less than minimum");
                        goto exit;
                    }}}}
                    else if (ival > UCHAR_MAX) {{{{
                        PyErr_SetString(PyExc_OverflowError,
                                        "unsigned byte integer is greater than maximum");
                        goto exit;
                    }}}}
                    else {{{{
                        {paramname} = (unsigned char) ival;
                    }}}}
                }}}}
                """.format(argname=argname, paramname=self.name)
        elif self.format_unit == 'B':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {{{{
                    long ival = PyLong_AsUnsignedLongMask({argname});
                    if (ival == -1 && PyErr_Occurred()) {{{{
                        goto exit;
                    }}}}
                    else {{{{
                        {paramname} = (unsigned char) ival;
                    }}}}
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

class byte_converter(unsigned_char_converter): pass

class short_converter(CConverter):
    type = 'short'
    default_type = int
    format_unit = 'h'
    c_ignored_default = "0"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'h':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {{{{
                    long ival = PyLong_AsLong({argname});
                    if (ival == -1 && PyErr_Occurred()) {{{{
                        goto exit;
                    }}}}
                    else if (ival < SHRT_MIN) {{{{
                        PyErr_SetString(PyExc_OverflowError,
                                        "signed short integer is less than minimum");
                        goto exit;
                    }}}}
                    else if (ival > SHRT_MAX) {{{{
                        PyErr_SetString(PyExc_OverflowError,
                                        "signed short integer is greater than maximum");
                        goto exit;
                    }}}}
                    else {{{{
                        {paramname} = (short) ival;
                    }}}}
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

class unsigned_short_converter(CConverter):
    type = 'unsigned short'
    default_type = int
    c_ignored_default = "0"

    def converter_init(self, *, bitwise=False):
        if bitwise:
            self.format_unit = 'H'
        else:
            self.converter = '_PyLong_UnsignedShort_Converter'

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'H':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {paramname} = (unsigned short)PyLong_AsUnsignedLongMask({argname});
                if ({paramname} == (unsigned short)-1 && PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

@add_legacy_c_converter('C', accept={str})
class int_converter(CConverter):
    type = 'int'
    default_type = int
    format_unit = 'i'
    c_ignored_default = "0"

    def converter_init(self, *, accept={int}, type=None):
        if accept == {str}:
            self.format_unit = 'C'
        elif accept != {int}:
            fail("int_converter: illegal 'accept' argument " + repr(accept))
        if type != None:
            self.type = type

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'i':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {paramname} = _PyLong_AsInt({argname});
                if ({paramname} == -1 && PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        elif self.format_unit == 'C':
            return """
                if (!PyUnicode_Check({argname})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "a unicode character", {argname});
                    goto exit;
                }}}}
                if (PyUnicode_READY({argname})) {{{{
                    goto exit;
                }}}}
                if (PyUnicode_GET_LENGTH({argname}) != 1) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "a unicode character", {argname});
                    goto exit;
                }}}}
                {paramname} = PyUnicode_READ_CHAR({argname}, 0);
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        return super().parse_arg(argname, displayname)

class unsigned_int_converter(CConverter):
    type = 'unsigned int'
    default_type = int
    c_ignored_default = "0"

    def converter_init(self, *, bitwise=False):
        if bitwise:
            self.format_unit = 'I'
        else:
            self.converter = '_PyLong_UnsignedInt_Converter'

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'I':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {paramname} = (unsigned int)PyLong_AsUnsignedLongMask({argname});
                if ({paramname} == (unsigned int)-1 && PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

class long_converter(CConverter):
    type = 'long'
    default_type = int
    format_unit = 'l'
    c_ignored_default = "0"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'l':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {paramname} = PyLong_AsLong({argname});
                if ({paramname} == -1 && PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

class unsigned_long_converter(CConverter):
    type = 'unsigned long'
    default_type = int
    c_ignored_default = "0"

    def converter_init(self, *, bitwise=False):
        if bitwise:
            self.format_unit = 'k'
        else:
            self.converter = '_PyLong_UnsignedLong_Converter'

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'k':
            return """
                if (!PyLong_Check({argname})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "int", {argname});
                    goto exit;
                }}}}
                {paramname} = PyLong_AsUnsignedLongMask({argname});
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        return super().parse_arg(argname, displayname)

class long_long_converter(CConverter):
    type = 'long long'
    default_type = int
    format_unit = 'L'
    c_ignored_default = "0"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'L':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {paramname} = PyLong_AsLongLong({argname});
                if ({paramname} == -1 && PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

class unsigned_long_long_converter(CConverter):
    type = 'unsigned long long'
    default_type = int
    c_ignored_default = "0"

    def converter_init(self, *, bitwise=False):
        if bitwise:
            self.format_unit = 'K'
        else:
            self.converter = '_PyLong_UnsignedLongLong_Converter'

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'K':
            return """
                if (!PyLong_Check({argname})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "int", {argname});
                    goto exit;
                }}}}
                {paramname} = PyLong_AsUnsignedLongLongMask({argname});
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        return super().parse_arg(argname, displayname)

class Py_ssize_t_converter(CConverter):
    type = 'Py_ssize_t'
    c_ignored_default = "0"

    def converter_init(self, *, accept={int}):
        if accept == {int}:
            self.format_unit = 'n'
            self.default_type = int
        elif accept == {int, NoneType}:
            self.converter = '_Py_convert_optional_to_ssize_t'
        else:
            fail("Py_ssize_t_converter: illegal 'accept' argument " + repr(accept))

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'n':
            return """
                if (PyFloat_Check({argname})) {{{{
                    PyErr_SetString(PyExc_TypeError,
                                    "integer argument expected, got float" );
                    goto exit;
                }}}}
                {{{{
                    Py_ssize_t ival = -1;
                    PyObject *iobj = PyNumber_Index({argname});
                    if (iobj != NULL) {{{{
                        ival = PyLong_AsSsize_t(iobj);
                        Py_DECREF(iobj);
                    }}}}
                    if (ival == -1 && PyErr_Occurred()) {{{{
                        goto exit;
                    }}}}
                    {paramname} = ival;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)


class slice_index_converter(CConverter):
    type = 'Py_ssize_t'

    def converter_init(self, *, accept={int, NoneType}):
        if accept == {int}:
            self.converter = '_PyEval_SliceIndexNotNone'
        elif accept == {int, NoneType}:
            self.converter = '_PyEval_SliceIndex'
        else:
            fail("slice_index_converter: illegal 'accept' argument " + repr(accept))

class size_t_converter(CConverter):
    type = 'size_t'
    converter = '_PyLong_Size_t_Converter'
    c_ignored_default = "0"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'n':
            return """
                {paramname} = PyNumber_AsSsize_t({argname}, PyExc_OverflowError);
                if ({paramname} == -1 && PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)


class float_converter(CConverter):
    type = 'float'
    default_type = float
    format_unit = 'f'
    c_ignored_default = "0.0"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'f':
            return """
                if (PyFloat_CheckExact({argname})) {{{{
                    {paramname} = (float) (PyFloat_AS_DOUBLE({argname}));
                }}}}
                else
                {{{{
                    {paramname} = (float) PyFloat_AsDouble({argname});
                    if ({paramname} == -1.0 && PyErr_Occurred()) {{{{
                        goto exit;
                    }}}}
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)

class double_converter(CConverter):
    type = 'double'
    default_type = float
    format_unit = 'd'
    c_ignored_default = "0.0"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'd':
            return """
                if (PyFloat_CheckExact({argname})) {{{{
                    {paramname} = PyFloat_AS_DOUBLE({argname});
                }}}}
                else
                {{{{
                    {paramname} = PyFloat_AsDouble({argname});
                    if ({paramname} == -1.0 && PyErr_Occurred()) {{{{
                        goto exit;
                    }}}}
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)


class Py_complex_converter(CConverter):
    type = 'Py_complex'
    default_type = complex
    format_unit = 'D'
    c_ignored_default = "{0.0, 0.0}"

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'D':
            return """
                {paramname} = PyComplex_AsCComplex({argname});
                if (PyErr_Occurred()) {{{{
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name)
        return super().parse_arg(argname, displayname)


class object_converter(CConverter):
    type = 'PyObject *'
    format_unit = 'O'

    def converter_init(self, *, converter=None, type=None, subclass_of=None):
        if converter:
            if subclass_of:
                fail("object: Cannot pass in both 'converter' and 'subclass_of'")
            self.format_unit = 'O&'
            self.converter = converter
        elif subclass_of:
            self.format_unit = 'O!'
            self.subclass_of = subclass_of

        if type is not None:
            self.type = type


#
# We define three conventions for buffer types in the 'accept' argument:
#
#  buffer  : any object supporting the buffer interface
#  rwbuffer: any object supporting the buffer interface, but must be writeable
#  robuffer: any object supporting the buffer interface, but must not be writeable
#

class buffer: pass
class rwbuffer: pass
class robuffer: pass

def str_converter_key(types, encoding, zeroes):
    return (frozenset(types), bool(encoding), bool(zeroes))

str_converter_argument_map = {}

class str_converter(CConverter):
    type = 'const char *'
    default_type = (str, Null, NoneType)
    format_unit = 's'

    def converter_init(self, *, accept={str}, encoding=None, zeroes=False):

        key = str_converter_key(accept, encoding, zeroes)
        format_unit = str_converter_argument_map.get(key)
        if not format_unit:
            fail("str_converter: illegal combination of arguments", key)

        self.format_unit = format_unit
        self.length = bool(zeroes)
        if encoding:
            if self.default not in (Null, None, unspecified):
                fail("str_converter: Argument Clinic doesn't support default values for encoded strings")
            self.encoding = encoding
            self.type = 'char *'
            # sorry, clinic can't support preallocated buffers
            # for es# and et#
            self.c_default = "NULL"
        if NoneType in accept and self.c_default == "Py_None":
            self.c_default = "NULL"

    def cleanup(self):
        if self.encoding:
            name = self.name
            return "".join(["if (", name, ") {\n   PyMem_FREE(", name, ");\n}\n"])

    def parse_arg(self, argname, displayname):
        if self.format_unit == 's':
            return """
                if (!PyUnicode_Check({argname})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "str", {argname});
                    goto exit;
                }}}}
                Py_ssize_t {paramname}_length;
                {paramname} = PyUnicode_AsUTF8AndSize({argname}, &{paramname}_length);
                if ({paramname} == NULL) {{{{
                    goto exit;
                }}}}
                if (strlen({paramname}) != (size_t){paramname}_length) {{{{
                    PyErr_SetString(PyExc_ValueError, "embedded null character");
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        if self.format_unit == 'z':
            return """
                if ({argname} == Py_None) {{{{
                    {paramname} = NULL;
                }}}}
                else if (PyUnicode_Check({argname})) {{{{
                    Py_ssize_t {paramname}_length;
                    {paramname} = PyUnicode_AsUTF8AndSize({argname}, &{paramname}_length);
                    if ({paramname} == NULL) {{{{
                        goto exit;
                    }}}}
                    if (strlen({paramname}) != (size_t){paramname}_length) {{{{
                        PyErr_SetString(PyExc_ValueError, "embedded null character");
                        goto exit;
                    }}}}
                }}}}
                else {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "str or None", {argname});
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        return super().parse_arg(argname, displayname)

#
# This is the fourth or fifth rewrite of registering all the
# string converter format units.  Previous approaches hid
# bugs--generally mismatches between the semantics of the format
# unit and the arguments necessary to represent those semantics
# properly.  Hopefully with this approach we'll get it 100% right.
#
# The r() function (short for "register") both registers the
# mapping from arguments to format unit *and* registers the
# legacy C converter for that format unit.
#
def r(format_unit, *, accept, encoding=False, zeroes=False):
    if not encoding and format_unit != 's':
        # add the legacy c converters here too.
        #
        # note: add_legacy_c_converter can't work for
        #   es, es#, et, or et#
        #   because of their extra encoding argument
        #
        # also don't add the converter for 's' because
        # the metaclass for CConverter adds it for us.
        kwargs = {}
        if accept != {str}:
            kwargs['accept'] = accept
        if zeroes:
            kwargs['zeroes'] = True
        added_f = functools.partial(str_converter, **kwargs)
        legacy_converters[format_unit] = added_f

    d = str_converter_argument_map
    key = str_converter_key(accept, encoding, zeroes)
    if key in d:
        sys.exit("Duplicate keys specified for str_converter_argument_map!")
    d[key] = format_unit

r('es',  encoding=True,              accept={str})
r('es#', encoding=True, zeroes=True, accept={str})
r('et',  encoding=True,              accept={bytes, bytearray, str})
r('et#', encoding=True, zeroes=True, accept={bytes, bytearray, str})
r('s',                               accept={str})
r('s#',                 zeroes=True, accept={robuffer, str})
r('y',                               accept={robuffer})
r('y#',                 zeroes=True, accept={robuffer})
r('z',                               accept={str, NoneType})
r('z#',                 zeroes=True, accept={robuffer, str, NoneType})
del r


class PyBytesObject_converter(CConverter):
    type = 'PyBytesObject *'
    format_unit = 'S'
    # accept = {bytes}

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'S':
            return """
                if (!PyBytes_Check({argname})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "bytes", {argname});
                    goto exit;
                }}}}
                {paramname} = ({type}){argname};
                """.format(argname=argname, paramname=self.name,
                           type=self.type, displayname=displayname)
        return super().parse_arg(argname, displayname)

class PyByteArrayObject_converter(CConverter):
    type = 'PyByteArrayObject *'
    format_unit = 'Y'
    # accept = {bytearray}

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'Y':
            return """
                if (!PyByteArray_Check({argname})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "bytearray", {argname});
                    goto exit;
                }}}}
                {paramname} = ({type}){argname};
                """.format(argname=argname, paramname=self.name,
                           type=self.type, displayname=displayname)
        return super().parse_arg(argname, displayname)

class unicode_converter(CConverter):
    type = 'PyObject *'
    default_type = (str, Null, NoneType)
    format_unit = 'U'

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'U':
            return """
                if (!PyUnicode_Check({argname})) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "str", {argname});
                    goto exit;
                }}}}
                if (PyUnicode_READY({argname}) == -1) {{{{
                    goto exit;
                }}}}
                {paramname} = {argname};
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        return super().parse_arg(argname, displayname)

@add_legacy_c_converter('u#', zeroes=True)
@add_legacy_c_converter('Z', accept={str, NoneType})
@add_legacy_c_converter('Z#', accept={str, NoneType}, zeroes=True)
class Py_UNICODE_converter(CConverter):
    type = 'const Py_UNICODE *'
    default_type = (str, Null, NoneType)
    format_unit = 'u'

    def converter_init(self, *, accept={str}, zeroes=False):
        format_unit = 'Z' if accept=={str, NoneType} else 'u'
        if zeroes:
            format_unit += '#'
            self.length = True
        self.format_unit = format_unit

@add_legacy_c_converter('s*', accept={str, buffer})
@add_legacy_c_converter('z*', accept={str, buffer, NoneType})
@add_legacy_c_converter('w*', accept={rwbuffer})
class Py_buffer_converter(CConverter):
    type = 'Py_buffer'
    format_unit = 'y*'
    impl_by_reference = True
    c_ignored_default = "{NULL, NULL}"

    def converter_init(self, *, accept={buffer}):
        if self.default not in (unspecified, None):
            fail("The only legal default value for Py_buffer is None.")

        self.c_default = self.c_ignored_default

        if accept == {str, buffer, NoneType}:
            format_unit = 'z*'
        elif accept == {str, buffer}:
            format_unit = 's*'
        elif accept == {buffer}:
            format_unit = 'y*'
        elif accept == {rwbuffer}:
            format_unit = 'w*'
        else:
            fail("Py_buffer_converter: illegal combination of arguments")

        self.format_unit = format_unit

    def cleanup(self):
        name = self.name
        return "".join(["if (", name, ".obj) {\n   PyBuffer_Release(&", name, ");\n}\n"])

    def parse_arg(self, argname, displayname):
        if self.format_unit == 'y*':
            return """
                if (PyObject_GetBuffer({argname}, &{paramname}, PyBUF_SIMPLE) != 0) {{{{
                    goto exit;
                }}}}
                if (!PyBuffer_IsContiguous(&{paramname}, 'C')) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "contiguous buffer", {argname});
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        elif self.format_unit == 's*':
            return """
                if (PyUnicode_Check({argname})) {{{{
                    Py_ssize_t len;
                    const char *ptr = PyUnicode_AsUTF8AndSize({argname}, &len);
                    if (ptr == NULL) {{{{
                        goto exit;
                    }}}}
                    PyBuffer_FillInfo(&{paramname}, {argname}, (void *)ptr, len, 1, 0);
                }}}}
                else {{{{ /* any bytes-like object */
                    if (PyObject_GetBuffer({argname}, &{paramname}, PyBUF_SIMPLE) != 0) {{{{
                        goto exit;
                    }}}}
                    if (!PyBuffer_IsContiguous(&{paramname}, 'C')) {{{{
                        _PyArg_BadArgument("{{name}}", {displayname}, "contiguous buffer", {argname});
                        goto exit;
                    }}}}
                }}}}
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        elif self.format_unit == 'w*':
            return """
                if (PyObject_GetBuffer({argname}, &{paramname}, PyBUF_WRITABLE) < 0) {{{{
                    PyErr_Clear();
                    _PyArg_BadArgument("{{name}}", {displayname}, "read-write bytes-like object", {argname});
                    goto exit;
                }}}}
                if (!PyBuffer_IsContiguous(&{paramname}, 'C')) {{{{
                    _PyArg_BadArgument("{{name}}", {displayname}, "contiguous buffer", {argname});
                    goto exit;
                }}}}
                """.format(argname=argname, paramname=self.name,
                           displayname=displayname)
        return super().parse_arg(argname, displayname)


def correct_name_for_self(f):
    if f.kind in (CALLABLE, METHOD_INIT):
        if f.cls:
            return "PyObject *", "self"
        return "PyObject *", "module"
    if f.kind == STATIC_METHOD:
        return "void *", "null"
    if f.kind in (CLASS_METHOD, METHOD_NEW):
        return "PyTypeObject *", "type"
    raise RuntimeError("Unhandled type of function f: " + repr(f.kind))

def required_type_for_self_for_parser(f):
    type, _ = correct_name_for_self(f)
    if f.kind in (METHOD_INIT, METHOD_NEW, STATIC_METHOD, CLASS_METHOD):
        return type
    return None


class self_converter(CConverter):
    """
    A special-case converter:
    this is the default converter used for "self".
    """
    type = None
    format_unit = ''

    def converter_init(self, *, type=None):
        self.specified_type = type

    def pre_render(self):
        f = self.function
        default_type, default_name = correct_name_for_self(f)
        self.signature_name = default_name
        self.type = self.specified_type or self.type or default_type

        kind = self.function.kind
        new_or_init = kind in (METHOD_NEW, METHOD_INIT)

        if (kind == STATIC_METHOD) or new_or_init:
            self.show_in_signature = False

    # tp_new (METHOD_NEW) functions are of type newfunc:
    #     typedef PyObject *(*newfunc)(struct _typeobject *, PyObject *, PyObject *);
    # PyTypeObject is a typedef for struct _typeobject.
    #
    # tp_init (METHOD_INIT) functions are of type initproc:
    #     typedef int (*initproc)(PyObject *, PyObject *, PyObject *);
    #
    # All other functions generated by Argument Clinic are stored in
    # PyMethodDef structures, in the ml_meth slot, which is of type PyCFunction:
    #     typedef PyObject *(*PyCFunction)(PyObject *, PyObject *);
    # However!  We habitually cast these functions to PyCFunction,
    # since functions that accept keyword arguments don't fit this signature
    # but are stored there anyway.  So strict type equality isn't important
    # for these functions.
    #
    # So:
    #
    # * The name of the first parameter to the impl and the parsing function will always
    #   be self.name.
    #
    # * The type of the first parameter to the impl will always be of self.type.
    #
    # * If the function is neither tp_new (METHOD_NEW) nor tp_init (METHOD_INIT):
    #   * The type of the first parameter to the parsing function is also self.type.
    #     This means that if you step into the parsing function, your "self" parameter
    #     is of the correct type, which may make debugging more pleasant.
    #
    # * Else if the function is tp_new (METHOD_NEW):
    #   * The type of the first parameter to the parsing function is "PyTypeObject *",
    #     so the type signature of the function call is an exact match.
    #   * If self.type != "PyTypeObject *", we cast the first parameter to self.type
    #     in the impl call.
    #
    # * Else if the function is tp_init (METHOD_INIT):
    #   * The type of the first parameter to the parsing function is "PyObject *",
    #     so the type signature of the function call is an exact match.
    #   * If self.type != "PyObject *", we cast the first parameter to self.type
    #     in the impl call.

    @property
    def parser_type(self):
        return required_type_for_self_for_parser(self.function) or self.type

    def render(self, parameter, data):
        """
        parameter is a clinic.Parameter instance.
        data is a CRenderData instance.
        """
        if self.function.kind == STATIC_METHOD:
            return

        self._render_self(parameter, data)

        if self.type != self.parser_type:
            # insert cast to impl_argument[0], aka self.
            # we know we're in the first slot in all the CRenderData lists,
            # because we render parameters in order, and self is always first.
            assert len(data.impl_arguments) == 1
            assert data.impl_arguments[0] == self.name
            data.impl_arguments[0] = '(' + self.type + ")" + data.impl_arguments[0]

    def set_template_dict(self, template_dict):
        template_dict['self_name'] = self.name
        template_dict['self_type'] = self.parser_type
        kind = self.function.kind
        cls = self.function.cls

        if ((kind in (METHOD_NEW, METHOD_INIT)) and cls and cls.typedef):
            type_object = self.function.cls.type_object
            if kind == METHOD_NEW:
                type_check = '({} == {})'.format(self.name, type_object)
            else:
                type_check = 'Py_IS_TYPE({}, {})'.format(self.name, type_object)

            line = '{} &&\n        '.format(type_check)
            template_dict['self_type_check'] = line



def add_c_return_converter(f, name=None):
    if not name:
        name = f.__name__
        if not name.endswith('_return_converter'):
            return f
        name = name[:-len('_return_converter')]
    return_converters[name] = f
    return f


class CReturnConverterAutoRegister(type):
    def __init__(cls, name, bases, classdict):
        add_c_return_converter(cls)

class CReturnConverter(metaclass=CReturnConverterAutoRegister):

    # The C type to use for this variable.
    # 'type' should be a Python string specifying the type, e.g. "int".
    # If this is a pointer type, the type string should end with ' *'.
    type = 'PyObject *'

    # The Python default value for this parameter, as a Python value.
    # Or the magic value "unspecified" if there is no default.
    default = None

    def __init__(self, *, py_default=None, **kwargs):
        self.py_default = py_default
        try:
            self.return_converter_init(**kwargs)
        except TypeError as e:
            s = ', '.join(name + '=' + repr(value) for name, value in kwargs.items())
            sys.exit(self.__class__.__name__ + '(' + s + ')\n' + str(e))

    def return_converter_init(self):
        pass

    def declare(self, data, name="_return_value"):
        line = []
        add = line.append
        add(self.type)
        if not self.type.endswith('*'):
            add(' ')
        add(name + ';')
        data.declarations.append(''.join(line))
        data.return_value = name

    def err_occurred_if(self, expr, data):
        data.return_conversion.append('if (({}) && PyErr_Occurred()) {{\n    goto exit;\n}}\n'.format(expr))

    def err_occurred_if_null_pointer(self, variable, data):
        data.return_conversion.append('if ({} == NULL) {{\n    goto exit;\n}}\n'.format(variable))

    def render(self, function, data):
        """
        function is a clinic.Function instance.
        data is a CRenderData instance.
        """
        pass

add_c_return_converter(CReturnConverter, 'object')

class NoneType_return_converter(CReturnConverter):
    def render(self, function, data):
        self.declare(data)
        data.return_conversion.append('''
if (_return_value != Py_None) {
    goto exit;
}
return_value = Py_None;
Py_INCREF(Py_None);
'''.strip())

class bool_return_converter(CReturnConverter):
    type = 'int'

    def render(self, function, data):
        self.declare(data)
        self.err_occurred_if("_return_value == -1", data)
        data.return_conversion.append('return_value = PyBool_FromLong((long)_return_value);\n')

class long_return_converter(CReturnConverter):
    type = 'long'
    conversion_fn = 'PyLong_FromLong'
    cast = ''
    unsigned_cast = ''

    def render(self, function, data):
        self.declare(data)
        self.err_occurred_if("_return_value == {}-1".format(self.unsigned_cast), data)
        data.return_conversion.append(
            ''.join(('return_value = ', self.conversion_fn, '(', self.cast, '_return_value);\n')))

class int_return_converter(long_return_converter):
    type = 'int'
    cast = '(long)'

class init_return_converter(long_return_converter):
    """
    Special return converter for __init__ functions.
    """
    type = 'int'
    cast = '(long)'

    def render(self, function, data):
        pass

class unsigned_long_return_converter(long_return_converter):
    type = 'unsigned long'
    conversion_fn = 'PyLong_FromUnsignedLong'
    unsigned_cast = '(unsigned long)'

class unsigned_int_return_converter(unsigned_long_return_converter):
    type = 'unsigned int'
    cast = '(unsigned long)'
    unsigned_cast = '(unsigned int)'

class Py_ssize_t_return_converter(long_return_converter):
    type = 'Py_ssize_t'
    conversion_fn = 'PyLong_FromSsize_t'

class size_t_return_converter(long_return_converter):
    type = 'size_t'
    conversion_fn = 'PyLong_FromSize_t'
    unsigned_cast = '(size_t)'


class double_return_converter(CReturnConverter):
    type = 'double'
    cast = ''

    def render(self, function, data):
        self.declare(data)
        self.err_occurred_if("_return_value == -1.0", data)
        data.return_conversion.append(
            'return_value = PyFloat_FromDouble(' + self.cast + '_return_value);\n')

class float_return_converter(double_return_converter):
    type = 'float'
    cast = '(double)'


def eval_ast_expr(node, globals, *, filename='-'):
    """
    Takes an ast.Expr node.  Compiles and evaluates it.
    Returns the result of the expression.

    globals represents the globals dict the expression
    should see.  (There's no equivalent for "locals" here.)
    """

    if isinstance(node, ast.Expr):
        node = node.value

    node = ast.Expression(node)
    co = compile(node, filename, 'eval')
    fn = types.FunctionType(co, globals)
    return fn()


class IndentStack:
    def __init__(self):
        self.indents = []
        self.margin = None

    def _ensure(self):
        if not self.indents:
            fail('IndentStack expected indents, but none are defined.')

    def measure(self, line):
        """
        Returns the length of the line's margin.
        """
        if '\t' in line:
            fail('Tab characters are illegal in the Argument Clinic DSL.')
        stripped = line.lstrip()
        if not len(stripped):
            # we can't tell anything from an empty line
            # so just pretend it's indented like our current indent
            self._ensure()
            return self.indents[-1]
        return len(line) - len(stripped)

    def infer(self, line):
        """
        Infer what is now the current margin based on this line.
        Returns:
            1 if we have indented (or this is the first margin)
            0 if the margin has not changed
           -N if we have dedented N times
        """
        indent = self.measure(line)
        margin = ' ' * indent
        if not self.indents:
            self.indents.append(indent)
            self.margin = margin
            return 1
        current = self.indents[-1]
        if indent == current:
            return 0
        if indent > current:
            self.indents.append(indent)
            self.margin = margin
            return 1
        # indent < current
        if indent not in self.indents:
            fail("Illegal outdent.")
        outdent_count = 0
        while indent != current:
            self.indents.pop()
            current = self.indents[-1]
            outdent_count -= 1
        self.margin = margin
        return outdent_count

    @property
    def depth(self):
        """
        Returns how many margins are currently defined.
        """
        return len(self.indents)

    def indent(self, line):
        """
        Indents a line by the currently defined margin.
        """
        return self.margin + line

    def dedent(self, line):
        """
        Dedents a line by the currently defined margin.
        (The inverse of 'indent'.)
        """
        margin = self.margin
        indent = self.indents[-1]
        if not line.startswith(margin):
            fail('Cannot dedent, line does not start with the previous margin:')
        return line[indent:]


class DSLParser:
    def __init__(self, clinic):
        self.clinic = clinic

        self.directives = {}
        for name in dir(self):
            # functions that start with directive_ are added to directives
            _, s, key = name.partition("directive_")
            if s:
                self.directives[key] = getattr(self, name)

            # functions that start with at_ are too, with an @ in front
            _, s, key = name.partition("at_")
            if s:
                self.directives['@' + key] = getattr(self, name)

        self.reset()

    def reset(self):
        self.function = None
        self.state = self.state_dsl_start
        self.parameter_indent = None
        self.keyword_only = False
        self.positional_only = False
        self.group = 0
        self.parameter_state = self.ps_start
        self.seen_positional_with_default = False
        self.indent = IndentStack()
        self.kind = CALLABLE
        self.coexist = False
        self.parameter_continuation = ''
        self.preserve_output = False

    def directive_version(self, required):
        global version
        if version_comparitor(version, required) < 0:
            fail("Insufficient Clinic version!\n  Version: " + version + "\n  Required: " + required)

    def directive_module(self, name):
        fields = name.split('.')
        new = fields.pop()
        module, cls = self.clinic._module_and_class(fields)
        if cls:
            fail("Can't nest a module inside a class!")

        if name in module.classes:
            fail("Already defined module " + repr(name) + "!")

        m = Module(name, module)
        module.modules[name] = m
        self.block.signatures.append(m)

    def directive_class(self, name, typedef, type_object):
        fields = name.split('.')
        in_classes = False
        parent = self
        name = fields.pop()
        so_far = []
        module, cls = self.clinic._module_and_class(fields)

        parent = cls or module
        if name in parent.classes:
            fail("Already defined class " + repr(name) + "!")

        c = Class(name, module, cls, typedef, type_object)
        parent.classes[name] = c
        self.block.signatures.append(c)

    def directive_set(self, name, value):
        if name not in ("line_prefix", "line_suffix"):
            fail("unknown variable", repr(name))

        value = value.format_map({
            'block comment start': '/*',
            'block comment end': '*/',
            })

        self.clinic.__dict__[name] = value

    def directive_destination(self, name, command, *args):
        if command == 'new':
            self.clinic.add_destination(name, *args)
            return

        if command == 'clear':
            self.clinic.get_destination(name).clear()
        fail("unknown destination command", repr(command))


    def directive_output(self, command_or_name, destination=''):
        fd = self.clinic.destination_buffers

        if command_or_name == "preset":
            preset = self.clinic.presets.get(destination)
            if not preset:
                fail("Unknown preset " + repr(destination) + "!")
            fd.update(preset)
            return

        if command_or_name == "push":
            self.clinic.destination_buffers_stack.append(fd.copy())
            return

        if command_or_name == "pop":
            if not self.clinic.destination_buffers_stack:
                fail("Can't 'output pop', stack is empty!")
            previous_fd = self.clinic.destination_buffers_stack.pop()
            fd.update(previous_fd)
            return

        # secret command for debugging!
        if command_or_name == "print":
            self.block.output.append(pprint.pformat(fd))
            self.block.output.append('\n')
            return

        d = self.clinic.get_destination(destination)

        if command_or_name == "everything":
            for name in list(fd):
                fd[name] = d
            return

        if command_or_name not in fd:
            fail("Invalid command / destination name " + repr(command_or_name) + ", must be one of:\n  preset push pop print everything " + " ".join(fd))
        fd[command_or_name] = d

    def directive_dump(self, name):
        self.block.output.append(self.clinic.get_destination(name).dump())

    def directive_print(self, *args):
        self.block.output.append(' '.join(args))
        self.block.output.append('\n')

    def directive_preserve(self):
        if self.preserve_output:
            fail("Can't have preserve twice in one block!")
        self.preserve_output = True

    def at_classmethod(self):
        if self.kind is not CALLABLE:
            fail("Can't set @classmethod, function is not a normal callable")
        self.kind = CLASS_METHOD

    def at_staticmethod(self):
        if self.kind is not CALLABLE:
            fail("Can't set @staticmethod, function is not a normal callable")
        self.kind = STATIC_METHOD

    def at_coexist(self):
        if self.coexist:
            fail("Called @coexist twice!")
        self.coexist = True

    def parse(self, block):
        self.reset()
        self.block = block
        self.saved_output = self.block.output
        block.output = []
        block_start = self.clinic.block_parser.line_number
        lines = block.input.split('\n')
        for line_number, line in enumerate(lines, self.clinic.block_parser.block_start_line_number):
            if '\t' in line:
                fail('Tab characters are illegal in the Clinic DSL.\n\t' + repr(line), line_number=block_start)
            self.state(line)

        self.next(self.state_terminal)
        self.state(None)

        block.output.extend(self.clinic.language.render(clinic, block.signatures))

        if self.preserve_output:
            if block.output:
                fail("'preserve' only works for blocks that don't produce any output!")
            block.output = self.saved_output

    @staticmethod
    def ignore_line(line):
        # ignore comment-only lines
        if line.lstrip().startswith('#'):
            return True

        # Ignore empty lines too
        # (but not in docstring sections!)
        if not line.strip():
            return True

        return False

    @staticmethod
    def calculate_indent(line):
        return len(line) - len(line.strip())

    def next(self, state, line=None):
        # real_print(self.state.__name__, "->", state.__name__, ", line=", line)
        self.state = state
        if line is not None:
            self.state(line)

    def state_dsl_start(self, line):
        # self.block = self.ClinicOutputBlock(self)
        if self.ignore_line(line):
            return

        # is it a directive?
        fields = shlex.split(line)
        directive_name = fields[0]
        directive = self.directives.get(directive_name, None)
        if directive:
            try:
                directive(*fields[1:])
            except TypeError as e:
                fail(str(e))
            return

        self.next(self.state_modulename_name, line)

    def state_modulename_name(self, line):
        # looking for declaration, which establishes the leftmost column
        # line should be
        #     modulename.fnname [as c_basename] [-> return annotation]
        # square brackets denote optional syntax.
        #
        # alternatively:
        #     modulename.fnname [as c_basename] = modulename.existing_fn_name
        # clones the parameters and return converter from that
        # function.  you can't modify them.  you must enter a
        # new docstring.
        #
        # (but we might find a directive first!)
        #
        # this line is permitted to start with whitespace.
        # we'll call this number of spaces F (for "function").

        if not line.strip():
            return

        self.indent.infer(line)

        # are we cloning?
        before, equals, existing = line.rpartition('=')
        if equals:
            full_name, _, c_basename = before.partition(' as ')
            full_name = full_name.strip()
            c_basename = c_basename.strip()
            existing = existing.strip()
            if (is_legal_py_identifier(full_name) and
                (not c_basename or is_legal_c_identifier(c_basename)) and
                is_legal_py_identifier(existing)):
                # we're cloning!
                fields = [x.strip() for x in existing.split('.')]
                function_name = fields.pop()
                module, cls = self.clinic._module_and_class(fields)

                for existing_function in (cls or module).functions:
                    if existing_function.name == function_name:
                        break
                else:
                    existing_function = None
                if not existing_function:
                    print("class", cls, "module", module, "existing", existing)
                    print("cls. functions", cls.functions)
                    fail("Couldn't find existing function " + repr(existing) + "!")

                fields = [x.strip() for x in full_name.split('.')]
                function_name = fields.pop()
                module, cls = self.clinic._module_and_class(fields)

                if not (existing_function.kind == self.kind and existing_function.coexist == self.coexist):
                    fail("'kind' of function and cloned function don't match!  (@classmethod/@staticmethod/@coexist)")
                self.function = existing_function.copy(name=function_name, full_name=full_name, module=module, cls=cls, c_basename=c_basename, docstring='')

                self.block.signatures.append(self.function)
                (cls or module).functions.append(self.function)
                self.next(self.state_function_docstring)
                return

        line, _, returns = line.partition('->')

        full_name, _, c_basename = line.partition(' as ')
        full_name = full_name.strip()
        c_basename = c_basename.strip() or None

        if not is_legal_py_identifier(full_name):
            fail("Illegal function name: {}".format(full_name))
        if c_basename and not is_legal_c_identifier(c_basename):
            fail("Illegal C basename: {}".format(c_basename))

        return_converter = None
        if returns:
            ast_input = "def x() -> {}: pass".format(returns)
            module = None
            try:
                module = ast.parse(ast_input)
            except SyntaxError:
                pass
            if not module:
                fail("Badly-formed annotation for " + full_name + ": " + returns)
            try:
                name, legacy, kwargs = self.parse_converter(module.body[0].returns)
                if legacy:
                    fail("Legacy converter {!r} not allowed as a return converter"
                         .format(name))
                if name not in return_converters:
                    fail("No available return converter called " + repr(name))
                return_converter = return_converters[name](**kwargs)
            except ValueError:
                fail("Badly-formed annotation for " + full_name + ": " + returns)

        fields = [x.strip() for x in full_name.split('.')]
        function_name = fields.pop()
        module, cls = self.clinic._module_and_class(fields)

        fields = full_name.split('.')
        if fields[-1] == '__new__':
            if (self.kind != CLASS_METHOD) or (not cls):
                fail("__new__ must be a class method!")
            self.kind = METHOD_NEW
        elif fields[-1] == '__init__':
            if (self.kind != CALLABLE) or (not cls):
                fail("__init__ must be a normal method, not a class or static method!")
            self.kind = METHOD_INIT
            if not return_converter:
                return_converter = init_return_converter()
        elif fields[-1] in unsupported_special_methods:
            fail(fields[-1] + " is a special method and cannot be converted to Argument Clinic!  (Yet.)")

        if not return_converter:
            return_converter = CReturnConverter()

        if not module:
            fail("Undefined module used in declaration of " + repr(full_name.strip()) + ".")
        self.function = Function(name=function_name, full_name=full_name, module=module, cls=cls, c_basename=c_basename,
                                 return_converter=return_converter, kind=self.kind, coexist=self.coexist)
        self.block.signatures.append(self.function)

        # insert a self converter automatically
        type, name = correct_name_for_self(self.function)
        kwargs = {}
        if cls and type == "PyObject *":
            kwargs['type'] = cls.typedef
        sc = self.function.self_converter = self_converter(name, name, self.function, **kwargs)
        p_self = Parameter(sc.name, inspect.Parameter.POSITIONAL_ONLY, function=self.function, converter=sc)
        self.function.parameters[sc.name] = p_self

        (cls or module).functions.append(self.function)
        self.next(self.state_parameters_start)

    # Now entering the parameters section.  The rules, formally stated:
    #
    #   * All lines must be indented with spaces only.
    #   * The first line must be a parameter declaration.
    #   * The first line must be indented.
    #       * This first line establishes the indent for parameters.
    #       * We'll call this number of spaces P (for "parameter").
    #   * Thenceforth:
    #       * Lines indented with P spaces specify a parameter.
    #       * Lines indented with > P spaces are docstrings for the previous
    #         parameter.
    #           * We'll call this number of spaces D (for "docstring").
    #           * All subsequent lines indented with >= D spaces are stored as
    #             part of the per-parameter docstring.
    #           * All lines will have the first D spaces of the indent stripped
    #             before they are stored.
    #           * It's illegal to have a line starting with a number of spaces X
    #             such that P < X < D.
    #       * A line with < P spaces is the first line of the function
    #         docstring, which ends processing for parameters and per-parameter
    #         docstrings.
    #           * The first line of the function docstring must be at the same
    #             indent as the function declaration.
    #       * It's illegal to have any line in the parameters section starting
    #         with X spaces such that F < X < P.  (As before, F is the indent
    #         of the function declaration.)
    #
    # Also, currently Argument Clinic places the following restrictions on groups:
    #   * Each group must contain at least one parameter.
    #   * Each group may contain at most one group, which must be the furthest
    #     thing in the group from the required parameters.  (The nested group
    #     must be the first in the group when it's before the required
    #     parameters, and the last thing in the group when after the required
    #     parameters.)
    #   * There may be at most one (top-level) group to the left or right of
    #     the required parameters.
    #   * You must specify a slash, and it must be after all parameters.
    #     (In other words: either all parameters are positional-only,
    #      or none are.)
    #
    #  Said another way:
    #   * Each group must contain at least one parameter.
    #   * All left square brackets before the required parameters must be
    #     consecutive.  (You can't have a left square bracket followed
    #     by a parameter, then another left square bracket.  You can't
    #     have a left square bracket, a parameter, a right square bracket,
    #     and then a left square bracket.)
    #   * All right square brackets after the required parameters must be
    #     consecutive.
    #
    # These rules are enforced with a single state variable:
    # "parameter_state".  (Previously the code was a miasma of ifs and
    # separate boolean state variables.)  The states are:
    #
    #  [ [ a, b, ] c, ] d, e, f=3, [ g, h, [ i ] ]   <- line
    # 01   2          3       4    5           6     <- state transitions
    #
    # 0: ps_start.  before we've seen anything.  legal transitions are to 1 or 3.
    # 1: ps_left_square_before.  left square brackets before required parameters.
    # 2: ps_group_before.  in a group, before required parameters.
    # 3: ps_required.  required parameters, positional-or-keyword or positional-only
    #     (we don't know yet).  (renumber left groups!)
    # 4: ps_optional.  positional-or-keyword or positional-only parameters that
    #    now must have default values.
    # 5: ps_group_after.  in a group, after required parameters.
    # 6: ps_right_square_after.  right square brackets after required parameters.
    ps_start, ps_left_square_before, ps_group_before, ps_required, \
    ps_optional, ps_group_after, ps_right_square_after = range(7)

    def state_parameters_start(self, line):
        if self.ignore_line(line):
            return

        # if this line is not indented, we have no parameters
        if not self.indent.infer(line):
            return self.next(self.state_function_docstring, line)

        self.parameter_continuation = ''
        return self.next(self.state_parameter, line)


    def to_required(self):
        """
        Transition to the "required" parameter state.
        """
        if self.parameter_state != self.ps_required:
            self.parameter_state = self.ps_required
            for p in self.function.parameters.values():
                p.group = -p.group

    def state_parameter(self, line):
        if self.parameter_continuation:
            line = self.parameter_continuation + ' ' + line.lstrip()
            self.parameter_continuation = ''

        if self.ignore_line(line):
            return

        assert self.indent.depth == 2
        indent = self.indent.infer(line)
        if indent == -1:
            # we outdented, must be to definition column
            return self.next(self.state_function_docstring, line)

        if indent == 1:
            # we indented, must be to new parameter docstring column
            return self.next(self.state_parameter_docstring_start, line)

        line = line.rstrip()
        if line.endswith('\\'):
            self.parameter_continuation = line[:-1]
            return

        line = line.lstrip()

        if line in ('*', '/', '[', ']'):
            self.parse_special_symbol(line)
            return

        if self.parameter_state in (self.ps_start, self.ps_required):
            self.to_required()
        elif self.parameter_state == self.ps_left_square_before:
            self.parameter_state = self.ps_group_before
        elif self.parameter_state == self.ps_group_before:
            if not self.group:
                self.to_required()
        elif self.parameter_state in (self.ps_group_after, self.ps_optional):
            pass
        else:
            fail("Function " + self.function.name + " has an unsupported group configuration. (Unexpected state " + str(self.parameter_state) + ".a)")

        # handle "as" for  parameters too
        c_name = None
        name, have_as_token, trailing = line.partition(' as ')
        if have_as_token:
            name = name.strip()
            if ' ' not in name:
                fields = trailing.strip().split(' ')
                if not fields:
                    fail("Invalid 'as' clause!")
                c_name = fields[0]
                if c_name.endswith(':'):
                    name += ':'
                    c_name = c_name[:-1]
                fields[0] = name
                line = ' '.join(fields)

        base, equals, default = line.rpartition('=')
        if not equals:
            base = default
            default = None

        module = None
        try:
            ast_input = "def x({}): pass".format(base)
            module = ast.parse(ast_input)
        except SyntaxError:
            try:
                # the last = was probably inside a function call, like
                #   c: int(accept={str})
                # so assume there was no actual default value.
                default = None
                ast_input = "def x({}): pass".format(line)
                module = ast.parse(ast_input)
            except SyntaxError:
                pass
        if not module:
            fail("Function " + self.function.name + " has an invalid parameter declaration:\n\t" + line)

        function_args = module.body[0].args

        if len(function_args.args) > 1:
            fail("Function " + self.function.name + " has an invalid parameter declaration (comma?):\n\t" + line)
        if function_args.defaults or function_args.kw_defaults:
            fail("Function " + self.function.name + " has an invalid parameter declaration (default value?):\n\t" + line)
        if function_args.vararg or function_args.kwarg:
            fail("Function " + self.function.name + " has an invalid parameter declaration (*args? **kwargs?):\n\t" + line)

        parameter = function_args.args[0]

        parameter_name = parameter.arg
        name, legacy, kwargs = self.parse_converter(parameter.annotation)

        if not default:
            if self.parameter_state == self.ps_optional:
                fail("Can't have a parameter without a default (" + repr(parameter_name) + ")\nafter a parameter with a default!")
            value = unspecified
            if 'py_default' in kwargs:
                fail("You can't specify py_default without specifying a default value!")
        else:
            if self.parameter_state == self.ps_required:
                self.parameter_state = self.ps_optional
            default = default.strip()
            bad = False
            ast_input = "x = {}".format(default)
            bad = False
            try:
                module = ast.parse(ast_input)

                if 'c_default' not in kwargs:
                    # we can only represent very simple data values in C.
                    # detect whether default is okay, via a blacklist
                    # of disallowed ast nodes.
                    class DetectBadNodes(ast.NodeVisitor):
                        bad = False
                        def bad_node(self, node):
                            self.bad = True

                        # inline function call
                        visit_Call = bad_node
                        # inline if statement ("x = 3 if y else z")
                        visit_IfExp = bad_node

                        # comprehensions and generator expressions
                        visit_ListComp = visit_SetComp = bad_node
                        visit_DictComp = visit_GeneratorExp = bad_node

                        # literals for advanced types
                        visit_Dict = visit_Set = bad_node
                        visit_List = visit_Tuple = bad_node

                        # "starred": "a = [1, 2, 3]; *a"
                        visit_Starred = bad_node

                    blacklist = DetectBadNodes()
                    blacklist.visit(module)
                    bad = blacklist.bad
                else:
                    # if they specify a c_default, we can be more lenient about the default value.
                    # but at least make an attempt at ensuring it's a valid expression.
                    try:
                        value = eval(default)
                        if value == unspecified:
                            fail("'unspecified' is not a legal default value!")
                    except NameError:
                        pass # probably a named constant
                    except Exception as e:
                        fail("Malformed expression given as default value\n"
                             "{!r} caused {!r}".format(default, e))
                if bad:
                    fail("Unsupported expression as default value: " + repr(default))

                expr = module.body[0].value
                # mild hack: explicitly support NULL as a default value
                if isinstance(expr, ast.Name) and expr.id == 'NULL':
                    value = NULL
                    py_default = '<unrepresentable>'
                    c_default = "NULL"
                elif (isinstance(expr, ast.BinOp) or
                    (isinstance(expr, ast.UnaryOp) and
                     not (isinstance(expr.operand, ast.Num) or
                          (hasattr(ast, 'Constant') and
                           isinstance(expr.operand, ast.Constant) and
                           type(expr.operand.value) in (int, float, complex)))
                    )):
                    c_default = kwargs.get("c_default")
                    if not (isinstance(c_default, str) and c_default):
                        fail("When you specify an expression (" + repr(default) + ") as your default value,\nyou MUST specify a valid c_default." + ast.dump(expr))
                    py_default = default
                    value = unknown
                elif isinstance(expr, ast.Attribute):
                    a = []
                    n = expr
                    while isinstance(n, ast.Attribute):
                        a.append(n.attr)
                        n = n.value
                    if not isinstance(n, ast.Name):
                        fail("Unsupported default value " + repr(default) + " (looked like a Python constant)")
                    a.append(n.id)
                    py_default = ".".join(reversed(a))

                    c_default = kwargs.get("c_default")
                    if not (isinstance(c_default, str) and c_default):
                        fail("When you specify a named constant (" + repr(py_default) + ") as your default value,\nyou MUST specify a valid c_default.")

                    try:
                        value = eval(py_default)
                    except NameError:
                        value = unknown
                else:
                    value = ast.literal_eval(expr)
                    py_default = repr(value)
                    if isinstance(value, (bool, None.__class__)):
                        c_default = "Py_" + py_default
                    elif isinstance(value, str):
                        c_default = c_repr(value)
                    else:
                        c_default = py_default

            except SyntaxError as e:
                fail("Syntax error: " + repr(e.text))
            except (ValueError, AttributeError):
                value = unknown
                c_default = kwargs.get("c_default")
                py_default = default
                if not (isinstance(c_default, str) and c_default):
                    fail("When you specify a named constant (" + repr(py_default) + ") as your default value,\nyou MUST specify a valid c_default.")

            kwargs.setdefault('c_default', c_default)
            kwargs.setdefault('py_default', py_default)

        dict = legacy_converters if legacy else converters
        legacy_str = "legacy " if legacy else ""
        if name not in dict:
            fail('{} is not a valid {}converter'.format(name, legacy_str))
        # if you use a c_name for the parameter, we just give that name to the converter
        # but the parameter object gets the python name
        converter = dict[name](c_name or parameter_name, parameter_name, self.function, value, **kwargs)

        kind = inspect.Parameter.KEYWORD_ONLY if self.keyword_only else inspect.Parameter.POSITIONAL_OR_KEYWORD

        if isinstance(converter, self_converter):
            if len(self.function.parameters) == 1:
                if (self.parameter_state != self.ps_required):
                    fail("A 'self' parameter cannot be marked optional.")
                if value is not unspecified:
                    fail("A 'self' parameter cannot have a default value.")
                if self.group:
                    fail("A 'self' parameter cannot be in an optional group.")
                kind = inspect.Parameter.POSITIONAL_ONLY
                self.parameter_state = self.ps_start
                self.function.parameters.clear()
            else:
                fail("A 'self' parameter, if specified, must be the very first thing in the parameter block.")

        if isinstance(converter, defining_class_converter):
            _lp = len(self.function.parameters)
            if _lp == 1:
                if (self.parameter_state != self.ps_required):
                    fail("A 'defining_class' parameter cannot be marked optional.")
                if value is not unspecified:
                    fail("A 'defining_class' parameter cannot have a default value.")
                if self.group:
                    fail("A 'defining_class' parameter cannot be in an optional group.")
            else:
                fail("A 'defining_class' parameter, if specified, must either be the first thing in the parameter block, or come just after 'self'.")


        p = Parameter(parameter_name, kind, function=self.function, converter=converter, default=value, group=self.group)

        if parameter_name in self.function.parameters:
            fail("You can't have two parameters named " + repr(parameter_name) + "!")
        self.function.parameters[parameter_name] = p

    def parse_converter(self, annotation):
        if (hasattr(ast, 'Constant') and
            isinstance(annotation, ast.Constant) and
            type(annotation.value) is str):
            return annotation.value, True, {}

        if isinstance(annotation, ast.Str):
            return annotation.s, True, {}

        if isinstance(annotation, ast.Name):
            return annotation.id, False, {}

        if not isinstance(annotation, ast.Call):
            fail("Annotations must be either a name, a function call, or a string.")

        name = annotation.func.id
        symbols = globals()

        kwargs = {node.arg: eval_ast_expr(node.value, symbols) for node in annotation.keywords}
        return name, False, kwargs

    def parse_special_symbol(self, symbol):
        if symbol == '*':
            if self.keyword_only:
                fail("Function " + self.function.name + " uses '*' more than once.")
            self.keyword_only = True
        elif symbol == '[':
            if self.parameter_state in (self.ps_start, self.ps_left_square_before):
                self.parameter_state = self.ps_left_square_before
            elif self.parameter_state in (self.ps_required, self.ps_group_after):
                self.parameter_state = self.ps_group_after
            else:
                fail("Function " + self.function.name + " has an unsupported group configuration. (Unexpected state " + str(self.parameter_state) + ".b)")
            self.group += 1
            self.function.docstring_only = True
        elif symbol == ']':
            if not self.group:
                fail("Function " + self.function.name + " has a ] without a matching [.")
            if not any(p.group == self.group for p in self.function.parameters.values()):
                fail("Function " + self.function.name + " has an empty group.\nAll groups must contain at least one parameter.")
            self.group -= 1
            if self.parameter_state in (self.ps_left_square_before, self.ps_group_before):
                self.parameter_state = self.ps_group_before
            elif self.parameter_state in (self.ps_group_after, self.ps_right_square_after):
                self.parameter_state = self.ps_right_square_after
            else:
                fail("Function " + self.function.name + " has an unsupported group configuration. (Unexpected state " + str(self.parameter_state) + ".c)")
        elif symbol == '/':
            if self.positional_only:
                fail("Function " + self.function.name + " uses '/' more than once.")
            self.positional_only = True
            # ps_required and ps_optional are allowed here, that allows positional-only without option groups
            # to work (and have default values!)
            if (self.parameter_state not in (self.ps_required, self.ps_optional, self.ps_right_square_after, self.ps_group_before)) or self.group:
                fail("Function " + self.function.name + " has an unsupported group configuration. (Unexpected state " + str(self.parameter_state) + ".d)")
            if self.keyword_only:
                fail("Function " + self.function.name + " mixes keyword-only and positional-only parameters, which is unsupported.")
            # fixup preceding parameters
            for p in self.function.parameters.values():
                if (p.kind != inspect.Parameter.POSITIONAL_OR_KEYWORD and not isinstance(p.converter, self_converter)):
                    fail("Function " + self.function.name + " mixes keyword-only and positional-only parameters, which is unsupported.")
                p.kind = inspect.Parameter.POSITIONAL_ONLY

    def state_parameter_docstring_start(self, line):
        self.parameter_docstring_indent = len(self.indent.margin)
        assert self.indent.depth == 3
        return self.next(self.state_parameter_docstring, line)

    # every line of the docstring must start with at least F spaces,
    # where F > P.
    # these F spaces will be stripped.
    def state_parameter_docstring(self, line):
        stripped = line.strip()
        if stripped.startswith('#'):
            return

        indent = self.indent.measure(line)
        if indent < self.parameter_docstring_indent:
            self.indent.infer(line)
            assert self.indent.depth < 3
            if self.indent.depth == 2:
                # back to a parameter
                return self.next(self.state_parameter, line)
            assert self.indent.depth == 1
            return self.next(self.state_function_docstring, line)

        assert self.function.parameters
        last_parameter = next(reversed(list(self.function.parameters.values())))

        new_docstring = last_parameter.docstring

        if new_docstring:
            new_docstring += '\n'
        if stripped:
            new_docstring += self.indent.dedent(line)

        last_parameter.docstring = new_docstring

    # the final stanza of the DSL is the docstring.
    def state_function_docstring(self, line):
        if self.group:
            fail("Function " + self.function.name + " has a ] without a matching [.")

        stripped = line.strip()
        if stripped.startswith('#'):
            return

        new_docstring = self.function.docstring
        if new_docstring:
            new_docstring += "\n"
        if stripped:
            line = self.indent.dedent(line).rstrip()
        else:
            line = ''
        new_docstring += line
        self.function.docstring = new_docstring

    def format_docstring(self):
        f = self.function

        new_or_init = f.kind in (METHOD_NEW, METHOD_INIT)
        if new_or_init and not f.docstring:
            # don't render a docstring at all, no signature, nothing.
            return f.docstring

        text, add, output = _text_accumulator()
        parameters = f.render_parameters

        ##
        ## docstring first line
        ##

        if new_or_init:
            # classes get *just* the name of the class
            # not __new__, not __init__, and not module.classname
            assert f.cls
            add(f.cls.name)
        else:
            add(f.name)
        add('(')

        # populate "right_bracket_count" field for every parameter
        assert parameters, "We should always have a self parameter. " + repr(f)
        assert isinstance(parameters[0].converter, self_converter)
        # self is always positional-only.
        assert parameters[0].is_positional_only()
        parameters[0].right_bracket_count = 0
        positional_only = True
        for p in parameters[1:]:
            if not p.is_positional_only():
                positional_only = False
            else:
                assert positional_only
            if positional_only:
                p.right_bracket_count = abs(p.group)
            else:
                # don't put any right brackets around non-positional-only parameters, ever.
                p.right_bracket_count = 0

        right_bracket_count = 0

        def fix_right_bracket_count(desired):
            nonlocal right_bracket_count
            s = ''
            while right_bracket_count < desired:
                s += '['
                right_bracket_count += 1
            while right_bracket_count > desired:
                s += ']'
                right_bracket_count -= 1
            return s

        need_slash = False
        added_slash = False
        need_a_trailing_slash = False

        # we only need a trailing slash:
        #   * if this is not a "docstring_only" signature
        #   * and if the last *shown* parameter is
        #     positional only
        if not f.docstring_only:
            for p in reversed(parameters):
                if not p.converter.show_in_signature:
                    continue
                if p.is_positional_only():
                    need_a_trailing_slash = True
                break


        added_star = False

        first_parameter = True
        last_p = parameters[-1]
        line_length = len(''.join(text))
        indent = " " * line_length
        def add_parameter(text):
            nonlocal line_length
            nonlocal first_parameter
            if first_parameter:
                s = text
                first_parameter = False
            else:
                s = ' ' + text
                if line_length + len(s) >= 72:
                    add('\n')
                    add(indent)
                    line_length = len(indent)
                    s = text
            line_length += len(s)
            add(s)

        for p in parameters:
            if not p.converter.show_in_signature:
                continue
            assert p.name

            is_self = isinstance(p.converter, self_converter)
            if is_self and f.docstring_only:
                # this isn't a real machine-parsable signature,
                # so let's not print the "self" parameter
                continue

            if p.is_positional_only():
                need_slash = not f.docstring_only
            elif need_slash and not (added_slash or p.is_positional_only()):
                added_slash = True
                add_parameter('/,')

            if p.is_keyword_only() and not added_star:
                added_star = True
                add_parameter('*,')

            p_add, p_output = text_accumulator()
            p_add(fix_right_bracket_count(p.right_bracket_count))

            if isinstance(p.converter, self_converter):
                # annotate first parameter as being a "self".
                #
                # if inspect.Signature gets this function,
                # and it's already bound, the self parameter
                # will be stripped off.
                #
                # if it's not bound, it should be marked
                # as positional-only.
                #
                # note: we don't print "self" for __init__,
                # because this isn't actually the signature
                # for __init__.  (it can't be, __init__ doesn't
                # have a docstring.)  if this is an __init__
                # (or __new__), then this signature is for
                # calling the class to construct a new instance.
                p_add('$')

            name = p.converter.signature_name or p.name
            p_add(name)

            if p.converter.is_optional():
                p_add('=')
                value = p.converter.py_default
                if not value:
                    value = repr(p.converter.default)
                p_add(value)

            if (p != last_p) or need_a_trailing_slash:
                p_add(',')

            add_parameter(p_output())

        add(fix_right_bracket_count(0))
        if need_a_trailing_slash:
            add_parameter('/')
        add(')')

        # PEP 8 says:
        #
        #     The Python standard library will not use function annotations
        #     as that would result in a premature commitment to a particular
        #     annotation style. Instead, the annotations are left for users
        #     to discover and experiment with useful annotation styles.
        #
        # therefore this is commented out:
        #
        # if f.return_converter.py_default:
        #     add(' -> ')
        #     add(f.return_converter.py_default)

        if not f.docstring_only:
            add("\n" + sig_end_marker + "\n")

        docstring_first_line = output()

        # now fix up the places where the brackets look wrong
        docstring_first_line = docstring_first_line.replace(', ]', ',] ')

        # okay.  now we're officially building the "parameters" section.
        # create substitution text for {parameters}
        spacer_line = False
        for p in parameters:
            if not p.docstring.strip():
                continue
            if spacer_line:
                add('\n')
            else:
                spacer_line = True
            add("  ")
            add(p.name)
            add('\n')
            add(textwrap.indent(rstrip_lines(p.docstring.rstrip()), "    "))
        parameters = output()
        if parameters:
            parameters += '\n'

        ##
        ## docstring body
        ##

        docstring = f.docstring.rstrip()
        lines = [line.rstrip() for line in docstring.split('\n')]

        # Enforce the summary line!
        # The first line of a docstring should be a summary of the function.
        # It should fit on one line (80 columns? 79 maybe?) and be a paragraph
        # by itself.
        #
        # Argument Clinic enforces the following rule:
        #  * either the docstring is empty,
        #  * or it must have a summary line.
        #
        # Guido said Clinic should enforce this:
        # http://mail.python.org/pipermail/python-dev/2013-June/127110.html

        if len(lines) >= 2:
            if lines[1]:
                fail("Docstring for " + f.full_name + " does not have a summary line!\n" +
                    "Every non-blank function docstring must start with\n" +
                    "a single line summary followed by an empty line.")
        elif len(lines) == 1:
            # the docstring is only one line right now--the summary line.
            # add an empty line after the summary line so we have space
            # between it and the {parameters} we're about to add.
            lines.append('')

        parameters_marker_count = len(docstring.split('{parameters}')) - 1
        if parameters_marker_count > 1:
            fail('You may not specify {parameters} more than once in a docstring!')

        if not parameters_marker_count:
            # insert after summary line
            lines.insert(2, '{parameters}')

        # insert at front of docstring
        lines.insert(0, docstring_first_line)

        docstring = "\n".join(lines)

        add(docstring)
        docstring = output()

        docstring = linear_format(docstring, parameters=parameters)
        docstring = docstring.rstrip()

        return docstring

    def state_terminal(self, line):
        """
        Called when processing the block is done.
        """
        assert not line

        if not self.function:
            return

        if self.keyword_only:
            values = self.function.parameters.values()
            if not values:
                no_parameter_after_star = True
            else:
                last_parameter = next(reversed(list(values)))
                no_parameter_after_star = last_parameter.kind != inspect.Parameter.KEYWORD_ONLY
            if no_parameter_after_star:
                fail("Function " + self.function.name + " specifies '*' without any parameters afterwards.")

        # remove trailing whitespace from all parameter docstrings
        for name, value in self.function.parameters.items():
            if not value:
                continue
            value.docstring = value.docstring.rstrip()

        self.function.docstring = self.format_docstring()




# maps strings to callables.
# the callable should return an object
# that implements the clinic parser
# interface (__init__ and parse).
#
# example parsers:
#   "clinic", handles the Clinic DSL
#   "python", handles running Python code
#
parsers = {'clinic' : DSLParser, 'python': PythonParser}


clinic = None


def main(argv):
    import sys

    if sys.version_info.major < 3 or sys.version_info.minor < 3:
        sys.exit("Error: clinic.py requires Python 3.3 or greater.")

    import argparse
    cmdline = argparse.ArgumentParser(
        description="""Preprocessor for CPython C files.

The purpose of the Argument Clinic is automating all the boilerplate involved
with writing argument parsing code for builtins and providing introspection
signatures ("docstrings") for CPython builtins.

For more information see https://docs.python.org/3/howto/clinic.html""")
    cmdline.add_argument("-f", "--force", action='store_true')
    cmdline.add_argument("-o", "--output", type=str)
    cmdline.add_argument("-v", "--verbose", action='store_true')
    cmdline.add_argument("--converters", action='store_true')
    cmdline.add_argument("--make", action='store_true',
                         help="Walk --srcdir to run over all relevant files.")
    cmdline.add_argument("--srcdir", type=str, default=os.curdir,
                         help="The directory tree to walk in --make mode.")
    cmdline.add_argument("filename", type=str, nargs="*")
    ns = cmdline.parse_args(argv)

    if ns.converters:
        if ns.filename:
            print("Usage error: can't specify --converters and a filename at the same time.")
            print()
            cmdline.print_usage()
            sys.exit(-1)
        converters = []
        return_converters = []
        ignored = set("""
            add_c_converter
            add_c_return_converter
            add_default_legacy_c_converter
            add_legacy_c_converter
            """.strip().split())
        module = globals()
        for name in module:
            for suffix, ids in (
                ("_return_converter", return_converters),
                ("_converter", converters),
            ):
                if name in ignored:
                    continue
                if name.endswith(suffix):
                    ids.append((name, name[:-len(suffix)]))
                    break
        print()

        print("Legacy converters:")
        legacy = sorted(legacy_converters)
        print('    ' + ' '.join(c for c in legacy if c[0].isupper()))
        print('    ' + ' '.join(c for c in legacy if c[0].islower()))
        print()

        for title, attribute, ids in (
            ("Converters", 'converter_init', converters),
            ("Return converters", 'return_converter_init', return_converters),
        ):
            print(title + ":")
            longest = -1
            for name, short_name in ids:
                longest = max(longest, len(short_name))
            for name, short_name in sorted(ids, key=lambda x: x[1].lower()):
                cls = module[name]
                callable = getattr(cls, attribute, None)
                if not callable:
                    continue
                signature = inspect.signature(callable)
                parameters = []
                for parameter_name, parameter in signature.parameters.items():
                    if parameter.kind == inspect.Parameter.KEYWORD_ONLY:
                        if parameter.default != inspect.Parameter.empty:
                            s = '{}={!r}'.format(parameter_name, parameter.default)
                        else:
                            s = parameter_name
                        parameters.append(s)
                print('    {}({})'.format(short_name, ', '.join(parameters)))
            print()
        print("All converters also accept (c_default=None, py_default=None, annotation=None).")
        print("All return converters also accept (py_default=None).")
        sys.exit(0)

    if ns.make:
        if ns.output or ns.filename:
            print("Usage error: can't use -o or filenames with --make.")
            print()
            cmdline.print_usage()
            sys.exit(-1)
        if not ns.srcdir:
            print("Usage error: --srcdir must not be empty with --make.")
            print()
            cmdline.print_usage()
            sys.exit(-1)
        for root, dirs, files in os.walk(ns.srcdir):
            for rcs_dir in ('.svn', '.git', '.hg', 'build', 'externals'):
                if rcs_dir in dirs:
                    dirs.remove(rcs_dir)
            for filename in files:
                if not (filename.endswith('.c') or filename.endswith('.h')):
                    continue
                path = os.path.join(root, filename)
                if ns.verbose:
                    print(path)
                parse_file(path, verify=not ns.force)
        return

    if not ns.filename:
        cmdline.print_usage()
        sys.exit(-1)

    if ns.output and len(ns.filename) > 1:
        print("Usage error: can't use -o with multiple filenames.")
        print()
        cmdline.print_usage()
        sys.exit(-1)

    for filename in ns.filename:
        if ns.verbose:
            print(filename)
        parse_file(filename, output=ns.output, verify=not ns.force)


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))

external/cpython_master/Tools/clinic/clinic.py
DragonHPC/dragon
# Parse Makefiles and Python Setup(.in) files.

import re


# Extract variable definitions from a Makefile.
# Return a dictionary mapping names to values.
# May raise IOError.

makevardef = re.compile('^([a-zA-Z0-9_]+)[ \t]*=(.*)')

def getmakevars(filename):
    variables = {}
    fp = open(filename)
    pendingline = ""
    try:
        while 1:
            line = fp.readline()
            if pendingline:
                line = pendingline + line
                pendingline = ""
            if not line:
                break
            if line.endswith('\\\n'):
                pendingline = line[:-2]
            matchobj = makevardef.match(line)
            if not matchobj:
                continue
            (name, value) = matchobj.group(1, 2)
            # Strip trailing comment
            i = value.find('#')
            if i >= 0:
                value = value[:i]
            value = value.strip()
            variables[name] = value
    finally:
        fp.close()
    return variables


# Parse a Python Setup(.in) file.
# Return two dictionaries, the first mapping modules to their
# definitions, the second mapping variable names to their values.
# May raise IOError.

setupvardef = re.compile('^([a-zA-Z0-9_]+)=(.*)')

def getsetupinfo(filename):
    modules = {}
    variables = {}
    fp = open(filename)
    pendingline = ""
    try:
        while 1:
            line = fp.readline()
            if pendingline:
                line = pendingline + line
                pendingline = ""
            if not line:
                break
            # Strip comments
            i = line.find('#')
            if i >= 0:
                line = line[:i]
            if line.endswith('\\\n'):
                pendingline = line[:-2]
                continue
            matchobj = setupvardef.match(line)
            if matchobj:
                (name, value) = matchobj.group(1, 2)
                variables[name] = value.strip()
            else:
                words = line.split()
                if words:
                    modules[words[0]] = words[1:]
    finally:
        fp.close()
    return modules, variables


# Test the above functions.

def test():
    import sys
    import os
    if not sys.argv[1:]:
        print('usage: python parsesetup.py Makefile*|Setup* ...')
        sys.exit(2)
    for arg in sys.argv[1:]:
        base = os.path.basename(arg)
        if base[:8] == 'Makefile':
            print('Make style parsing:', arg)
            v = getmakevars(arg)
            prdict(v)
        elif base[:5] == 'Setup':
            print('Setup style parsing:', arg)
            m, v = getsetupinfo(arg)
            prdict(m)
            prdict(v)
        else:
            print(arg, 'is neither a Makefile nor a Setup file')
            print('(name must begin with "Makefile" or "Setup")')

def prdict(d):
    keys = sorted(d.keys())
    for key in keys:
        value = d[key]
        print("%-15s" % key, str(value))

if __name__ == '__main__':
    test()

external/cpython_master/Tools/freeze/parsesetup.py
DragonHPC/dragon
import ast
from dataclasses import field, dataclass
import re
from typing import Any, Dict, IO, Optional, List, Text, Tuple, Set
from enum import Enum

from pegen import grammar
from pegen.grammar import (
    Alt,
    Cut,
    Gather,
    GrammarVisitor,
    Group,
    Lookahead,
    NamedItem,
    NameLeaf,
    NegativeLookahead,
    Opt,
    PositiveLookahead,
    Repeat0,
    Repeat1,
    Rhs,
    Rule,
    StringLeaf,
)
from pegen.parser_generator import ParserGenerator


EXTENSION_PREFIX = """\
#include "pegen.h"

#if defined(Py_DEBUG) && defined(Py_BUILD_CORE)
extern int Py_DebugFlag;
#define D(x) if (Py_DebugFlag) x;
#else
#define D(x)
#endif

"""


EXTENSION_SUFFIX = """
void *
_PyPegen_parse(Parser *p)
{
    // Initialize keywords
    p->keywords = reserved_keywords;
    p->n_keyword_lists = n_keyword_lists;

    return start_rule(p);
}
"""


class NodeTypes(Enum):
    NAME_TOKEN = 0
    NUMBER_TOKEN = 1
    STRING_TOKEN = 2
    GENERIC_TOKEN = 3
    KEYWORD = 4
    SOFT_KEYWORD = 5
    CUT_OPERATOR = 6


BASE_NODETYPES = {
    "NAME": NodeTypes.NAME_TOKEN,
    "NUMBER": NodeTypes.NUMBER_TOKEN,
    "STRING": NodeTypes.STRING_TOKEN,
}


@dataclass
class FunctionCall:
    function: str
    arguments: List[Any] = field(default_factory=list)
    assigned_variable: Optional[str] = None
    return_type: Optional[str] = None
    nodetype: Optional[NodeTypes] = None
    force_true: bool = False
    comment: Optional[str] = None

    def __str__(self) -> str:
        parts = []
        parts.append(self.function)
        if self.arguments:
            parts.append(f"({', '.join(map(str, self.arguments))})")
        if self.force_true:
            parts.append(", 1")
        if self.assigned_variable:
            parts = ["(", self.assigned_variable, " = ", *parts, ")"]
        if self.comment:
            parts.append(f"  // {self.comment}")
        return "".join(parts)


class CCallMakerVisitor(GrammarVisitor):
    def __init__(
        self,
        parser_generator: ParserGenerator,
        exact_tokens: Dict[str, int],
        non_exact_tokens: Set[str],
    ):
        self.gen = parser_generator
        self.exact_tokens = exact_tokens
        self.non_exact_tokens = non_exact_tokens
        self.cache: Dict[Any, FunctionCall] = {}
        self.keyword_cache: Dict[str, int] = {}
        self.soft_keywords: Set[str] = set()

    def keyword_helper(self, keyword: str) -> FunctionCall:
        if keyword not in self.keyword_cache:
            self.keyword_cache[keyword] = self.gen.keyword_type()
        return FunctionCall(
            assigned_variable="_keyword",
            function="_PyPegen_expect_token",
            arguments=["p", self.keyword_cache[keyword]],
            return_type="Token *",
            nodetype=NodeTypes.KEYWORD,
            comment=f"token='{keyword}'",
        )

    def soft_keyword_helper(self, value: str) -> FunctionCall:
        self.soft_keywords.add(value.replace('"', ""))
        return FunctionCall(
            assigned_variable="_keyword",
            function="_PyPegen_expect_soft_keyword",
            arguments=["p", value],
            return_type="expr_ty",
            nodetype=NodeTypes.SOFT_KEYWORD,
            comment=f"soft_keyword='{value}'",
        )

    def visit_NameLeaf(self, node: NameLeaf) -> FunctionCall:
        name = node.value
        if name in self.non_exact_tokens:
            if name in BASE_NODETYPES:
                return FunctionCall(
                    assigned_variable=f"{name.lower()}_var",
                    function=f"_PyPegen_{name.lower()}_token",
                    arguments=["p"],
                    nodetype=BASE_NODETYPES[name],
                    return_type="expr_ty",
                    comment=name,
                )
            return FunctionCall(
                assigned_variable=f"{name.lower()}_var",
                function=f"_PyPegen_expect_token",
                arguments=["p", name],
                nodetype=NodeTypes.GENERIC_TOKEN,
                return_type="Token *",
                comment=f"token='{name}'",
            )

        type = None
        rule = self.gen.all_rules.get(name.lower())
        if rule is not None:
            type = "asdl_seq *" if rule.is_loop() or rule.is_gather() else rule.type

        return FunctionCall(
            assigned_variable=f"{name}_var",
            function=f"{name}_rule",
            arguments=["p"],
            return_type=type,
            comment=f"{node}",
        )

    def visit_StringLeaf(self, node: StringLeaf) -> FunctionCall:
        val = ast.literal_eval(node.value)
        if re.match(r"[a-zA-Z_]\w*\Z", val):  # This is a keyword
            if node.value.endswith("'"):
                return self.keyword_helper(val)
            else:
                return self.soft_keyword_helper(node.value)
        else:
            assert val in self.exact_tokens, f"{node.value} is not a known literal"
            type = self.exact_tokens[val]
            return FunctionCall(
                assigned_variable="_literal",
                function=f"_PyPegen_expect_token",
                arguments=["p", type],
                nodetype=NodeTypes.GENERIC_TOKEN,
                return_type="Token *",
                comment=f"token='{val}'",
            )

    def visit_Rhs(self, node: Rhs) -> FunctionCall:
        def can_we_inline(node: Rhs) -> int:
            if len(node.alts) != 1 or len(node.alts[0].items) != 1:
                return False
            # If the alternative has an action we cannot inline
            if getattr(node.alts[0], "action", None) is not None:
                return False
            return True

        if node in self.cache:
            return self.cache[node]
        if can_we_inline(node):
            self.cache[node] = self.generate_call(node.alts[0].items[0])
        else:
            name = self.gen.name_node(node)
            self.cache[node] = FunctionCall(
                assigned_variable=f"{name}_var",
                function=f"{name}_rule",
                arguments=["p"],
                comment=f"{node}",
            )
        return self.cache[node]

    def visit_NamedItem(self, node: NamedItem) -> FunctionCall:
        call = self.generate_call(node.item)
        if node.name:
            call.assigned_variable = node.name
        return call

    def lookahead_call_helper(self, node: Lookahead, positive: int) -> FunctionCall:
        call = self.generate_call(node.node)
        if call.nodetype == NodeTypes.NAME_TOKEN:
            return FunctionCall(
                function=f"_PyPegen_lookahead_with_name",
                arguments=[positive, call.function, *call.arguments],
                return_type="int",
            )
        elif call.nodetype == NodeTypes.SOFT_KEYWORD:
            return FunctionCall(
                function=f"_PyPegen_lookahead_with_string",
                arguments=[positive, call.function, *call.arguments],
                return_type="int",
            )
        elif call.nodetype in {NodeTypes.GENERIC_TOKEN, NodeTypes.KEYWORD}:
            return FunctionCall(
                function=f"_PyPegen_lookahead_with_int",
                arguments=[positive, call.function, *call.arguments],
                return_type="int",
                comment=f"token={node.node}",
            )
        else:
            return FunctionCall(
                function=f"_PyPegen_lookahead",
                arguments=[positive, call.function, *call.arguments],
                return_type="int",
            )

    def visit_PositiveLookahead(self, node: PositiveLookahead) -> FunctionCall:
        return self.lookahead_call_helper(node, 1)

    def visit_NegativeLookahead(self, node: NegativeLookahead) -> FunctionCall:
        return self.lookahead_call_helper(node, 0)

    def visit_Opt(self, node: Opt) -> FunctionCall:
        call = self.generate_call(node.node)
        return FunctionCall(
            assigned_variable="_opt_var",
            function=call.function,
            arguments=call.arguments,
            force_true=True,
            comment=f"{node}",
        )

    def visit_Repeat0(self, node: Repeat0) -> FunctionCall:
        if node in self.cache:
            return self.cache[node]
        name = self.gen.name_loop(node.node, False)
        self.cache[node] = FunctionCall(
            assigned_variable=f"{name}_var",
            function=f"{name}_rule",
            arguments=["p"],
            return_type="asdl_seq *",
            comment=f"{node}",
        )
        return self.cache[node]

    def visit_Repeat1(self, node: Repeat1) -> FunctionCall:
        if node in self.cache:
            return self.cache[node]
        name = self.gen.name_loop(node.node, True)
        self.cache[node] = FunctionCall(
            assigned_variable=f"{name}_var",
            function=f"{name}_rule",
            arguments=["p"],
            return_type="asdl_seq *",
            comment=f"{node}",
        )
        return self.cache[node]

    def visit_Gather(self, node: Gather) -> FunctionCall:
        if node in self.cache:
            return self.cache[node]
        name = self.gen.name_gather(node)
        self.cache[node] = FunctionCall(
            assigned_variable=f"{name}_var",
            function=f"{name}_rule",
            arguments=["p"],
            return_type="asdl_seq *",
            comment=f"{node}",
        )
        return self.cache[node]

    def visit_Group(self, node: Group) -> FunctionCall:
        return self.generate_call(node.rhs)

    def visit_Cut(self, node: Cut) -> FunctionCall:
        return FunctionCall(
            assigned_variable="_cut_var",
            return_type="int",
            function="1",
            nodetype=NodeTypes.CUT_OPERATOR,
        )

    def generate_call(self, node: Any) -> FunctionCall:
        return super().visit(node)


class CParserGenerator(ParserGenerator, GrammarVisitor):
    def __init__(
        self,
        grammar: grammar.Grammar,
        tokens: Dict[int, str],
        exact_tokens: Dict[str, int],
        non_exact_tokens: Set[str],
        file: Optional[IO[Text]],
        debug: bool = False,
        skip_actions: bool = False,
    ):
        super().__init__(grammar, tokens, file)
        self.callmakervisitor: CCallMakerVisitor = CCallMakerVisitor(
            self, exact_tokens, non_exact_tokens
        )
        self._varname_counter = 0
        self.debug = debug
        self.skip_actions = skip_actions

    def add_level(self) -> None:
        self.print("D(p->level++);")

    def remove_level(self) -> None:
        self.print("D(p->level--);")

    def add_return(self, ret_val: str) -> None:
        self.remove_level()
        self.print(f"return {ret_val};")

    def unique_varname(self, name: str = "tmpvar") -> str:
        new_var = name + "_" + str(self._varname_counter)
        self._varname_counter += 1
        return new_var

    def call_with_errorcheck_return(self, call_text: str, returnval: str) -> None:
        error_var = self.unique_varname()
        self.print(f"int {error_var} = {call_text};")
        self.print(f"if ({error_var}) {{")
        with self.indent():
            self.add_return(returnval)
        self.print("}")

    def call_with_errorcheck_goto(self, call_text: str, goto_target: str) -> None:
        error_var = self.unique_varname()
        self.print(f"int {error_var} = {call_text};")
        self.print(f"if ({error_var}) {{")
        with self.indent():
            self.print(f"goto {goto_target};")
        self.print(f"}}")

    def out_of_memory_return(self, expr: str, cleanup_code: Optional[str] = None,) -> None:
        self.print(f"if ({expr}) {{")
        with self.indent():
            if cleanup_code is not None:
                self.print(cleanup_code)
            self.print("p->error_indicator = 1;")
            self.print("PyErr_NoMemory();")
            self.add_return("NULL")
        self.print(f"}}")

    def out_of_memory_goto(self, expr: str, goto_target: str) -> None:
        self.print(f"if ({expr}) {{")
        with self.indent():
            self.print("PyErr_NoMemory();")
            self.print(f"goto {goto_target};")
        self.print(f"}}")

    def generate(self, filename: str) -> None:
        self.collect_todo()
        self.print(f"// @generated by pegen.py from {filename}")
        header = self.grammar.metas.get("header", EXTENSION_PREFIX)
        if header:
            self.print(header.rstrip("\n"))
        subheader = self.grammar.metas.get("subheader", "")
        if subheader:
            self.print(subheader)
        self._setup_keywords()
        for i, (rulename, rule) in enumerate(self.todo.items(), 1000):
            comment = "  // Left-recursive" if rule.left_recursive else ""
            self.print(f"#define {rulename}_type {i}{comment}")
        self.print()
        for rulename, rule in self.todo.items():
            if rule.is_loop() or rule.is_gather():
                type = "asdl_seq *"
            elif rule.type:
                type = rule.type + " "
            else:
                type = "void *"
            self.print(f"static {type}{rulename}_rule(Parser *p);")
        self.print()
        while self.todo:
            for rulename, rule in list(self.todo.items()):
                del self.todo[rulename]
                self.print()
                if rule.left_recursive:
                    self.print("// Left-recursive")
                self.visit(rule)
        if self.skip_actions:
            mode = 0
        else:
            mode = int(self.rules["start"].type == "mod_ty") if "start" in self.rules else 1
            if mode == 1 and self.grammar.metas.get("bytecode"):
                mode += 1
        modulename = self.grammar.metas.get("modulename", "parse")
        trailer = self.grammar.metas.get("trailer", EXTENSION_SUFFIX)
        if trailer:
            self.print(trailer.rstrip("\n") % dict(mode=mode, modulename=modulename))

    def _group_keywords_by_length(self) -> Dict[int, List[Tuple[str, int]]]:
        groups: Dict[int, List[Tuple[str, int]]] = {}
        for keyword_str, keyword_type in self.callmakervisitor.keyword_cache.items():
            length = len(keyword_str)
            if length in groups:
                groups[length].append((keyword_str, keyword_type))
            else:
                groups[length] = [(keyword_str, keyword_type)]
        return groups

    def _setup_keywords(self) -> None:
        keyword_cache = self.callmakervisitor.keyword_cache
        n_keyword_lists = (
            len(max(keyword_cache.keys(), key=len)) + 1 if len(keyword_cache) > 0 else 0
        )
        self.print(f"static const int n_keyword_lists = {n_keyword_lists};")
        groups = self._group_keywords_by_length()
        self.print("static KeywordToken *reserved_keywords[] = {")
        with self.indent():
            num_groups = max(groups) + 1 if groups else 1
            for keywords_length in range(num_groups):
                if keywords_length not in groups.keys():
                    self.print("(KeywordToken[]) {{NULL, -1}},")
                else:
                    self.print("(KeywordToken[]) {")
                    with self.indent():
                        for keyword_str, keyword_type in groups[keywords_length]:
                            self.print(f'{{"{keyword_str}", {keyword_type}}},')
                        self.print("{NULL, -1},")
                    self.print("},")
        self.print("};")

    def _set_up_token_start_metadata_extraction(self) -> None:
        self.print("if (p->mark == p->fill && _PyPegen_fill_token(p) < 0) {")
        with self.indent():
            self.print("p->error_indicator = 1;")
            self.add_return("NULL")
        self.print("}")
        self.print("int _start_lineno = p->tokens[_mark]->lineno;")
        self.print("UNUSED(_start_lineno); // Only used by EXTRA macro")
        self.print("int _start_col_offset = p->tokens[_mark]->col_offset;")
        self.print("UNUSED(_start_col_offset); // Only used by EXTRA macro")

    def _set_up_token_end_metadata_extraction(self) -> None:
        self.print("Token *_token = _PyPegen_get_last_nonnwhitespace_token(p);")
        self.print("if (_token == NULL) {")
        with self.indent():
            self.add_return("NULL")
        self.print("}")
        self.print("int _end_lineno = _token->end_lineno;")
        self.print("UNUSED(_end_lineno); // Only used by EXTRA macro")
        self.print("int _end_col_offset = _token->end_col_offset;")
        self.print("UNUSED(_end_col_offset); // Only used by EXTRA macro")

    def _check_for_errors(self) -> None:
        self.print("if (p->error_indicator) {")
        with self.indent():
            self.add_return("NULL")
        self.print("}")

    def _set_up_rule_memoization(self, node: Rule, result_type: str) -> None:
        self.print("{")
        with self.indent():
            self.add_level()
            self.print(f"{result_type} _res = NULL;")
            self.print(f"if (_PyPegen_is_memoized(p, {node.name}_type, &_res)) {{")
            with self.indent():
                self.add_return("_res")
            self.print("}")
            self.print("int _mark = p->mark;")
            self.print("int _resmark = p->mark;")
            self.print("while (1) {")
            with self.indent():
                self.call_with_errorcheck_return(
                    f"_PyPegen_update_memo(p, _mark, {node.name}_type, _res)", "_res"
                )
                self.print("p->mark = _mark;")
                self.print(f"void *_raw = {node.name}_raw(p);")
                self.print("if (p->error_indicator)")
                with self.indent():
                    self.print("return NULL;")
                self.print("if (_raw == NULL || p->mark <= _resmark)")
                with self.indent():
                    self.print("break;")
                self.print(f"_resmark = p->mark;")
                self.print("_res = _raw;")
            self.print("}")
            self.print(f"p->mark = _resmark;")
            self.add_return("_res")
        self.print("}")
        self.print(f"static {result_type}")
        self.print(f"{node.name}_raw(Parser *p)")

    def _should_memoize(self, node: Rule) -> bool:
        return node.memo and not node.left_recursive

    def _handle_default_rule_body(self, node: Rule, rhs: Rhs, result_type: str) -> None:
        memoize = self._should_memoize(node)

        with self.indent():
            self.add_level()
            self._check_for_errors()
            self.print(f"{result_type} _res = NULL;")
            if memoize:
                self.print(f"if (_PyPegen_is_memoized(p, {node.name}_type, &_res)) {{")
                with self.indent():
                    self.add_return("_res")
                self.print("}")
            self.print("int _mark = p->mark;")
            if any(alt.action and "EXTRA" in alt.action for alt in rhs.alts):
                self._set_up_token_start_metadata_extraction()
            self.visit(
                rhs, is_loop=False, is_gather=node.is_gather(), rulename=node.name,
            )
            if self.debug:
                self.print(f'D(fprintf(stderr, "Fail at %d: {node.name}\\n", p->mark));')
            self.print("_res = NULL;")
        self.print("  done:")
        with self.indent():
            if memoize:
                self.print(f"_PyPegen_insert_memo(p, _mark, {node.name}_type, _res);")
            self.add_return("_res")

    def _handle_loop_rule_body(self, node: Rule, rhs: Rhs) -> None:
        memoize = self._should_memoize(node)
        is_repeat1 = node.name.startswith("_loop1")

        with self.indent():
            self.add_level()
            self._check_for_errors()
            self.print("void *_res = NULL;")
            if memoize:
                self.print(f"if (_PyPegen_is_memoized(p, {node.name}_type, &_res)) {{")
                with self.indent():
                    self.add_return("_res")
                self.print("}")
            self.print("int _mark = p->mark;")
            self.print("int _start_mark = p->mark;")
            self.print("void **_children = PyMem_Malloc(sizeof(void *));")
            self.out_of_memory_return(f"!_children")
            self.print("ssize_t _children_capacity = 1;")
            self.print("ssize_t _n = 0;")
            if any(alt.action and "EXTRA" in alt.action for alt in rhs.alts):
                self._set_up_token_start_metadata_extraction()
            self.visit(
                rhs, is_loop=True, is_gather=node.is_gather(), rulename=node.name,
            )
            if is_repeat1:
                self.print("if (_n == 0 || p->error_indicator) {")
                with self.indent():
                    self.print("PyMem_Free(_children);")
                    self.add_return("NULL")
                self.print("}")
            self.print("asdl_seq *_seq = _Py_asdl_seq_new(_n, p->arena);")
            self.out_of_memory_return(f"!_seq", cleanup_code="PyMem_Free(_children);")
            self.print("for (int i = 0; i < _n; i++) asdl_seq_SET(_seq, i, _children[i]);")
            self.print("PyMem_Free(_children);")
            if node.name:
                self.print(f"_PyPegen_insert_memo(p, _start_mark, {node.name}_type, _seq);")
            self.add_return("_seq")

    def visit_Rule(self, node: Rule) -> None:
        is_loop = node.is_loop()
        is_gather = node.is_gather()
        rhs = node.flatten()
        if is_loop or is_gather:
            result_type = "asdl_seq *"
        elif node.type:
            result_type = node.type
        else:
            result_type = "void *"

        for line in str(node).splitlines():
            self.print(f"// {line}")
        if node.left_recursive and node.leader:
            self.print(f"static {result_type} {node.name}_raw(Parser *);")

        self.print(f"static {result_type}")
        self.print(f"{node.name}_rule(Parser *p)")

        if node.left_recursive and node.leader:
            self._set_up_rule_memoization(node, result_type)

        self.print("{")
        if is_loop:
            self._handle_loop_rule_body(node, rhs)
        else:
            self._handle_default_rule_body(node, rhs, result_type)
        self.print("}")

    def visit_NamedItem(self, node: NamedItem) -> None:
        call = self.callmakervisitor.generate_call(node)
        if call.assigned_variable:
            call.assigned_variable = self.dedupe(call.assigned_variable)
        self.print(call)

    def visit_Rhs(
        self, node: Rhs, is_loop: bool, is_gather: bool, rulename: Optional[str]
    ) -> None:
        if is_loop:
            assert len(node.alts) == 1
        for alt in node.alts:
            self.visit(alt, is_loop=is_loop, is_gather=is_gather, rulename=rulename)

    def join_conditions(self, keyword: str, node: Any) -> None:
        self.print(f"{keyword} (")
        with self.indent():
            first = True
            for item in node.items:
                if first:
                    first = False
                else:
                    self.print("&&")
                self.visit(item)
        self.print(")")

    def emit_action(self, node: Alt, cleanup_code: Optional[str] = None) -> None:
        self.print(f"_res = {node.action};")

        self.print("if (_res == NULL && PyErr_Occurred()) {")
        with self.indent():
            self.print("p->error_indicator = 1;")
            if cleanup_code:
                self.print(cleanup_code)
            self.add_return("NULL")
        self.print("}")

        if self.debug:
            self.print(
                f'D(fprintf(stderr, "Hit with action [%d-%d]: %s\\n", _mark, p->mark, "{node}"));'
            )

    def emit_default_action(self, is_gather: bool, node: Alt) -> None:
        if len(self.local_variable_names) > 1:
            if is_gather:
                assert len(self.local_variable_names) == 2
                self.print(
                    f"_res = _PyPegen_seq_insert_in_front(p, "
                    f"{self.local_variable_names[0]}, {self.local_variable_names[1]});"
                )
            else:
                if self.debug:
                    self.print(
                        f'D(fprintf(stderr, "Hit without action [%d:%d]: %s\\n", _mark, p->mark, "{node}"));'
                    )
                self.print(
                    f"_res = _PyPegen_dummy_name(p, {', '.join(self.local_variable_names)});"
                )
        else:
            if self.debug:
                self.print(
                    f'D(fprintf(stderr, "Hit with default action [%d:%d]: %s\\n", _mark, p->mark, "{node}"));'
                )
            self.print(f"_res = {self.local_variable_names[0]};")

    def emit_dummy_action(self) -> None:
        self.print("_res = _PyPegen_dummy_name(p);")

    def handle_alt_normal(self, node: Alt, is_gather: bool, rulename: Optional[str]) -> None:
        self.join_conditions(keyword="if", node=node)
        self.print("{")
        # We have parsed successfully all the conditions for the option.
        with self.indent():
            node_str = str(node).replace('"', '\\"')
            self.print(
                f'D(fprintf(stderr, "%*c+ {rulename}[%d-%d]: %s succeeded!\\n", p->level, \' \', _mark, p->mark, "{node_str}"));'
            )
            # Prepare to emmit the rule action and do so
            if node.action and "EXTRA" in node.action:
                self._set_up_token_end_metadata_extraction()
            if self.skip_actions:
                self.emit_dummy_action()
            elif node.action:
                self.emit_action(node)
            else:
                self.emit_default_action(is_gather, node)

            # As the current option has parsed correctly, do not continue with the rest.
            self.print(f"goto done;")
        self.print("}")

    def handle_alt_loop(self, node: Alt, is_gather: bool, rulename: Optional[str]) -> None:
        # Condition of the main body of the alternative
        self.join_conditions(keyword="while", node=node)
        self.print("{")
        # We have parsed successfully one item!
        with self.indent():
            # Prepare to emit the rule action and do so
            if node.action and "EXTRA" in node.action:
                self._set_up_token_end_metadata_extraction()
            if self.skip_actions:
                self.emit_dummy_action()
            elif node.action:
                self.emit_action(node, cleanup_code="PyMem_Free(_children);")
            else:
                self.emit_default_action(is_gather, node)

            # Add the result of rule to the temporary buffer of children. This buffer
            # will populate later an asdl_seq with all elements to return.
            self.print("if (_n == _children_capacity) {")
            with self.indent():
                self.print("_children_capacity *= 2;")
                self.print(
                    "void **_new_children = PyMem_Realloc(_children, _children_capacity*sizeof(void *));"
                )
                self.out_of_memory_return(f"!_new_children")
                self.print("_children = _new_children;")
            self.print("}")
            self.print("_children[_n++] = _res;")
            self.print("_mark = p->mark;")
        self.print("}")

    def visit_Alt(
        self, node: Alt, is_loop: bool, is_gather: bool, rulename: Optional[str]
    ) -> None:
        if len(node.items) == 1 and str(node.items[0]).startswith('invalid_'):
            self.print(f"if (p->call_invalid_rules) {{ // {node}")
        else:
            self.print(f"{{ // {node}")
        with self.indent():
            self._check_for_errors()
            node_str = str(node).replace('"', '\\"')
            self.print(
                f'D(fprintf(stderr, "%*c> {rulename}[%d-%d]: %s\\n", p->level, \' \', _mark, p->mark, "{node_str}"));'
            )
            # Prepare variable declarations for the alternative
            vars = self.collect_vars(node)
            for v, var_type in sorted(item for item in vars.items() if item[0] is not None):
                if not var_type:
                    var_type = "void *"
                else:
                    var_type += " "
                if v == "_cut_var":
                    v += " = 0"  # cut_var must be initialized
                self.print(f"{var_type}{v};")
                if v.startswith("_opt_var"):
                    self.print(f"UNUSED({v}); // Silence compiler warnings")

            with self.local_variable_context():
                if is_loop:
                    self.handle_alt_loop(node, is_gather, rulename)
                else:
                    self.handle_alt_normal(node, is_gather, rulename)

            self.print("p->mark = _mark;")
            node_str = str(node).replace('"', '\\"')
            self.print(
                f"D(fprintf(stderr, \"%*c%s {rulename}[%d-%d]: %s failed!\\n\", p->level, ' ',\n"
                f'                  p->error_indicator ? "ERROR!" : "-", _mark, p->mark, "{node_str}"));'
            )
            if "_cut_var" in vars:
                self.print("if (_cut_var) {")
                with self.indent():
                    self.add_return("NULL")
                self.print("}")
        self.print("}")

    def collect_vars(self, node: Alt) -> Dict[Optional[str], Optional[str]]:
        types = {}
        with self.local_variable_context():
            for item in node.items:
                name, type = self.add_var(item)
                types[name] = type
        return types

    def add_var(self, node: NamedItem) -> Tuple[Optional[str], Optional[str]]:
        call = self.callmakervisitor.generate_call(node.item)
        name = node.name if node.name else call.assigned_variable
        if name is not None:
            name = self.dedupe(name)
        return name, call.return_type

external/cpython_master/Tools/peg_generator/pegen/c_generator.py
DragonHPC/dragon
#!/usr/bin/env python3.8
# @generated by pegen from ./Tools/peg_generator/pegen/metagrammar.gram

import ast
import sys
import tokenize

from typing import Any, Optional

from pegen.parser import memoize, memoize_left_rec, logger, Parser
from ast import literal_eval

from pegen.grammar import (
    Alt,
    Cut,
    Gather,
    Group,
    Item,
    Lookahead,
    LookaheadOrCut,
    MetaTuple,
    MetaList,
    NameLeaf,
    NamedItem,
    NamedItemList,
    NegativeLookahead,
    Opt,
    Plain,
    PositiveLookahead,
    Repeat0,
    Repeat1,
    Rhs,
    Rule,
    RuleList,
    RuleName,
    Grammar,
    StringLeaf,
)

class GeneratedParser(Parser):

    @memoize
    def start(self) -> Optional[Grammar]:
        # start: grammar $
        mark = self.mark()
        cut = False
        if (
            (grammar := self.grammar())
            and
            (endmarker := self.expect('ENDMARKER'))
        ):
            return grammar
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def grammar(self) -> Optional[Grammar]:
        # grammar: metas rules | rules
        mark = self.mark()
        cut = False
        if (
            (metas := self.metas())
            and
            (rules := self.rules())
        ):
            return Grammar ( rules , metas )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (rules := self.rules())
        ):
            return Grammar ( rules , [ ] )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def metas(self) -> Optional[MetaList]:
        # metas: meta metas | meta
        mark = self.mark()
        cut = False
        if (
            (meta := self.meta())
            and
            (metas := self.metas())
        ):
            return [ meta ] + metas
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (meta := self.meta())
        ):
            return [ meta ]
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def meta(self) -> Optional[MetaTuple]:
        # meta: "@" NAME NEWLINE | "@" NAME NAME NEWLINE | "@" NAME STRING NEWLINE
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect("@"))
            and
            (name := self.name())
            and
            (newline := self.expect('NEWLINE'))
        ):
            return ( name . string , None )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (literal := self.expect("@"))
            and
            (a := self.name())
            and
            (b := self.name())
            and
            (newline := self.expect('NEWLINE'))
        ):
            return ( a . string , b . string )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (literal := self.expect("@"))
            and
            (name := self.name())
            and
            (string := self.string())
            and
            (newline := self.expect('NEWLINE'))
        ):
            return ( name . string , literal_eval ( string . string ) )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def rules(self) -> Optional[RuleList]:
        # rules: rule rules | rule
        mark = self.mark()
        cut = False
        if (
            (rule := self.rule())
            and
            (rules := self.rules())
        ):
            return [ rule ] + rules
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (rule := self.rule())
        ):
            return [ rule ]
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def rule(self) -> Optional[Rule]:
        # rule: rulename memoflag? ":" alts NEWLINE INDENT more_alts DEDENT | rulename memoflag? ":" NEWLINE INDENT more_alts DEDENT | rulename memoflag? ":" alts NEWLINE
        mark = self.mark()
        cut = False
        if (
            (rulename := self.rulename())
            and
            (opt := self.memoflag(),)
            and
            (literal := self.expect(":"))
            and
            (alts := self.alts())
            and
            (newline := self.expect('NEWLINE'))
            and
            (indent := self.expect('INDENT'))
            and
            (more_alts := self.more_alts())
            and
            (dedent := self.expect('DEDENT'))
        ):
            return Rule ( rulename [ 0 ] , rulename [ 1 ] , Rhs ( alts . alts + more_alts . alts ) , memo = opt )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (rulename := self.rulename())
            and
            (opt := self.memoflag(),)
            and
            (literal := self.expect(":"))
            and
            (newline := self.expect('NEWLINE'))
            and
            (indent := self.expect('INDENT'))
            and
            (more_alts := self.more_alts())
            and
            (dedent := self.expect('DEDENT'))
        ):
            return Rule ( rulename [ 0 ] , rulename [ 1 ] , more_alts , memo = opt )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (rulename := self.rulename())
            and
            (opt := self.memoflag(),)
            and
            (literal := self.expect(":"))
            and
            (alts := self.alts())
            and
            (newline := self.expect('NEWLINE'))
        ):
            return Rule ( rulename [ 0 ] , rulename [ 1 ] , alts , memo = opt )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def rulename(self) -> Optional[RuleName]:
        # rulename: NAME '[' NAME '*' ']' | NAME '[' NAME ']' | NAME
        mark = self.mark()
        cut = False
        if (
            (name := self.name())
            and
            (literal := self.expect('['))
            and
            (type := self.name())
            and
            (literal_1 := self.expect('*'))
            and
            (literal_2 := self.expect(']'))
        ):
            return ( name . string , type . string + "*" )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (name := self.name())
            and
            (literal := self.expect('['))
            and
            (type := self.name())
            and
            (literal_1 := self.expect(']'))
        ):
            return ( name . string , type . string )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (name := self.name())
        ):
            return ( name . string , None )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def memoflag(self) -> Optional[str]:
        # memoflag: '(' 'memo' ')'
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect('('))
            and
            (literal_1 := self.expect('memo'))
            and
            (literal_2 := self.expect(')'))
        ):
            return "memo"
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def alts(self) -> Optional[Rhs]:
        # alts: alt "|" alts | alt
        mark = self.mark()
        cut = False
        if (
            (alt := self.alt())
            and
            (literal := self.expect("|"))
            and
            (alts := self.alts())
        ):
            return Rhs ( [ alt ] + alts . alts )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (alt := self.alt())
        ):
            return Rhs ( [ alt ] )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def more_alts(self) -> Optional[Rhs]:
        # more_alts: "|" alts NEWLINE more_alts | "|" alts NEWLINE
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect("|"))
            and
            (alts := self.alts())
            and
            (newline := self.expect('NEWLINE'))
            and
            (more_alts := self.more_alts())
        ):
            return Rhs ( alts . alts + more_alts . alts )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (literal := self.expect("|"))
            and
            (alts := self.alts())
            and
            (newline := self.expect('NEWLINE'))
        ):
            return Rhs ( alts . alts )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def alt(self) -> Optional[Alt]:
        # alt: items '$' action | items '$' | items action | items
        mark = self.mark()
        cut = False
        if (
            (items := self.items())
            and
            (literal := self.expect('$'))
            and
            (action := self.action())
        ):
            return Alt ( items + [ NamedItem ( None , NameLeaf ( 'ENDMARKER' ) ) ] , action = action )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (items := self.items())
            and
            (literal := self.expect('$'))
        ):
            return Alt ( items + [ NamedItem ( None , NameLeaf ( 'ENDMARKER' ) ) ] , action = None )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (items := self.items())
            and
            (action := self.action())
        ):
            return Alt ( items , action = action )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (items := self.items())
        ):
            return Alt ( items , action = None )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def items(self) -> Optional[NamedItemList]:
        # items: named_item items | named_item
        mark = self.mark()
        cut = False
        if (
            (named_item := self.named_item())
            and
            (items := self.items())
        ):
            return [ named_item ] + items
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (named_item := self.named_item())
        ):
            return [ named_item ]
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def named_item(self) -> Optional[NamedItem]:
        # named_item: NAME '=' ~ item | item | lookahead
        mark = self.mark()
        cut = False
        if (
            (name := self.name())
            and
            (literal := self.expect('='))
            and
            (cut := True)
            and
            (item := self.item())
        ):
            return NamedItem ( name . string , item )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (item := self.item())
        ):
            return NamedItem ( None , item )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (it := self.lookahead())
        ):
            return NamedItem ( None , it )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def lookahead(self) -> Optional[LookaheadOrCut]:
        # lookahead: '&' ~ atom | '!' ~ atom | '~'
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect('&'))
            and
            (cut := True)
            and
            (atom := self.atom())
        ):
            return PositiveLookahead ( atom )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (literal := self.expect('!'))
            and
            (cut := True)
            and
            (atom := self.atom())
        ):
            return NegativeLookahead ( atom )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (literal := self.expect('~'))
        ):
            return Cut ( )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def item(self) -> Optional[Item]:
        # item: '[' ~ alts ']' | atom '?' | atom '*' | atom '+' | atom '.' atom '+' | atom
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect('['))
            and
            (cut := True)
            and
            (alts := self.alts())
            and
            (literal_1 := self.expect(']'))
        ):
            return Opt ( alts )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (atom := self.atom())
            and
            (literal := self.expect('?'))
        ):
            return Opt ( atom )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (atom := self.atom())
            and
            (literal := self.expect('*'))
        ):
            return Repeat0 ( atom )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (atom := self.atom())
            and
            (literal := self.expect('+'))
        ):
            return Repeat1 ( atom )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (sep := self.atom())
            and
            (literal := self.expect('.'))
            and
            (node := self.atom())
            and
            (literal_1 := self.expect('+'))
        ):
            return Gather ( sep , node )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (atom := self.atom())
        ):
            return atom
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def atom(self) -> Optional[Plain]:
        # atom: '(' ~ alts ')' | NAME | STRING
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect('('))
            and
            (cut := True)
            and
            (alts := self.alts())
            and
            (literal_1 := self.expect(')'))
        ):
            return Group ( alts )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (name := self.name())
        ):
            return NameLeaf ( name . string )
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (string := self.string())
        ):
            return StringLeaf ( string . string )
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def action(self) -> Optional[str]:
        # action: "{" ~ target_atoms "}"
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect("{"))
            and
            (cut := True)
            and
            (target_atoms := self.target_atoms())
            and
            (literal_1 := self.expect("}"))
        ):
            return target_atoms
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def target_atoms(self) -> Optional[str]:
        # target_atoms: target_atom target_atoms | target_atom
        mark = self.mark()
        cut = False
        if (
            (target_atom := self.target_atom())
            and
            (target_atoms := self.target_atoms())
        ):
            return target_atom + " " + target_atoms
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (target_atom := self.target_atom())
        ):
            return target_atom
        self.reset(mark)
        if cut: return None
        return None

    @memoize
    def target_atom(self) -> Optional[str]:
        # target_atom: "{" ~ target_atoms "}" | NAME | NUMBER | STRING | "?" | ":" | !"}" OP
        mark = self.mark()
        cut = False
        if (
            (literal := self.expect("{"))
            and
            (cut := True)
            and
            (target_atoms := self.target_atoms())
            and
            (literal_1 := self.expect("}"))
        ):
            return "{" + target_atoms + "}"
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (name := self.name())
        ):
            return name . string
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (number := self.number())
        ):
            return number . string
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (string := self.string())
        ):
            return string . string
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (literal := self.expect("?"))
        ):
            return "?"
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            (literal := self.expect(":"))
        ):
            return ":"
        self.reset(mark)
        if cut: return None
        cut = False
        if (
            self.negative_lookahead(self.expect, "}")
            and
            (op := self.op())
        ):
            return op . string
        self.reset(mark)
        if cut: return None
        return None


if __name__ == '__main__':
    from pegen.parser import simple_parser_main
    simple_parser_main(GeneratedParser)

external/cpython_master/Tools/peg_generator/pegen/grammar_parser.py
DragonHPC/dragon
import token
from typing import Any, Dict, Optional, IO, Text, Tuple

from pegen.grammar import (
    Cut,
    GrammarVisitor,
    NameLeaf,
    StringLeaf,
    Rhs,
    NamedItem,
    Lookahead,
    PositiveLookahead,
    NegativeLookahead,
    Opt,
    Repeat0,
    Repeat1,
    Gather,
    Group,
    Rule,
    Alt,
)
from pegen import grammar
from pegen.parser_generator import ParserGenerator

MODULE_PREFIX = """\
#!/usr/bin/env python3.8
# @generated by pegen from {filename}

import ast
import sys
import tokenize

from typing import Any, Optional

from pegen.parser import memoize, memoize_left_rec, logger, Parser

"""
MODULE_SUFFIX = """

if __name__ == '__main__':
    from pegen.parser import simple_parser_main
    simple_parser_main(GeneratedParser)
"""


class PythonCallMakerVisitor(GrammarVisitor):
    def __init__(self, parser_generator: ParserGenerator):
        self.gen = parser_generator
        self.cache: Dict[Any, Any] = {}

    def visit_NameLeaf(self, node: NameLeaf) -> Tuple[Optional[str], str]:
        name = node.value
        if name in ("NAME", "NUMBER", "STRING", "OP"):
            name = name.lower()
            return name, f"self.{name}()"
        if name in ("NEWLINE", "DEDENT", "INDENT", "ENDMARKER", "ASYNC", "AWAIT"):
            return name.lower(), f"self.expect({name!r})"
        return name, f"self.{name}()"

    def visit_StringLeaf(self, node: StringLeaf) -> Tuple[str, str]:
        return "literal", f"self.expect({node.value})"

    def visit_Rhs(self, node: Rhs) -> Tuple[Optional[str], str]:
        if node in self.cache:
            return self.cache[node]
        if len(node.alts) == 1 and len(node.alts[0].items) == 1:
            self.cache[node] = self.visit(node.alts[0].items[0])
        else:
            name = self.gen.name_node(node)
            self.cache[node] = name, f"self.{name}()"
        return self.cache[node]

    def visit_NamedItem(self, node: NamedItem) -> Tuple[Optional[str], str]:
        name, call = self.visit(node.item)
        if node.name:
            name = node.name
        return name, call

    def lookahead_call_helper(self, node: Lookahead) -> Tuple[str, str]:
        name, call = self.visit(node.node)
        head, tail = call.split("(", 1)
        assert tail[-1] == ")"
        tail = tail[:-1]
        return head, tail

    def visit_PositiveLookahead(self, node: PositiveLookahead) -> Tuple[None, str]:
        head, tail = self.lookahead_call_helper(node)
        return None, f"self.positive_lookahead({head}, {tail})"

    def visit_NegativeLookahead(self, node: NegativeLookahead) -> Tuple[None, str]:
        head, tail = self.lookahead_call_helper(node)
        return None, f"self.negative_lookahead({head}, {tail})"

    def visit_Opt(self, node: Opt) -> Tuple[str, str]:
        name, call = self.visit(node.node)
        # Note trailing comma (the call may already have one comma
        # at the end, for example when rules have both repeat0 and optional
        # markers, e.g: [rule*])
        if call.endswith(","):
            return "opt", call
        else:
            return "opt", f"{call},"

    def visit_Repeat0(self, node: Repeat0) -> Tuple[str, str]:
        if node in self.cache:
            return self.cache[node]
        name = self.gen.name_loop(node.node, False)
        self.cache[node] = name, f"self.{name}(),"  # Also a trailing comma!
        return self.cache[node]

    def visit_Repeat1(self, node: Repeat1) -> Tuple[str, str]:
        if node in self.cache:
            return self.cache[node]
        name = self.gen.name_loop(node.node, True)
        self.cache[node] = name, f"self.{name}()"  # But no trailing comma here!
        return self.cache[node]

    def visit_Gather(self, node: Gather) -> Tuple[str, str]:
        if node in self.cache:
            return self.cache[node]
        name = self.gen.name_gather(node)
        self.cache[node] = name, f"self.{name}()"  # No trailing comma here either!
        return self.cache[node]

    def visit_Group(self, node: Group) -> Tuple[Optional[str], str]:
        return self.visit(node.rhs)

    def visit_Cut(self, node: Cut) -> Tuple[str, str]:
        return "cut", "True"


class PythonParserGenerator(ParserGenerator, GrammarVisitor):
    def __init__(
        self,
        grammar: grammar.Grammar,
        file: Optional[IO[Text]],
        tokens: Dict[int, str] = token.tok_name,
    ):
        super().__init__(grammar, tokens, file)
        self.callmakervisitor = PythonCallMakerVisitor(self)

    def generate(self, filename: str) -> None:
        header = self.grammar.metas.get("header", MODULE_PREFIX)
        if header is not None:
            self.print(header.rstrip("\n").format(filename=filename))
        subheader = self.grammar.metas.get("subheader", "")
        if subheader:
            self.print(subheader.format(filename=filename))
        self.print("class GeneratedParser(Parser):")
        while self.todo:
            for rulename, rule in list(self.todo.items()):
                del self.todo[rulename]
                self.print()
                with self.indent():
                    self.visit(rule)
        trailer = self.grammar.metas.get("trailer", MODULE_SUFFIX)
        if trailer is not None:
            self.print(trailer.rstrip("\n"))

    def visit_Rule(self, node: Rule) -> None:
        is_loop = node.is_loop()
        is_gather = node.is_gather()
        rhs = node.flatten()
        if node.left_recursive:
            if node.leader:
                self.print("@memoize_left_rec")
            else:
                # Non-leader rules in a cycle are not memoized,
                # but they must still be logged.
                self.print("@logger")
        else:
            self.print("@memoize")
        node_type = node.type or "Any"
        self.print(f"def {node.name}(self) -> Optional[{node_type}]:")
        with self.indent():
            self.print(f"# {node.name}: {rhs}")
            if node.nullable:
                self.print(f"# nullable={node.nullable}")
            self.print("mark = self.mark()")
            if is_loop:
                self.print("children = []")
            self.visit(rhs, is_loop=is_loop, is_gather=is_gather)
            if is_loop:
                self.print("return children")
            else:
                self.print("return None")

    def visit_NamedItem(self, node: NamedItem) -> None:
        name, call = self.callmakervisitor.visit(node.item)
        if node.name:
            name = node.name
        if not name:
            self.print(call)
        else:
            if name != "cut":
                name = self.dedupe(name)
            self.print(f"({name} := {call})")

    def visit_Rhs(self, node: Rhs, is_loop: bool = False, is_gather: bool = False) -> None:
        if is_loop:
            assert len(node.alts) == 1
        for alt in node.alts:
            self.visit(alt, is_loop=is_loop, is_gather=is_gather)

    def visit_Alt(self, node: Alt, is_loop: bool, is_gather: bool) -> None:
        with self.local_variable_context():
            self.print("cut = False")  # TODO: Only if needed.
            if is_loop:
                self.print("while (")
            else:
                self.print("if (")
            with self.indent():
                first = True
                for item in node.items:
                    if first:
                        first = False
                    else:
                        self.print("and")
                    self.visit(item)
                    if is_gather:
                        self.print("is not None")

            self.print("):")
            with self.indent():
                action = node.action
                if not action:
                    if is_gather:
                        assert len(self.local_variable_names) == 2
                        action = (
                            f"[{self.local_variable_names[0]}] + {self.local_variable_names[1]}"
                        )
                    else:
                        action = f"[{', '.join(self.local_variable_names)}]"
                if is_loop:
                    self.print(f"children.append({action})")
                    self.print(f"mark = self.mark()")
                else:
                    self.print(f"return {action}")
            self.print("self.reset(mark)")
            # Skip remaining alternatives if a cut was reached.
            self.print("if cut: return None")  # TODO: Only if needed.

external/cpython_master/Tools/peg_generator/pegen/python_generator.py
DragonHPC/dragon
"""Switchboard class.

This class is used to coordinate updates among all Viewers.  Every Viewer must
conform to the following interface:

    - it must include a method called update_yourself() which takes three
      arguments; the red, green, and blue values of the selected color.

    - When a Viewer selects a color and wishes to update all other Views, it
      should call update_views() on the Switchboard object.  Note that the
      Viewer typically does *not* update itself before calling update_views(),
      since this would cause it to get updated twice.

Optionally, Viewers can also implement:

    - save_options() which takes an optiondb (a dictionary).  Store into this
      dictionary any values the Viewer wants to save in the persistent
      ~/.pynche file.  This dictionary is saved using marshal.  The namespace
      for the keys is ad-hoc; make sure you don't clobber some other Viewer's
      keys!

    - withdraw() which takes no arguments.  This is called when Pynche is
      unmapped.  All Viewers should implement this.

    - colordb_changed() which takes a single argument, an instance of
      ColorDB.  This is called whenever the color name database is changed and
      gives a chance for the Viewers to do something on those events.  See
      ListViewer for details.

External Viewers are found dynamically.  Viewer modules should have names such
as FooViewer.py.  If such a named module has a module global variable called
ADDTOVIEW and this variable is true, the Viewer will be added dynamically to
the `View' menu.  ADDTOVIEW contains a string which is used as the menu item
to display the Viewer (one kludge: if the string contains a `%', this is used
to indicate that the next character will get an underline in the menu,
otherwise the first character is underlined).

FooViewer.py should contain a class called FooViewer, and its constructor
should take two arguments, an instance of Switchboard, and optionally a Tk
master window.

"""

import sys
import marshal



class Switchboard:
    def __init__(self, initfile):
        self.__initfile = initfile
        self.__colordb = None
        self.__optiondb = {}
        self.__views = []
        self.__red = 0
        self.__green = 0
        self.__blue = 0
        self.__canceled = 0
        # read the initialization file
        fp = None
        if initfile:
            try:
                try:
                    fp = open(initfile, 'rb')
                    self.__optiondb = marshal.load(fp)
                    if not isinstance(self.__optiondb, dict):
                        print('Problem reading options from file:', initfile,
                              file=sys.stderr)
                        self.__optiondb = {}
                except (IOError, EOFError, ValueError):
                    pass
            finally:
                if fp:
                    fp.close()

    def add_view(self, view):
        self.__views.append(view)

    def update_views(self, red, green, blue):
        self.__red = red
        self.__green = green
        self.__blue = blue
        for v in self.__views:
            v.update_yourself(red, green, blue)

    def update_views_current(self):
        self.update_views(self.__red, self.__green, self.__blue)

    def current_rgb(self):
        return self.__red, self.__green, self.__blue

    def colordb(self):
        return self.__colordb

    def set_colordb(self, colordb):
        self.__colordb = colordb
        for v in self.__views:
            if hasattr(v, 'colordb_changed'):
                v.colordb_changed(colordb)
        self.update_views_current()

    def optiondb(self):
        return self.__optiondb

    def save_views(self):
        # save the current color
        self.__optiondb['RED'] = self.__red
        self.__optiondb['GREEN'] = self.__green
        self.__optiondb['BLUE'] = self.__blue
        for v in self.__views:
            if hasattr(v, 'save_options'):
                v.save_options(self.__optiondb)
        # save the name of the file used for the color database.  we'll try to
        # load this first.
        self.__optiondb['DBFILE'] = self.__colordb.filename()
        fp = None
        try:
            try:
                fp = open(self.__initfile, 'wb')
            except IOError:
                print('Cannot write options to file:', \
                      self.__initfile, file=sys.stderr)
            else:
                marshal.dump(self.__optiondb, fp)
        finally:
            if fp:
                fp.close()

    def withdraw_views(self):
        for v in self.__views:
            if hasattr(v, 'withdraw'):
                v.withdraw()

    def canceled(self, flag=1):
        self.__canceled = flag

    def canceled_p(self):
        return self.__canceled

external/cpython_master/Tools/pynche/Switchboard.py
DragonHPC/dragon
#! /usr/bin/env python3

"""
This script should be called *manually* when we want to upgrade SSLError
`library` and `reason` mnemonics to a more recent OpenSSL version.

It takes two arguments:
- the path to the OpenSSL source tree (e.g. git checkout)
- the path to the C file to be generated
  (probably Modules/_ssl_data.h)
"""

import datetime
import glob
import os
import re
import sys
import _ssl


def parse_error_codes(h_file, prefix, libcode):
    pat = re.compile(r"#\s*define\W+(%s([\w]+))\W+(\d+)\b" % re.escape(prefix))
    codes = []
    with open(h_file, "r", encoding="latin1") as f:
        for line in f:
            match = pat.search(line)
            if match:
                code, name, num = match.groups()
                num = int(num)
                # e.g. ("SSL_R_BAD_DATA", ("ERR_LIB_SSL", "BAD_DATA", 390))
                codes.append((code, (libcode, name, num)))
    assert codes, f"no codes found in {h_file}"
    return codes

if __name__ == "__main__":
    openssl_inc = sys.argv[1]
    outfile = sys.argv[2]
    use_stdout = outfile == '-'
    f = sys.stdout if use_stdout else open(outfile, "w")
    # mnemonic -> (library code, error prefix, header file)
    error_libraries = {}
    for error_header in glob.glob(os.path.join(glob.escape(openssl_inc), 'include/openssl/*err.h')):
        base = os.path.basename(error_header)
        if base in ('buffererr.h', 'objectserr.h', 'storeerr.h'):
            # Deprecated in 3.0.
            continue
        mnemonic = base[:-5].upper()
        if mnemonic == "":
            # err.h
            lib_codes = {
                code: num
                for (code, (_, _, num)) in parse_error_codes(error_header, 'ERR_LIB_', None)
            }
        else:
            error_libraries[mnemonic] = (f'ERR_LIB_{mnemonic}', f'{mnemonic}_R_', error_header)

    # Read codes from libraries
    new_codes = []
    for libcode, prefix, h_file in sorted(error_libraries.values()):
        new_codes += parse_error_codes(h_file, prefix, libcode)
    new_code_nums = set((libcode, num)
                        for (code, (libcode, name, num)) in new_codes)

    # Merge with existing codes (in case some old codes disappeared).
    codes = {}
    for errname, (libnum, errnum) in _ssl.err_names_to_codes.items():
        lib = error_libraries[_ssl.lib_codes_to_names[libnum]]
        libcode = lib[0]              # e.g. ERR_LIB_PEM
        errcode = lib[1] + errname    # e.g. SSL_R_BAD_SSL_SESSION_ID_LENGTH
        # Only keep it if the numeric codes weren't reused
        if (libcode, errnum) not in new_code_nums:
            codes[errcode] = libcode, errname, errnum
    codes.update(dict(new_codes))

    def w(l):
        f.write(l + "\n")
    w("/* File generated by Tools/ssl/make_ssl_data.py */")
    w("/* Generated on %s */" % datetime.datetime.now().isoformat())
    w("")

    w("static struct py_ssl_library_code library_codes[] = {")
    for mnemo, (libcode, _, _) in sorted(error_libraries.items()):
        w(f'#ifdef {libcode}')
        w('    {"%s", %s},' % (mnemo, libcode))
        w('#endif')
    w('    { NULL }')
    w('};')
    w("")

    w("static struct py_ssl_error_code error_codes[] = {")
    for errcode, (libcode, name, num) in sorted(codes.items()):
        w('  #ifdef %s' % (errcode))
        w('    {"%s", %s, %s},' % (name, libcode, errcode))
        w('  #else')
        w('    {"%s", %s, %d},' % (name, lib_codes[libcode], num))
        w('  #endif')
    w('    { NULL }')
    w('};')
    if not use_stdout:
        f.close()

external/cpython_master/Tools/ssl/make_ssl_data.py
DragonHPC/dragon
# Autodetecting setup.py script for building the Python extensions

import argparse
import importlib._bootstrap
import importlib.machinery
import importlib.util
import os
import re
import sys
import sysconfig
from glob import glob, escape
import _osx_support


try:
    import subprocess
    del subprocess
    SUBPROCESS_BOOTSTRAP = False
except ImportError:
    # Bootstrap Python: distutils.spawn uses subprocess to build C extensions,
    # subprocess requires C extensions built by setup.py like _posixsubprocess.
    #
    # Use _bootsubprocess which only uses the os module.
    #
    # It is dropped from sys.modules as soon as all C extension modules
    # are built.
    import _bootsubprocess
    sys.modules['subprocess'] = _bootsubprocess
    del _bootsubprocess
    SUBPROCESS_BOOTSTRAP = True


from distutils import log
from distutils.command.build_ext import build_ext
from distutils.command.build_scripts import build_scripts
from distutils.command.install import install
from distutils.command.install_lib import install_lib
from distutils.core import Extension, setup
from distutils.errors import CCompilerError, DistutilsError
from distutils.spawn import find_executable


# Compile extensions used to test Python?
TEST_EXTENSIONS = True

# This global variable is used to hold the list of modules to be disabled.
DISABLED_MODULE_LIST = []


def get_platform():
    # Cross compiling
    if "_PYTHON_HOST_PLATFORM" in os.environ:
        return os.environ["_PYTHON_HOST_PLATFORM"]

    # Get value of sys.platform
    if sys.platform.startswith('osf1'):
        return 'osf1'
    return sys.platform


CROSS_COMPILING = ("_PYTHON_HOST_PLATFORM" in os.environ)
HOST_PLATFORM = get_platform()
MS_WINDOWS = (HOST_PLATFORM == 'win32')
CYGWIN = (HOST_PLATFORM == 'cygwin')
MACOS = (HOST_PLATFORM == 'darwin')
AIX = (HOST_PLATFORM.startswith('aix'))
VXWORKS = ('vxworks' in HOST_PLATFORM)


SUMMARY = """
Python is an interpreted, interactive, object-oriented programming
language. It is often compared to Tcl, Perl, Scheme or Java.

Python combines remarkable power with very clear syntax. It has
modules, classes, exceptions, very high level dynamic data types, and
dynamic typing. There are interfaces to many system calls and
libraries, as well as to various windowing systems (X11, Motif, Tk,
Mac, MFC). New built-in modules are easily written in C or C++. Python
is also usable as an extension language for applications that need a
programmable interface.

The Python implementation is portable: it runs on many brands of UNIX,
on Windows, DOS, Mac, Amiga... If your favorite system isn't
listed here, it may still be supported, if there's a C compiler for
it. Ask around on comp.lang.python -- or just try compiling Python
yourself.
"""

CLASSIFIERS = """
Development Status :: 6 - Mature
License :: OSI Approved :: Python Software Foundation License
Natural Language :: English
Programming Language :: C
Programming Language :: Python
Topic :: Software Development
"""


def run_command(cmd):
    status = os.system(cmd)
    return os.waitstatus_to_exitcode(status)


# Set common compiler and linker flags derived from the Makefile,
# reserved for building the interpreter and the stdlib modules.
# See bpo-21121 and bpo-35257
def set_compiler_flags(compiler_flags, compiler_py_flags_nodist):
    flags = sysconfig.get_config_var(compiler_flags)
    py_flags_nodist = sysconfig.get_config_var(compiler_py_flags_nodist)
    sysconfig.get_config_vars()[compiler_flags] = flags + ' ' + py_flags_nodist


def add_dir_to_list(dirlist, dir):
    """Add the directory 'dir' to the list 'dirlist' (after any relative
    directories) if:

    1) 'dir' is not already in 'dirlist'
    2) 'dir' actually exists, and is a directory.
    """
    if dir is None or not os.path.isdir(dir) or dir in dirlist:
        return
    for i, path in enumerate(dirlist):
        if not os.path.isabs(path):
            dirlist.insert(i + 1, dir)
            return
    dirlist.insert(0, dir)


def sysroot_paths(make_vars, subdirs):
    """Get the paths of sysroot sub-directories.

    * make_vars: a sequence of names of variables of the Makefile where
      sysroot may be set.
    * subdirs: a sequence of names of subdirectories used as the location for
      headers or libraries.
    """

    dirs = []
    for var_name in make_vars:
        var = sysconfig.get_config_var(var_name)
        if var is not None:
            m = re.search(r'--sysroot=([^"]\S*|"[^"]+")', var)
            if m is not None:
                sysroot = m.group(1).strip('"')
                for subdir in subdirs:
                    if os.path.isabs(subdir):
                        subdir = subdir[1:]
                    path = os.path.join(sysroot, subdir)
                    if os.path.isdir(path):
                        dirs.append(path)
                break
    return dirs


MACOS_SDK_ROOT = None
MACOS_SDK_SPECIFIED = None

def macosx_sdk_root():
    """Return the directory of the current macOS SDK.

    If no SDK was explicitly configured, call the compiler to find which
    include files paths are being searched by default.  Use '/' if the
    compiler is searching /usr/include (meaning system header files are
    installed) or use the root of an SDK if that is being searched.
    (The SDK may be supplied via Xcode or via the Command Line Tools).
    The SDK paths used by Apple-supplied tool chains depend on the
    setting of various variables; see the xcrun man page for more info.
    Also sets MACOS_SDK_SPECIFIED for use by macosx_sdk_specified().
    """
    global MACOS_SDK_ROOT, MACOS_SDK_SPECIFIED

    # If already called, return cached result.
    if MACOS_SDK_ROOT:
        return MACOS_SDK_ROOT

    cflags = sysconfig.get_config_var('CFLAGS')
    m = re.search(r'-isysroot\s*(\S+)', cflags)
    if m is not None:
        MACOS_SDK_ROOT = m.group(1)
        MACOS_SDK_SPECIFIED = MACOS_SDK_ROOT != '/'
    else:
        MACOS_SDK_ROOT = _osx_support._default_sysroot(
            sysconfig.get_config_var('CC'))
        MACOS_SDK_SPECIFIED = False

    return MACOS_SDK_ROOT


def macosx_sdk_specified():
    """Returns true if an SDK was explicitly configured.

    True if an SDK was selected at configure time, either by specifying
    --enable-universalsdk=(something other than no or /) or by adding a
    -isysroot option to CFLAGS.  In some cases, like when making
    decisions about macOS Tk framework paths, we need to be able to
    know whether the user explicitly asked to build with an SDK versus
    the implicit use of an SDK when header files are no longer
    installed on a running system by the Command Line Tools.
    """
    global MACOS_SDK_SPECIFIED

    # If already called, return cached result.
    if MACOS_SDK_SPECIFIED:
        return MACOS_SDK_SPECIFIED

    # Find the sdk root and set MACOS_SDK_SPECIFIED
    macosx_sdk_root()
    return MACOS_SDK_SPECIFIED


def is_macosx_sdk_path(path):
    """
    Returns True if 'path' can be located in an OSX SDK
    """
    return ( (path.startswith('/usr/') and not path.startswith('/usr/local'))
                or path.startswith('/System/')
                or path.startswith('/Library/') )


def grep_headers_for(function, headers):
    for header in headers:
        with open(header, 'r', errors='surrogateescape') as f:
            if function in f.read():
                return True
    return False

def find_file(filename, std_dirs, paths):
    """Searches for the directory where a given file is located,
    and returns a possibly-empty list of additional directories, or None
    if the file couldn't be found at all.

    'filename' is the name of a file, such as readline.h or libcrypto.a.
    'std_dirs' is the list of standard system directories; if the
        file is found in one of them, no additional directives are needed.
    'paths' is a list of additional locations to check; if the file is
        found in one of them, the resulting list will contain the directory.
    """
    if MACOS:
        # Honor the MacOSX SDK setting when one was specified.
        # An SDK is a directory with the same structure as a real
        # system, but with only header files and libraries.
        sysroot = macosx_sdk_root()

    # Check the standard locations
    for dir in std_dirs:
        f = os.path.join(dir, filename)

        if MACOS and is_macosx_sdk_path(dir):
            f = os.path.join(sysroot, dir[1:], filename)

        if os.path.exists(f): return []

    # Check the additional directories
    for dir in paths:
        f = os.path.join(dir, filename)

        if MACOS and is_macosx_sdk_path(dir):
            f = os.path.join(sysroot, dir[1:], filename)

        if os.path.exists(f):
            return [dir]

    # Not found anywhere
    return None


def find_library_file(compiler, libname, std_dirs, paths):
    result = compiler.find_library_file(std_dirs + paths, libname)
    if result is None:
        return None

    if MACOS:
        sysroot = macosx_sdk_root()

    # Check whether the found file is in one of the standard directories
    dirname = os.path.dirname(result)
    for p in std_dirs:
        # Ensure path doesn't end with path separator
        p = p.rstrip(os.sep)

        if MACOS and is_macosx_sdk_path(p):
            # Note that, as of Xcode 7, Apple SDKs may contain textual stub
            # libraries with .tbd extensions rather than the normal .dylib
            # shared libraries installed in /.  The Apple compiler tool
            # chain handles this transparently but it can cause problems
            # for programs that are being built with an SDK and searching
            # for specific libraries.  Distutils find_library_file() now
            # knows to also search for and return .tbd files.  But callers
            # of find_library_file need to keep in mind that the base filename
            # of the returned SDK library file might have a different extension
            # from that of the library file installed on the running system,
            # for example:
            #   /Applications/Xcode.app/Contents/Developer/Platforms/
            #       MacOSX.platform/Developer/SDKs/MacOSX10.11.sdk/
            #       usr/lib/libedit.tbd
            # vs
            #   /usr/lib/libedit.dylib
            if os.path.join(sysroot, p[1:]) == dirname:
                return [ ]

        if p == dirname:
            return [ ]

    # Otherwise, it must have been in one of the additional directories,
    # so we have to figure out which one.
    for p in paths:
        # Ensure path doesn't end with path separator
        p = p.rstrip(os.sep)

        if MACOS and is_macosx_sdk_path(p):
            if os.path.join(sysroot, p[1:]) == dirname:
                return [ p ]

        if p == dirname:
            return [p]
    else:
        assert False, "Internal error: Path not found in std_dirs or paths"

def validate_tzpath():
    base_tzpath = sysconfig.get_config_var('TZPATH')
    if not base_tzpath:
        return

    tzpaths = base_tzpath.split(os.pathsep)
    bad_paths = [tzpath for tzpath in tzpaths if not os.path.isabs(tzpath)]
    if bad_paths:
        raise ValueError('TZPATH must contain only absolute paths, '
                         + f'found:\n{tzpaths!r}\nwith invalid paths:\n'
                         + f'{bad_paths!r}')

def find_module_file(module, dirlist):
    """Find a module in a set of possible folders. If it is not found
    return the unadorned filename"""
    list = find_file(module, [], dirlist)
    if not list:
        return module
    if len(list) > 1:
        log.info("WARNING: multiple copies of %s found", module)
    return os.path.join(list[0], module)


class PyBuildExt(build_ext):

    def __init__(self, dist):
        build_ext.__init__(self, dist)
        self.srcdir = None
        self.lib_dirs = None
        self.inc_dirs = None
        self.config_h_vars = None
        self.failed = []
        self.failed_on_import = []
        self.missing = []
        self.disabled_configure = []
        if '-j' in os.environ.get('MAKEFLAGS', ''):
            self.parallel = True

    def add(self, ext):
        self.extensions.append(ext)

    def set_srcdir(self):
        self.srcdir = sysconfig.get_config_var('srcdir')
        if not self.srcdir:
            # Maybe running on Windows but not using CYGWIN?
            raise ValueError("No source directory; cannot proceed.")
        self.srcdir = os.path.abspath(self.srcdir)

    def remove_disabled(self):
        # Remove modules that are present on the disabled list
        extensions = [ext for ext in self.extensions
                      if ext.name not in DISABLED_MODULE_LIST]
        # move ctypes to the end, it depends on other modules
        ext_map = dict((ext.name, i) for i, ext in enumerate(extensions))
        if "_ctypes" in ext_map:
            ctypes = extensions.pop(ext_map["_ctypes"])
            extensions.append(ctypes)
        self.extensions = extensions

    def update_sources_depends(self):
        # Fix up the autodetected modules, prefixing all the source files
        # with Modules/.
        moddirlist = [os.path.join(self.srcdir, 'Modules')]

        # Fix up the paths for scripts, too
        self.distribution.scripts = [os.path.join(self.srcdir, filename)
                                     for filename in self.distribution.scripts]

        # Python header files
        headers = [sysconfig.get_config_h_filename()]
        headers += glob(os.path.join(escape(sysconfig.get_path('include')), "*.h"))

        for ext in self.extensions:
            ext.sources = [ find_module_file(filename, moddirlist)
                            for filename in ext.sources ]
            if ext.depends is not None:
                ext.depends = [find_module_file(filename, moddirlist)
                               for filename in ext.depends]
            else:
                ext.depends = []
            # re-compile extensions if a header file has been changed
            ext.depends.extend(headers)

    def remove_configured_extensions(self):
        # The sysconfig variables built by makesetup that list the already
        # built modules and the disabled modules as configured by the Setup
        # files.
        sysconf_built = sysconfig.get_config_var('MODBUILT_NAMES').split()
        sysconf_dis = sysconfig.get_config_var('MODDISABLED_NAMES').split()

        mods_built = []
        mods_disabled = []
        for ext in self.extensions:
            # If a module has already been built or has been disabled in the
            # Setup files, don't build it here.
            if ext.name in sysconf_built:
                mods_built.append(ext)
            if ext.name in sysconf_dis:
                mods_disabled.append(ext)

        mods_configured = mods_built + mods_disabled
        if mods_configured:
            self.extensions = [x for x in self.extensions if x not in
                               mods_configured]
            # Remove the shared libraries built by a previous build.
            for ext in mods_configured:
                fullpath = self.get_ext_fullpath(ext.name)
                if os.path.exists(fullpath):
                    os.unlink(fullpath)

        return (mods_built, mods_disabled)

    def set_compiler_executables(self):
        # When you run "make CC=altcc" or something similar, you really want
        # those environment variables passed into the setup.py phase.  Here's
        # a small set of useful ones.
        compiler = os.environ.get('CC')
        args = {}
        # unfortunately, distutils doesn't let us provide separate C and C++
        # compilers
        if compiler is not None:
            (ccshared,cflags) = sysconfig.get_config_vars('CCSHARED','CFLAGS')
            args['compiler_so'] = compiler + ' ' + ccshared + ' ' + cflags
        self.compiler.set_executables(**args)

    def build_extensions(self):
        self.set_srcdir()

        # Detect which modules should be compiled
        self.detect_modules()

        self.remove_disabled()

        self.update_sources_depends()
        mods_built, mods_disabled = self.remove_configured_extensions()
        self.set_compiler_executables()

        build_ext.build_extensions(self)

        if SUBPROCESS_BOOTSTRAP:
            # Drop our custom subprocess module:
            # use the newly built subprocess module
            del sys.modules['subprocess']

        for ext in self.extensions:
            self.check_extension_import(ext)

        self.summary(mods_built, mods_disabled)

    def summary(self, mods_built, mods_disabled):
        longest = max([len(e.name) for e in self.extensions], default=0)
        if self.failed or self.failed_on_import:
            all_failed = self.failed + self.failed_on_import
            longest = max(longest, max([len(name) for name in all_failed]))

        def print_three_column(lst):
            lst.sort(key=str.lower)
            # guarantee zip() doesn't drop anything
            while len(lst) % 3:
                lst.append("")
            for e, f, g in zip(lst[::3], lst[1::3], lst[2::3]):
                print("%-*s   %-*s   %-*s" % (longest, e, longest, f,
                                              longest, g))

        if self.missing:
            print()
            print("Python build finished successfully!")
            print("The necessary bits to build these optional modules were not "
                  "found:")
            print_three_column(self.missing)
            print("To find the necessary bits, look in setup.py in"
                  " detect_modules() for the module's name.")
            print()

        if mods_built:
            print()
            print("The following modules found by detect_modules() in"
            " setup.py, have been")
            print("built by the Makefile instead, as configured by the"
            " Setup files:")
            print_three_column([ext.name for ext in mods_built])
            print()

        if mods_disabled:
            print()
            print("The following modules found by detect_modules() in"
            " setup.py have not")
            print("been built, they are *disabled* in the Setup files:")
            print_three_column([ext.name for ext in mods_disabled])
            print()

        if self.disabled_configure:
            print()
            print("The following modules found by detect_modules() in"
            " setup.py have not")
            print("been built, they are *disabled* by configure:")
            print_three_column(self.disabled_configure)
            print()

        if self.failed:
            failed = self.failed[:]
            print()
            print("Failed to build these modules:")
            print_three_column(failed)
            print()

        if self.failed_on_import:
            failed = self.failed_on_import[:]
            print()
            print("Following modules built successfully"
                  " but were removed because they could not be imported:")
            print_three_column(failed)
            print()

        if any('_ssl' in l
               for l in (self.missing, self.failed, self.failed_on_import)):
            print()
            print("Could not build the ssl module!")
            print("Python requires an OpenSSL 1.0.2 or 1.1 compatible "
                  "libssl with X509_VERIFY_PARAM_set1_host().")
            print("LibreSSL 2.6.4 and earlier do not provide the necessary "
                  "APIs, https://github.com/libressl-portable/portable/issues/381")
            print()

    def build_extension(self, ext):

        if ext.name == '_ctypes':
            if not self.configure_ctypes(ext):
                self.failed.append(ext.name)
                return

        try:
            build_ext.build_extension(self, ext)
        except (CCompilerError, DistutilsError) as why:
            self.announce('WARNING: building of extension "%s" failed: %s' %
                          (ext.name, why))
            self.failed.append(ext.name)
            return

    def check_extension_import(self, ext):
        # Don't try to import an extension that has failed to compile
        if ext.name in self.failed:
            self.announce(
                'WARNING: skipping import check for failed build "%s"' %
                ext.name, level=1)
            return

        # Workaround for Mac OS X: The Carbon-based modules cannot be
        # reliably imported into a command-line Python
        if 'Carbon' in ext.extra_link_args:
            self.announce(
                'WARNING: skipping import check for Carbon-based "%s"' %
                ext.name)
            return

        if MACOS and (
                sys.maxsize > 2**32 and '-arch' in ext.extra_link_args):
            # Don't bother doing an import check when an extension was
            # build with an explicit '-arch' flag on OSX. That's currently
            # only used to build 32-bit only extensions in a 4-way
            # universal build and loading 32-bit code into a 64-bit
            # process will fail.
            self.announce(
                'WARNING: skipping import check for "%s"' %
                ext.name)
            return

        # Workaround for Cygwin: Cygwin currently has fork issues when many
        # modules have been imported
        if CYGWIN:
            self.announce('WARNING: skipping import check for Cygwin-based "%s"'
                % ext.name)
            return
        ext_filename = os.path.join(
            self.build_lib,
            self.get_ext_filename(self.get_ext_fullname(ext.name)))

        # If the build directory didn't exist when setup.py was
        # started, sys.path_importer_cache has a negative result
        # cached.  Clear that cache before trying to import.
        sys.path_importer_cache.clear()

        # Don't try to load extensions for cross builds
        if CROSS_COMPILING:
            return

        loader = importlib.machinery.ExtensionFileLoader(ext.name, ext_filename)
        spec = importlib.util.spec_from_file_location(ext.name, ext_filename,
                                                      loader=loader)
        try:
            importlib._bootstrap._load(spec)
        except ImportError as why:
            self.failed_on_import.append(ext.name)
            self.announce('*** WARNING: renaming "%s" since importing it'
                          ' failed: %s' % (ext.name, why), level=3)
            assert not self.inplace
            basename, tail = os.path.splitext(ext_filename)
            newname = basename + "_failed" + tail
            if os.path.exists(newname):
                os.remove(newname)
            os.rename(ext_filename, newname)

        except:
            exc_type, why, tb = sys.exc_info()
            self.announce('*** WARNING: importing extension "%s" '
                          'failed with %s: %s' % (ext.name, exc_type, why),
                          level=3)
            self.failed.append(ext.name)

    def add_multiarch_paths(self):
        # Debian/Ubuntu multiarch support.
        # https://wiki.ubuntu.com/MultiarchSpec
        cc = sysconfig.get_config_var('CC')
        tmpfile = os.path.join(self.build_temp, 'multiarch')
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)
        ret = run_command(
            '%s -print-multiarch > %s 2> /dev/null' % (cc, tmpfile))
        multiarch_path_component = ''
        try:
            if ret == 0:
                with open(tmpfile) as fp:
                    multiarch_path_component = fp.readline().strip()
        finally:
            os.unlink(tmpfile)

        if multiarch_path_component != '':
            add_dir_to_list(self.compiler.library_dirs,
                            '/usr/lib/' + multiarch_path_component)
            add_dir_to_list(self.compiler.include_dirs,
                            '/usr/include/' + multiarch_path_component)
            return

        if not find_executable('dpkg-architecture'):
            return
        opt = ''
        if CROSS_COMPILING:
            opt = '-t' + sysconfig.get_config_var('HOST_GNU_TYPE')
        tmpfile = os.path.join(self.build_temp, 'multiarch')
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)
        ret = run_command(
            'dpkg-architecture %s -qDEB_HOST_MULTIARCH > %s 2> /dev/null' %
            (opt, tmpfile))
        try:
            if ret == 0:
                with open(tmpfile) as fp:
                    multiarch_path_component = fp.readline().strip()
                add_dir_to_list(self.compiler.library_dirs,
                                '/usr/lib/' + multiarch_path_component)
                add_dir_to_list(self.compiler.include_dirs,
                                '/usr/include/' + multiarch_path_component)
        finally:
            os.unlink(tmpfile)

    def add_cross_compiling_paths(self):
        cc = sysconfig.get_config_var('CC')
        tmpfile = os.path.join(self.build_temp, 'ccpaths')
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)
        ret = run_command('%s -E -v - </dev/null 2>%s 1>/dev/null' % (cc, tmpfile))
        is_gcc = False
        is_clang = False
        in_incdirs = False
        try:
            if ret == 0:
                with open(tmpfile) as fp:
                    for line in fp.readlines():
                        if line.startswith("gcc version"):
                            is_gcc = True
                        elif line.startswith("clang version"):
                            is_clang = True
                        elif line.startswith("#include <...>"):
                            in_incdirs = True
                        elif line.startswith("End of search list"):
                            in_incdirs = False
                        elif (is_gcc or is_clang) and line.startswith("LIBRARY_PATH"):
                            for d in line.strip().split("=")[1].split(":"):
                                d = os.path.normpath(d)
                                if '/gcc/' not in d:
                                    add_dir_to_list(self.compiler.library_dirs,
                                                    d)
                        elif (is_gcc or is_clang) and in_incdirs and '/gcc/' not in line and '/clang/' not in line:
                            add_dir_to_list(self.compiler.include_dirs,
                                            line.strip())
        finally:
            os.unlink(tmpfile)

    def add_ldflags_cppflags(self):
        # Add paths specified in the environment variables LDFLAGS and
        # CPPFLAGS for header and library files.
        # We must get the values from the Makefile and not the environment
        # directly since an inconsistently reproducible issue comes up where
        # the environment variable is not set even though the value were passed
        # into configure and stored in the Makefile (issue found on OS X 10.3).
        for env_var, arg_name, dir_list in (
                ('LDFLAGS', '-R', self.compiler.runtime_library_dirs),
                ('LDFLAGS', '-L', self.compiler.library_dirs),
                ('CPPFLAGS', '-I', self.compiler.include_dirs)):
            env_val = sysconfig.get_config_var(env_var)
            if env_val:
                parser = argparse.ArgumentParser()
                parser.add_argument(arg_name, dest="dirs", action="append")
                options, _ = parser.parse_known_args(env_val.split())
                if options.dirs:
                    for directory in reversed(options.dirs):
                        add_dir_to_list(dir_list, directory)

    def configure_compiler(self):
        # Ensure that /usr/local is always used, but the local build
        # directories (i.e. '.' and 'Include') must be first.  See issue
        # 10520.
        if not CROSS_COMPILING:
            add_dir_to_list(self.compiler.library_dirs, '/usr/local/lib')
            add_dir_to_list(self.compiler.include_dirs, '/usr/local/include')
        # only change this for cross builds for 3.3, issues on Mageia
        if CROSS_COMPILING:
            self.add_cross_compiling_paths()
        self.add_multiarch_paths()
        self.add_ldflags_cppflags()

    def init_inc_lib_dirs(self):
        if (not CROSS_COMPILING and
                os.path.normpath(sys.base_prefix) != '/usr' and
                not sysconfig.get_config_var('PYTHONFRAMEWORK')):
            # OSX note: Don't add LIBDIR and INCLUDEDIR to building a framework
            # (PYTHONFRAMEWORK is set) to avoid # linking problems when
            # building a framework with different architectures than
            # the one that is currently installed (issue #7473)
            add_dir_to_list(self.compiler.library_dirs,
                            sysconfig.get_config_var("LIBDIR"))
            add_dir_to_list(self.compiler.include_dirs,
                            sysconfig.get_config_var("INCLUDEDIR"))

        system_lib_dirs = ['/lib64', '/usr/lib64', '/lib', '/usr/lib']
        system_include_dirs = ['/usr/include']
        # lib_dirs and inc_dirs are used to search for files;
        # if a file is found in one of those directories, it can
        # be assumed that no additional -I,-L directives are needed.
        if not CROSS_COMPILING:
            self.lib_dirs = self.compiler.library_dirs + system_lib_dirs
            self.inc_dirs = self.compiler.include_dirs + system_include_dirs
        else:
            # Add the sysroot paths. 'sysroot' is a compiler option used to
            # set the logical path of the standard system headers and
            # libraries.
            self.lib_dirs = (self.compiler.library_dirs +
                             sysroot_paths(('LDFLAGS', 'CC'), system_lib_dirs))
            self.inc_dirs = (self.compiler.include_dirs +
                             sysroot_paths(('CPPFLAGS', 'CFLAGS', 'CC'),
                                           system_include_dirs))

        config_h = sysconfig.get_config_h_filename()
        with open(config_h) as file:
            self.config_h_vars = sysconfig.parse_config_h(file)

        # OSF/1 and Unixware have some stuff in /usr/ccs/lib (like -ldb)
        if HOST_PLATFORM in ['osf1', 'unixware7', 'openunix8']:
            self.lib_dirs += ['/usr/ccs/lib']

        # HP-UX11iv3 keeps files in lib/hpux folders.
        if HOST_PLATFORM == 'hp-ux11':
            self.lib_dirs += ['/usr/lib/hpux64', '/usr/lib/hpux32']

        if MACOS:
            # This should work on any unixy platform ;-)
            # If the user has bothered specifying additional -I and -L flags
            # in OPT and LDFLAGS we might as well use them here.
            #
            # NOTE: using shlex.split would technically be more correct, but
            # also gives a bootstrap problem. Let's hope nobody uses
            # directories with whitespace in the name to store libraries.
            cflags, ldflags = sysconfig.get_config_vars(
                    'CFLAGS', 'LDFLAGS')
            for item in cflags.split():
                if item.startswith('-I'):
                    self.inc_dirs.append(item[2:])

            for item in ldflags.split():
                if item.startswith('-L'):
                    self.lib_dirs.append(item[2:])

    def detect_simple_extensions(self):
        #
        # The following modules are all pretty straightforward, and compile
        # on pretty much any POSIXish platform.
        #

        # array objects
        self.add(Extension('array', ['arraymodule.c']))

        # Context Variables
        self.add(Extension('_contextvars', ['_contextvarsmodule.c']))

        shared_math = 'Modules/_math.o'

        # math library functions, e.g. sin()
        self.add(Extension('math',  ['mathmodule.c'],
                           extra_compile_args=['-DPy_BUILD_CORE_MODULE'],
                           extra_objects=[shared_math],
                           depends=['_math.h', shared_math],
                           libraries=['m']))

        # complex math library functions
        self.add(Extension('cmath', ['cmathmodule.c'],
                           extra_compile_args=['-DPy_BUILD_CORE_MODULE'],
                           extra_objects=[shared_math],
                           depends=['_math.h', shared_math],
                           libraries=['m']))

        # time libraries: librt may be needed for clock_gettime()
        time_libs = []
        lib = sysconfig.get_config_var('TIMEMODULE_LIB')
        if lib:
            time_libs.append(lib)

        # time operations and variables
        self.add(Extension('time', ['timemodule.c'],
                           libraries=time_libs))
        # libm is needed by delta_new() that uses round() and by accum() that
        # uses modf().
        self.add(Extension('_datetime', ['_datetimemodule.c'],
                           libraries=['m']))
        # zoneinfo module
        self.add(Extension('_zoneinfo', ['_zoneinfo.c'])),
        # random number generator implemented in C
        self.add(Extension("_random", ["_randommodule.c"],
                           extra_compile_args=['-DPy_BUILD_CORE_MODULE']))
        # bisect
        self.add(Extension("_bisect", ["_bisectmodule.c"]))
        # heapq
        self.add(Extension("_heapq", ["_heapqmodule.c"]))
        # C-optimized pickle replacement
        self.add(Extension("_pickle", ["_pickle.c"],
                           extra_compile_args=['-DPy_BUILD_CORE_MODULE']))
        # atexit
        self.add(Extension("atexit", ["atexitmodule.c"]))
        # _json speedups
        self.add(Extension("_json", ["_json.c"],
                           extra_compile_args=['-DPy_BUILD_CORE_MODULE']))

        # profiler (_lsprof is for cProfile.py)
        self.add(Extension('_lsprof', ['_lsprof.c', 'rotatingtree.c']))
        # static Unicode character database
        self.add(Extension('unicodedata', ['unicodedata.c'],
                           depends=['unicodedata_db.h', 'unicodename_db.h']))
        # _opcode module
        self.add(Extension('_opcode', ['_opcode.c']))
        # asyncio speedups
        self.add(Extension("_asyncio", ["_asynciomodule.c"],
                           extra_compile_args=['-DPy_BUILD_CORE_MODULE']))
        # _abc speedups
        self.add(Extension("_abc", ["_abc.c"]))
        # _queue module
        self.add(Extension("_queue", ["_queuemodule.c"]))
        # _statistics module
        self.add(Extension("_statistics", ["_statisticsmodule.c"]))

        # Modules with some UNIX dependencies -- on by default:
        # (If you have a really backward UNIX, select and socket may not be
        # supported...)

        # fcntl(2) and ioctl(2)
        libs = []
        if (self.config_h_vars.get('FLOCK_NEEDS_LIBBSD', False)):
            # May be necessary on AIX for flock function
            libs = ['bsd']
        self.add(Extension('fcntl', ['fcntlmodule.c'],
                           libraries=libs))
        # pwd(3)
        self.add(Extension('pwd', ['pwdmodule.c']))
        # grp(3)
        if not VXWORKS:
            self.add(Extension('grp', ['grpmodule.c']))
        # spwd, shadow passwords
        if (self.config_h_vars.get('HAVE_GETSPNAM', False) or
                self.config_h_vars.get('HAVE_GETSPENT', False)):
            self.add(Extension('spwd', ['spwdmodule.c']))
        # AIX has shadow passwords, but access is not via getspent(), etc.
        # module support is not expected so it not 'missing'
        elif not AIX:
            self.missing.append('spwd')

        # select(2); not on ancient System V
        self.add(Extension('select', ['selectmodule.c']))

        # Fred Drake's interface to the Python parser
        self.add(Extension('parser', ['parsermodule.c']))

        # Memory-mapped files (also works on Win32).
        self.add(Extension('mmap', ['mmapmodule.c']))

        # Lance Ellinghaus's syslog module
        # syslog daemon interface
        self.add(Extension('syslog', ['syslogmodule.c']))

        # Python interface to subinterpreter C-API.
        self.add(Extension('_xxsubinterpreters', ['_xxsubinterpretersmodule.c']))

        #
        # Here ends the simple stuff.  From here on, modules need certain
        # libraries, are platform-specific, or present other surprises.
        #

        # Multimedia modules
        # These don't work for 64-bit platforms!!!
        # These represent audio samples or images as strings:
        #
        # Operations on audio samples
        # According to #993173, this one should actually work fine on
        # 64-bit platforms.
        #
        # audioop needs libm for floor() in multiple functions.
        self.add(Extension('audioop', ['audioop.c'],
                           libraries=['m']))

        # CSV files
        self.add(Extension('_csv', ['_csv.c']))

        # POSIX subprocess module helper.
        self.add(Extension('_posixsubprocess', ['_posixsubprocess.c']))

    def detect_test_extensions(self):
        # Python C API test module
        self.add(Extension('_testcapi', ['_testcapimodule.c'],
                           depends=['testcapi_long.h']))

        # Python Internal C API test module
        self.add(Extension('_testinternalcapi', ['_testinternalcapi.c'],
                           extra_compile_args=['-DPy_BUILD_CORE_MODULE']))

        # Python PEP-3118 (buffer protocol) test module
        self.add(Extension('_testbuffer', ['_testbuffer.c']))

        # Test loading multiple modules from one compiled file (http://bugs.python.org/issue16421)
        self.add(Extension('_testimportmultiple', ['_testimportmultiple.c']))

        # Test multi-phase extension module init (PEP 489)
        self.add(Extension('_testmultiphase', ['_testmultiphase.c']))

        # Fuzz tests.
        self.add(Extension('_xxtestfuzz',
                           ['_xxtestfuzz/_xxtestfuzz.c',
                            '_xxtestfuzz/fuzzer.c']))

    def detect_readline_curses(self):
        # readline
        do_readline = self.compiler.find_library_file(self.lib_dirs, 'readline')
        readline_termcap_library = ""
        curses_library = ""
        # Cannot use os.popen here in py3k.
        tmpfile = os.path.join(self.build_temp, 'readline_termcap_lib')
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)
        # Determine if readline is already linked against curses or tinfo.
        if do_readline:
            if CROSS_COMPILING:
                ret = run_command("%s -d %s | grep '(NEEDED)' > %s"
                                % (sysconfig.get_config_var('READELF'),
                                   do_readline, tmpfile))
            elif find_executable('ldd'):
                ret = run_command("ldd %s > %s" % (do_readline, tmpfile))
            else:
                ret = 1
            if ret == 0:
                with open(tmpfile) as fp:
                    for ln in fp:
                        if 'curses' in ln:
                            readline_termcap_library = re.sub(
                                r'.*lib(n?cursesw?)\.so.*', r'\1', ln
                            ).rstrip()
                            break
                        # termcap interface split out from ncurses
                        if 'tinfo' in ln:
                            readline_termcap_library = 'tinfo'
                            break
            if os.path.exists(tmpfile):
                os.unlink(tmpfile)
        # Issue 7384: If readline is already linked against curses,
        # use the same library for the readline and curses modules.
        if 'curses' in readline_termcap_library:
            curses_library = readline_termcap_library
        elif self.compiler.find_library_file(self.lib_dirs, 'ncursesw'):
            curses_library = 'ncursesw'
        # Issue 36210: OSS provided ncurses does not link on AIX
        # Use IBM supplied 'curses' for successful build of _curses
        elif AIX and self.compiler.find_library_file(self.lib_dirs, 'curses'):
            curses_library = 'curses'
        elif self.compiler.find_library_file(self.lib_dirs, 'ncurses'):
            curses_library = 'ncurses'
        elif self.compiler.find_library_file(self.lib_dirs, 'curses'):
            curses_library = 'curses'

        if MACOS:
            os_release = int(os.uname()[2].split('.')[0])
            dep_target = sysconfig.get_config_var('MACOSX_DEPLOYMENT_TARGET')
            if (dep_target and
                    (tuple(int(n) for n in dep_target.split('.')[0:2])
                        < (10, 5) ) ):
                os_release = 8
            if os_release < 9:
                # MacOSX 10.4 has a broken readline. Don't try to build
                # the readline module unless the user has installed a fixed
                # readline package
                if find_file('readline/rlconf.h', self.inc_dirs, []) is None:
                    do_readline = False
        if do_readline:
            if MACOS and os_release < 9:
                # In every directory on the search path search for a dynamic
                # library and then a static library, instead of first looking
                # for dynamic libraries on the entire path.
                # This way a statically linked custom readline gets picked up
                # before the (possibly broken) dynamic library in /usr/lib.
                readline_extra_link_args = ('-Wl,-search_paths_first',)
            else:
                readline_extra_link_args = ()

            readline_libs = ['readline']
            if readline_termcap_library:
                pass # Issue 7384: Already linked against curses or tinfo.
            elif curses_library:
                readline_libs.append(curses_library)
            elif self.compiler.find_library_file(self.lib_dirs +
                                                     ['/usr/lib/termcap'],
                                                     'termcap'):
                readline_libs.append('termcap')
            self.add(Extension('readline', ['readline.c'],
                               library_dirs=['/usr/lib/termcap'],
                               extra_link_args=readline_extra_link_args,
                               libraries=readline_libs))
        else:
            self.missing.append('readline')

        # Curses support, requiring the System V version of curses, often
        # provided by the ncurses library.
        curses_defines = []
        curses_includes = []
        panel_library = 'panel'
        if curses_library == 'ncursesw':
            curses_defines.append(('HAVE_NCURSESW', '1'))
            if not CROSS_COMPILING:
                curses_includes.append('/usr/include/ncursesw')
            # Bug 1464056: If _curses.so links with ncursesw,
            # _curses_panel.so must link with panelw.
            panel_library = 'panelw'
            if MACOS:
                # On OS X, there is no separate /usr/lib/libncursesw nor
                # libpanelw.  If we are here, we found a locally-supplied
                # version of libncursesw.  There should also be a
                # libpanelw.  _XOPEN_SOURCE defines are usually excluded
                # for OS X but we need _XOPEN_SOURCE_EXTENDED here for
                # ncurses wide char support
                curses_defines.append(('_XOPEN_SOURCE_EXTENDED', '1'))
        elif MACOS and curses_library == 'ncurses':
            # Building with the system-suppied combined libncurses/libpanel
            curses_defines.append(('HAVE_NCURSESW', '1'))
            curses_defines.append(('_XOPEN_SOURCE_EXTENDED', '1'))

        curses_enabled = True
        if curses_library.startswith('ncurses'):
            curses_libs = [curses_library]
            self.add(Extension('_curses', ['_cursesmodule.c'],
                               include_dirs=curses_includes,
                               define_macros=curses_defines,
                               libraries=curses_libs))
        elif curses_library == 'curses' and not MACOS:
                # OSX has an old Berkeley curses, not good enough for
                # the _curses module.
            if (self.compiler.find_library_file(self.lib_dirs, 'terminfo')):
                curses_libs = ['curses', 'terminfo']
            elif (self.compiler.find_library_file(self.lib_dirs, 'termcap')):
                curses_libs = ['curses', 'termcap']
            else:
                curses_libs = ['curses']

            self.add(Extension('_curses', ['_cursesmodule.c'],
                               define_macros=curses_defines,
                               libraries=curses_libs))
        else:
            curses_enabled = False
            self.missing.append('_curses')

        # If the curses module is enabled, check for the panel module
        # _curses_panel needs some form of ncurses
        skip_curses_panel = True if AIX else False
        if (curses_enabled and not skip_curses_panel and
                self.compiler.find_library_file(self.lib_dirs, panel_library)):
            self.add(Extension('_curses_panel', ['_curses_panel.c'],
                           include_dirs=curses_includes,
                           define_macros=curses_defines,
                           libraries=[panel_library, *curses_libs]))
        elif not skip_curses_panel:
            self.missing.append('_curses_panel')

    def detect_crypt(self):
        # crypt module.
        if VXWORKS:
            # bpo-31904: crypt() function is not provided by VxWorks.
            # DES_crypt() OpenSSL provides is too weak to implement
            # the encryption.
            return

        if self.compiler.find_library_file(self.lib_dirs, 'crypt'):
            libs = ['crypt']
        else:
            libs = []

        self.add(Extension('_crypt', ['_cryptmodule.c'],
                               libraries=libs))

    def detect_socket(self):
        # socket(2)
        if not VXWORKS:
            kwargs = {'depends': ['socketmodule.h']}
            if MACOS:
                # Issue #35569: Expose RFC 3542 socket options.
                kwargs['extra_compile_args'] = ['-D__APPLE_USE_RFC_3542']

            self.add(Extension('_socket', ['socketmodule.c'], **kwargs))
        elif self.compiler.find_library_file(self.lib_dirs, 'net'):
            libs = ['net']
            self.add(Extension('_socket', ['socketmodule.c'],
                               depends=['socketmodule.h'],
                               libraries=libs))

    def detect_dbm_gdbm(self):
        # Modules that provide persistent dictionary-like semantics.  You will
        # probably want to arrange for at least one of them to be available on
        # your machine, though none are defined by default because of library
        # dependencies.  The Python module dbm/__init__.py provides an
        # implementation independent wrapper for these; dbm/dumb.py provides
        # similar functionality (but slower of course) implemented in Python.

        # Sleepycat^WOracle Berkeley DB interface.
        #  http://www.oracle.com/database/berkeley-db/db/index.html
        #
        # This requires the Sleepycat^WOracle DB code. The supported versions
        # are set below.  Visit the URL above to download
        # a release.  Most open source OSes come with one or more
        # versions of BerkeleyDB already installed.

        max_db_ver = (5, 3)
        min_db_ver = (3, 3)
        db_setup_debug = False   # verbose debug prints from this script?

        def allow_db_ver(db_ver):
            """Returns a boolean if the given BerkeleyDB version is acceptable.

            Args:
              db_ver: A tuple of the version to verify.
            """
            if not (min_db_ver <= db_ver <= max_db_ver):
                return False
            return True

        def gen_db_minor_ver_nums(major):
            if major == 4:
                for x in range(max_db_ver[1]+1):
                    if allow_db_ver((4, x)):
                        yield x
            elif major == 3:
                for x in (3,):
                    if allow_db_ver((3, x)):
                        yield x
            else:
                raise ValueError("unknown major BerkeleyDB version", major)

        # construct a list of paths to look for the header file in on
        # top of the normal inc_dirs.
        db_inc_paths = [
            '/usr/include/db4',
            '/usr/local/include/db4',
            '/opt/sfw/include/db4',
            '/usr/include/db3',
            '/usr/local/include/db3',
            '/opt/sfw/include/db3',
            # Fink defaults (http://fink.sourceforge.net/)
            '/sw/include/db4',
            '/sw/include/db3',
        ]
        # 4.x minor number specific paths
        for x in gen_db_minor_ver_nums(4):
            db_inc_paths.append('/usr/include/db4%d' % x)
            db_inc_paths.append('/usr/include/db4.%d' % x)
            db_inc_paths.append('/usr/local/BerkeleyDB.4.%d/include' % x)
            db_inc_paths.append('/usr/local/include/db4%d' % x)
            db_inc_paths.append('/pkg/db-4.%d/include' % x)
            db_inc_paths.append('/opt/db-4.%d/include' % x)
            # MacPorts default (http://www.macports.org/)
            db_inc_paths.append('/opt/local/include/db4%d' % x)
        # 3.x minor number specific paths
        for x in gen_db_minor_ver_nums(3):
            db_inc_paths.append('/usr/include/db3%d' % x)
            db_inc_paths.append('/usr/local/BerkeleyDB.3.%d/include' % x)
            db_inc_paths.append('/usr/local/include/db3%d' % x)
            db_inc_paths.append('/pkg/db-3.%d/include' % x)
            db_inc_paths.append('/opt/db-3.%d/include' % x)

        if CROSS_COMPILING:
            db_inc_paths = []

        # Add some common subdirectories for Sleepycat DB to the list,
        # based on the standard include directories. This way DB3/4 gets
        # picked up when it is installed in a non-standard prefix and
        # the user has added that prefix into inc_dirs.
        std_variants = []
        for dn in self.inc_dirs:
            std_variants.append(os.path.join(dn, 'db3'))
            std_variants.append(os.path.join(dn, 'db4'))
            for x in gen_db_minor_ver_nums(4):
                std_variants.append(os.path.join(dn, "db4%d"%x))
                std_variants.append(os.path.join(dn, "db4.%d"%x))
            for x in gen_db_minor_ver_nums(3):
                std_variants.append(os.path.join(dn, "db3%d"%x))
                std_variants.append(os.path.join(dn, "db3.%d"%x))

        db_inc_paths = std_variants + db_inc_paths
        db_inc_paths = [p for p in db_inc_paths if os.path.exists(p)]

        db_ver_inc_map = {}

        if MACOS:
            sysroot = macosx_sdk_root()

        class db_found(Exception): pass
        try:
            # See whether there is a Sleepycat header in the standard
            # search path.
            for d in self.inc_dirs + db_inc_paths:
                f = os.path.join(d, "db.h")
                if MACOS and is_macosx_sdk_path(d):
                    f = os.path.join(sysroot, d[1:], "db.h")

                if db_setup_debug: print("db: looking for db.h in", f)
                if os.path.exists(f):
                    with open(f, 'rb') as file:
                        f = file.read()
                    m = re.search(br"#define\WDB_VERSION_MAJOR\W(\d+)", f)
                    if m:
                        db_major = int(m.group(1))
                        m = re.search(br"#define\WDB_VERSION_MINOR\W(\d+)", f)
                        db_minor = int(m.group(1))
                        db_ver = (db_major, db_minor)

                        # Avoid 4.6 prior to 4.6.21 due to a BerkeleyDB bug
                        if db_ver == (4, 6):
                            m = re.search(br"#define\WDB_VERSION_PATCH\W(\d+)", f)
                            db_patch = int(m.group(1))
                            if db_patch < 21:
                                print("db.h:", db_ver, "patch", db_patch,
                                      "being ignored (4.6.x must be >= 4.6.21)")
                                continue

                        if ( (db_ver not in db_ver_inc_map) and
                            allow_db_ver(db_ver) ):
                            # save the include directory with the db.h version
                            # (first occurrence only)
                            db_ver_inc_map[db_ver] = d
                            if db_setup_debug:
                                print("db.h: found", db_ver, "in", d)
                        else:
                            # we already found a header for this library version
                            if db_setup_debug: print("db.h: ignoring", d)
                    else:
                        # ignore this header, it didn't contain a version number
                        if db_setup_debug:
                            print("db.h: no version number version in", d)

            db_found_vers = list(db_ver_inc_map.keys())
            db_found_vers.sort()

            while db_found_vers:
                db_ver = db_found_vers.pop()
                db_incdir = db_ver_inc_map[db_ver]

                # check lib directories parallel to the location of the header
                db_dirs_to_check = [
                    db_incdir.replace("include", 'lib64'),
                    db_incdir.replace("include", 'lib'),
                ]

                if not MACOS:
                    db_dirs_to_check = list(filter(os.path.isdir, db_dirs_to_check))

                else:
                    # Same as other branch, but takes OSX SDK into account
                    tmp = []
                    for dn in db_dirs_to_check:
                        if is_macosx_sdk_path(dn):
                            if os.path.isdir(os.path.join(sysroot, dn[1:])):
                                tmp.append(dn)
                        else:
                            if os.path.isdir(dn):
                                tmp.append(dn)
                    db_dirs_to_check = tmp

                    db_dirs_to_check = tmp

                # Look for a version specific db-X.Y before an ambiguous dbX
                # XXX should we -ever- look for a dbX name?  Do any
                # systems really not name their library by version and
                # symlink to more general names?
                for dblib in (('db-%d.%d' % db_ver),
                              ('db%d%d' % db_ver),
                              ('db%d' % db_ver[0])):
                    dblib_file = self.compiler.find_library_file(
                                    db_dirs_to_check + self.lib_dirs, dblib )
                    if dblib_file:
                        dblib_dir = [ os.path.abspath(os.path.dirname(dblib_file)) ]
                        raise db_found
                    else:
                        if db_setup_debug: print("db lib: ", dblib, "not found")

        except db_found:
            if db_setup_debug:
                print("bsddb using BerkeleyDB lib:", db_ver, dblib)
                print("bsddb lib dir:", dblib_dir, " inc dir:", db_incdir)
            dblibs = [dblib]
            # Only add the found library and include directories if they aren't
            # already being searched. This avoids an explicit runtime library
            # dependency.
            if db_incdir in self.inc_dirs:
                db_incs = None
            else:
                db_incs = [db_incdir]
            if dblib_dir[0] in self.lib_dirs:
                dblib_dir = None
        else:
            if db_setup_debug: print("db: no appropriate library found")
            db_incs = None
            dblibs = []
            dblib_dir = None

        dbm_setup_debug = False   # verbose debug prints from this script?
        dbm_order = ['gdbm']
        # The standard Unix dbm module:
        if not CYGWIN:
            config_args = [arg.strip("'")
                           for arg in sysconfig.get_config_var("CONFIG_ARGS").split()]
            dbm_args = [arg for arg in config_args
                        if arg.startswith('--with-dbmliborder=')]
            if dbm_args:
                dbm_order = [arg.split('=')[-1] for arg in dbm_args][-1].split(":")
            else:
                dbm_order = "ndbm:gdbm:bdb".split(":")
            dbmext = None
            for cand in dbm_order:
                if cand == "ndbm":
                    if find_file("ndbm.h", self.inc_dirs, []) is not None:
                        # Some systems have -lndbm, others have -lgdbm_compat,
                        # others don't have either
                        if self.compiler.find_library_file(self.lib_dirs,
                                                               'ndbm'):
                            ndbm_libs = ['ndbm']
                        elif self.compiler.find_library_file(self.lib_dirs,
                                                             'gdbm_compat'):
                            ndbm_libs = ['gdbm_compat']
                        else:
                            ndbm_libs = []
                        if dbm_setup_debug: print("building dbm using ndbm")
                        dbmext = Extension('_dbm', ['_dbmmodule.c'],
                                           define_macros=[
                                               ('HAVE_NDBM_H',None),
                                               ],
                                           libraries=ndbm_libs)
                        break

                elif cand == "gdbm":
                    if self.compiler.find_library_file(self.lib_dirs, 'gdbm'):
                        gdbm_libs = ['gdbm']
                        if self.compiler.find_library_file(self.lib_dirs,
                                                               'gdbm_compat'):
                            gdbm_libs.append('gdbm_compat')
                        if find_file("gdbm/ndbm.h", self.inc_dirs, []) is not None:
                            if dbm_setup_debug: print("building dbm using gdbm")
                            dbmext = Extension(
                                '_dbm', ['_dbmmodule.c'],
                                define_macros=[
                                    ('HAVE_GDBM_NDBM_H', None),
                                    ],
                                libraries = gdbm_libs)
                            break
                        if find_file("gdbm-ndbm.h", self.inc_dirs, []) is not None:
                            if dbm_setup_debug: print("building dbm using gdbm")
                            dbmext = Extension(
                                '_dbm', ['_dbmmodule.c'],
                                define_macros=[
                                    ('HAVE_GDBM_DASH_NDBM_H', None),
                                    ],
                                libraries = gdbm_libs)
                            break
                elif cand == "bdb":
                    if dblibs:
                        if dbm_setup_debug: print("building dbm using bdb")
                        dbmext = Extension('_dbm', ['_dbmmodule.c'],
                                           library_dirs=dblib_dir,
                                           runtime_library_dirs=dblib_dir,
                                           include_dirs=db_incs,
                                           define_macros=[
                                               ('HAVE_BERKDB_H', None),
                                               ('DB_DBM_HSEARCH', None),
                                               ],
                                           libraries=dblibs)
                        break
            if dbmext is not None:
                self.add(dbmext)
            else:
                self.missing.append('_dbm')

        # Anthony Baxter's gdbm module.  GNU dbm(3) will require -lgdbm:
        if ('gdbm' in dbm_order and
            self.compiler.find_library_file(self.lib_dirs, 'gdbm')):
            self.add(Extension('_gdbm', ['_gdbmmodule.c'],
                               libraries=['gdbm']))
        else:
            self.missing.append('_gdbm')

    def detect_sqlite(self):
        # The sqlite interface
        sqlite_setup_debug = False   # verbose debug prints from this script?

        # We hunt for #define SQLITE_VERSION "n.n.n"
        # We need to find >= sqlite version 3.3.9, for sqlite3_prepare_v2
        sqlite_incdir = sqlite_libdir = None
        sqlite_inc_paths = [ '/usr/include',
                             '/usr/include/sqlite',
                             '/usr/include/sqlite3',
                             '/usr/local/include',
                             '/usr/local/include/sqlite',
                             '/usr/local/include/sqlite3',
                             ]
        if CROSS_COMPILING:
            sqlite_inc_paths = []
        MIN_SQLITE_VERSION_NUMBER = (3, 7, 2)
        MIN_SQLITE_VERSION = ".".join([str(x)
                                    for x in MIN_SQLITE_VERSION_NUMBER])

        # Scan the default include directories before the SQLite specific
        # ones. This allows one to override the copy of sqlite on OSX,
        # where /usr/include contains an old version of sqlite.
        if MACOS:
            sysroot = macosx_sdk_root()

        for d_ in self.inc_dirs + sqlite_inc_paths:
            d = d_
            if MACOS and is_macosx_sdk_path(d):
                d = os.path.join(sysroot, d[1:])

            f = os.path.join(d, "sqlite3.h")
            if os.path.exists(f):
                if sqlite_setup_debug: print("sqlite: found %s"%f)
                with open(f) as file:
                    incf = file.read()
                m = re.search(
                    r'\s*.*#\s*.*define\s.*SQLITE_VERSION\W*"([\d\.]*)"', incf)
                if m:
                    sqlite_version = m.group(1)
                    sqlite_version_tuple = tuple([int(x)
                                        for x in sqlite_version.split(".")])
                    if sqlite_version_tuple >= MIN_SQLITE_VERSION_NUMBER:
                        # we win!
                        if sqlite_setup_debug:
                            print("%s/sqlite3.h: version %s"%(d, sqlite_version))
                        sqlite_incdir = d
                        break
                    else:
                        if sqlite_setup_debug:
                            print("%s: version %s is too old, need >= %s"%(d,
                                        sqlite_version, MIN_SQLITE_VERSION))
                elif sqlite_setup_debug:
                    print("sqlite: %s had no SQLITE_VERSION"%(f,))

        if sqlite_incdir:
            sqlite_dirs_to_check = [
                os.path.join(sqlite_incdir, '..', 'lib64'),
                os.path.join(sqlite_incdir, '..', 'lib'),
                os.path.join(sqlite_incdir, '..', '..', 'lib64'),
                os.path.join(sqlite_incdir, '..', '..', 'lib'),
            ]
            sqlite_libfile = self.compiler.find_library_file(
                                sqlite_dirs_to_check + self.lib_dirs, 'sqlite3')
            if sqlite_libfile:
                sqlite_libdir = [os.path.abspath(os.path.dirname(sqlite_libfile))]

        if sqlite_incdir and sqlite_libdir:
            sqlite_srcs = ['_sqlite/cache.c',
                '_sqlite/connection.c',
                '_sqlite/cursor.c',
                '_sqlite/microprotocols.c',
                '_sqlite/module.c',
                '_sqlite/prepare_protocol.c',
                '_sqlite/row.c',
                '_sqlite/statement.c',
                '_sqlite/util.c', ]

            sqlite_defines = []
            if not MS_WINDOWS:
                sqlite_defines.append(('MODULE_NAME', '"sqlite3"'))
            else:
                sqlite_defines.append(('MODULE_NAME', '\\"sqlite3\\"'))

            # Enable support for loadable extensions in the sqlite3 module
            # if --enable-loadable-sqlite-extensions configure option is used.
            if '--enable-loadable-sqlite-extensions' not in sysconfig.get_config_var("CONFIG_ARGS"):
                sqlite_defines.append(("SQLITE_OMIT_LOAD_EXTENSION", "1"))

            if MACOS:
                # In every directory on the search path search for a dynamic
                # library and then a static library, instead of first looking
                # for dynamic libraries on the entire path.
                # This way a statically linked custom sqlite gets picked up
                # before the dynamic library in /usr/lib.
                sqlite_extra_link_args = ('-Wl,-search_paths_first',)
            else:
                sqlite_extra_link_args = ()

            include_dirs = ["Modules/_sqlite"]
            # Only include the directory where sqlite was found if it does
            # not already exist in set include directories, otherwise you
            # can end up with a bad search path order.
            if sqlite_incdir not in self.compiler.include_dirs:
                include_dirs.append(sqlite_incdir)
            # avoid a runtime library path for a system library dir
            if sqlite_libdir and sqlite_libdir[0] in self.lib_dirs:
                sqlite_libdir = None
            self.add(Extension('_sqlite3', sqlite_srcs,
                               define_macros=sqlite_defines,
                               include_dirs=include_dirs,
                               library_dirs=sqlite_libdir,
                               extra_link_args=sqlite_extra_link_args,
                               libraries=["sqlite3",]))
        else:
            self.missing.append('_sqlite3')

    def detect_platform_specific_exts(self):
        # Unix-only modules
        if not MS_WINDOWS:
            if not VXWORKS:
                # Steen Lumholt's termios module
                self.add(Extension('termios', ['termios.c']))
                # Jeremy Hylton's rlimit interface
            self.add(Extension('resource', ['resource.c']))
        else:
            self.missing.extend(['resource', 'termios'])

        # Platform-specific libraries
        if HOST_PLATFORM.startswith(('linux', 'freebsd', 'gnukfreebsd')):
            self.add(Extension('ossaudiodev', ['ossaudiodev.c']))
        elif not AIX:
            self.missing.append('ossaudiodev')

        if MACOS:
            self.add(Extension('_scproxy', ['_scproxy.c'],
                               extra_link_args=[
                                   '-framework', 'SystemConfiguration',
                                   '-framework', 'CoreFoundation']))

    def detect_compress_exts(self):
        # Andrew Kuchling's zlib module.  Note that some versions of zlib
        # 1.1.3 have security problems.  See CERT Advisory CA-2002-07:
        # http://www.cert.org/advisories/CA-2002-07.html
        #
        # zlib 1.1.4 is fixed, but at least one vendor (RedHat) has decided to
        # patch its zlib 1.1.3 package instead of upgrading to 1.1.4.  For
        # now, we still accept 1.1.3, because we think it's difficult to
        # exploit this in Python, and we'd rather make it RedHat's problem
        # than our problem <wink>.
        #
        # You can upgrade zlib to version 1.1.4 yourself by going to
        # http://www.gzip.org/zlib/
        zlib_inc = find_file('zlib.h', [], self.inc_dirs)
        have_zlib = False
        if zlib_inc is not None:
            zlib_h = zlib_inc[0] + '/zlib.h'
            version = '"0.0.0"'
            version_req = '"1.1.3"'
            if MACOS and is_macosx_sdk_path(zlib_h):
                zlib_h = os.path.join(macosx_sdk_root(), zlib_h[1:])
            with open(zlib_h) as fp:
                while 1:
                    line = fp.readline()
                    if not line:
                        break
                    if line.startswith('#define ZLIB_VERSION'):
                        version = line.split()[2]
                        break
            if version >= version_req:
                if (self.compiler.find_library_file(self.lib_dirs, 'z')):
                    if MACOS:
                        zlib_extra_link_args = ('-Wl,-search_paths_first',)
                    else:
                        zlib_extra_link_args = ()
                    self.add(Extension('zlib', ['zlibmodule.c'],
                                       libraries=['z'],
                                       extra_link_args=zlib_extra_link_args))
                    have_zlib = True
                else:
                    self.missing.append('zlib')
            else:
                self.missing.append('zlib')
        else:
            self.missing.append('zlib')

        # Helper module for various ascii-encoders.  Uses zlib for an optimized
        # crc32 if we have it.  Otherwise binascii uses its own.
        if have_zlib:
            extra_compile_args = ['-DUSE_ZLIB_CRC32']
            libraries = ['z']
            extra_link_args = zlib_extra_link_args
        else:
            extra_compile_args = []
            libraries = []
            extra_link_args = []
        self.add(Extension('binascii', ['binascii.c'],
                           extra_compile_args=extra_compile_args,
                           libraries=libraries,
                           extra_link_args=extra_link_args))

        # Gustavo Niemeyer's bz2 module.
        if (self.compiler.find_library_file(self.lib_dirs, 'bz2')):
            if MACOS:
                bz2_extra_link_args = ('-Wl,-search_paths_first',)
            else:
                bz2_extra_link_args = ()
            self.add(Extension('_bz2', ['_bz2module.c'],
                               libraries=['bz2'],
                               extra_link_args=bz2_extra_link_args))
        else:
            self.missing.append('_bz2')

        # LZMA compression support.
        if self.compiler.find_library_file(self.lib_dirs, 'lzma'):
            self.add(Extension('_lzma', ['_lzmamodule.c'],
                               libraries=['lzma']))
        else:
            self.missing.append('_lzma')

    def detect_expat_elementtree(self):
        # Interface to the Expat XML parser
        #
        # Expat was written by James Clark and is now maintained by a group of
        # developers on SourceForge; see www.libexpat.org for more information.
        # The pyexpat module was written by Paul Prescod after a prototype by
        # Jack Jansen.  The Expat source is included in Modules/expat/.  Usage
        # of a system shared libexpat.so is possible with --with-system-expat
        # configure option.
        #
        # More information on Expat can be found at www.libexpat.org.
        #
        if '--with-system-expat' in sysconfig.get_config_var("CONFIG_ARGS"):
            expat_inc = []
            define_macros = []
            extra_compile_args = []
            expat_lib = ['expat']
            expat_sources = []
            expat_depends = []
        else:
            expat_inc = [os.path.join(self.srcdir, 'Modules', 'expat')]
            define_macros = [
                ('HAVE_EXPAT_CONFIG_H', '1'),
                # bpo-30947: Python uses best available entropy sources to
                # call XML_SetHashSalt(), expat entropy sources are not needed
                ('XML_POOR_ENTROPY', '1'),
            ]
            extra_compile_args = []
            expat_lib = []
            expat_sources = ['expat/xmlparse.c',
                             'expat/xmlrole.c',
                             'expat/xmltok.c']
            expat_depends = ['expat/ascii.h',
                             'expat/asciitab.h',
                             'expat/expat.h',
                             'expat/expat_config.h',
                             'expat/expat_external.h',
                             'expat/internal.h',
                             'expat/latin1tab.h',
                             'expat/utf8tab.h',
                             'expat/xmlrole.h',
                             'expat/xmltok.h',
                             'expat/xmltok_impl.h'
                             ]

            cc = sysconfig.get_config_var('CC').split()[0]
            ret = run_command(
                      '"%s" -Werror -Wno-unreachable-code -E -xc /dev/null >/dev/null 2>&1' % cc)
            if ret == 0:
                extra_compile_args.append('-Wno-unreachable-code')

        self.add(Extension('pyexpat',
                           define_macros=define_macros,
                           extra_compile_args=extra_compile_args,
                           include_dirs=expat_inc,
                           libraries=expat_lib,
                           sources=['pyexpat.c'] + expat_sources,
                           depends=expat_depends))

        # Fredrik Lundh's cElementTree module.  Note that this also
        # uses expat (via the CAPI hook in pyexpat).

        if os.path.isfile(os.path.join(self.srcdir, 'Modules', '_elementtree.c')):
            define_macros.append(('USE_PYEXPAT_CAPI', None))
            self.add(Extension('_elementtree',
                               define_macros=define_macros,
                               include_dirs=expat_inc,
                               libraries=expat_lib,
                               sources=['_elementtree.c'],
                               depends=['pyexpat.c', *expat_sources,
                                        *expat_depends]))
        else:
            self.missing.append('_elementtree')

    def detect_multibytecodecs(self):
        # Hye-Shik Chang's CJKCodecs modules.
        self.add(Extension('_multibytecodec',
                           ['cjkcodecs/multibytecodec.c']))
        for loc in ('kr', 'jp', 'cn', 'tw', 'hk', 'iso2022'):
            self.add(Extension('_codecs_%s' % loc,
                               ['cjkcodecs/_codecs_%s.c' % loc]))

    def detect_multiprocessing(self):
        # Richard Oudkerk's multiprocessing module
        if MS_WINDOWS:
            multiprocessing_srcs = ['_multiprocessing/multiprocessing.c',
                                    '_multiprocessing/semaphore.c']

        else:
            multiprocessing_srcs = ['_multiprocessing/multiprocessing.c']
            if (sysconfig.get_config_var('HAVE_SEM_OPEN') and not
                sysconfig.get_config_var('POSIX_SEMAPHORES_NOT_ENABLED')):
                multiprocessing_srcs.append('_multiprocessing/semaphore.c')
            if (sysconfig.get_config_var('HAVE_SHM_OPEN') and
                sysconfig.get_config_var('HAVE_SHM_UNLINK')):
                posixshmem_srcs = ['_multiprocessing/posixshmem.c']
                libs = []
                if sysconfig.get_config_var('SHM_NEEDS_LIBRT'):
                    # need to link with librt to get shm_open()
                    libs.append('rt')
                self.add(Extension('_posixshmem', posixshmem_srcs,
                                   define_macros={},
                                   libraries=libs,
                                   include_dirs=["Modules/_multiprocessing"]))

        self.add(Extension('_multiprocessing', multiprocessing_srcs,
                           include_dirs=["Modules/_multiprocessing"]))

    def detect_uuid(self):
        # Build the _uuid module if possible
        uuid_incs = find_file("uuid.h", self.inc_dirs, ["/usr/include/uuid"])
        if uuid_incs is not None:
            if self.compiler.find_library_file(self.lib_dirs, 'uuid'):
                uuid_libs = ['uuid']
            else:
                uuid_libs = []
            self.add(Extension('_uuid', ['_uuidmodule.c'],
                               libraries=uuid_libs,
                               include_dirs=uuid_incs))
        else:
            self.missing.append('_uuid')

    def detect_modules(self):
        self.configure_compiler()
        self.init_inc_lib_dirs()

        self.detect_simple_extensions()
        if TEST_EXTENSIONS:
            self.detect_test_extensions()
        self.detect_readline_curses()
        self.detect_crypt()
        self.detect_socket()
        self.detect_openssl_hashlib()
        self.detect_hash_builtins()
        self.detect_dbm_gdbm()
        self.detect_sqlite()
        self.detect_platform_specific_exts()
        self.detect_nis()
        self.detect_compress_exts()
        self.detect_expat_elementtree()
        self.detect_multibytecodecs()
        self.detect_decimal()
        self.detect_ctypes()
        self.detect_multiprocessing()
        if not self.detect_tkinter():
            self.missing.append('_tkinter')
        self.detect_uuid()

##         # Uncomment these lines if you want to play with xxmodule.c
##         self.add(Extension('xx', ['xxmodule.c']))

        if 'd' not in sysconfig.get_config_var('ABIFLAGS'):
            self.add(Extension('xxlimited', ['xxlimited.c'],
                               define_macros=[('Py_LIMITED_API', '0x03050000')]))

    def detect_tkinter_explicitly(self):
        # Build _tkinter using explicit locations for Tcl/Tk.
        #
        # This is enabled when both arguments are given to ./configure:
        #
        #     --with-tcltk-includes="-I/path/to/tclincludes \
        #                            -I/path/to/tkincludes"
        #     --with-tcltk-libs="-L/path/to/tcllibs -ltclm.n \
        #                        -L/path/to/tklibs -ltkm.n"
        #
        # These values can also be specified or overridden via make:
        #    make TCLTK_INCLUDES="..." TCLTK_LIBS="..."
        #
        # This can be useful for building and testing tkinter with multiple
        # versions of Tcl/Tk.  Note that a build of Tk depends on a particular
        # build of Tcl so you need to specify both arguments and use care when
        # overriding.

        # The _TCLTK variables are created in the Makefile sharedmods target.
        tcltk_includes = os.environ.get('_TCLTK_INCLUDES')
        tcltk_libs = os.environ.get('_TCLTK_LIBS')
        if not (tcltk_includes and tcltk_libs):
            # Resume default configuration search.
            return False

        extra_compile_args = tcltk_includes.split()
        extra_link_args = tcltk_libs.split()
        self.add(Extension('_tkinter', ['_tkinter.c', 'tkappinit.c'],
                           define_macros=[('WITH_APPINIT', 1)],
                           extra_compile_args = extra_compile_args,
                           extra_link_args = extra_link_args))
        return True

    def detect_tkinter_darwin(self):
        # Build default _tkinter on macOS using Tcl and Tk frameworks.
        #
        # The macOS native Tk (AKA Aqua Tk) and Tcl are most commonly
        # built and installed as macOS framework bundles.  However,
        # for several reasons, we cannot take full advantage of the
        # Apple-supplied compiler chain's -framework options here.
        # Instead, we need to find and pass to the compiler the
        # absolute paths of the Tcl and Tk headers files we want to use
        # and the absolute path to the directory containing the Tcl
        # and Tk frameworks for linking.
        #
        # We want to handle here two common use cases on macOS:
        # 1. Build and link with system-wide third-party or user-built
        #    Tcl and Tk frameworks installed in /Library/Frameworks.
        # 2. Build and link using a user-specified macOS SDK so that the
        #    built Python can be exported to other systems.  In this case,
        #    search only the SDK's /Library/Frameworks (normally empty)
        #    and /System/Library/Frameworks.
        #
        # Any other use case should be able to be handled explicitly by
        # using the options described above in detect_tkinter_explicitly().
        # In particular it would be good to handle here the case where
        # you want to build and link with a framework build of Tcl and Tk
        # that is not in /Library/Frameworks, say, in your private
        # $HOME/Library/Frameworks directory or elsewhere. It turns
        # out to be difficult to make that work automatically here
        # without bringing into play more tools and magic. That case
        # can be handled using a recipe with the right arguments
        # to detect_tkinter_explicitly().
        #
        # Note also that the fallback case here is to try to use the
        # Apple-supplied Tcl and Tk frameworks in /System/Library but
        # be forewarned that they are deprecated by Apple and typically
        # out-of-date and buggy; their use should be avoided if at
        # all possible by installing a newer version of Tcl and Tk in
        # /Library/Frameworks before building Python without
        # an explicit SDK or by configuring build arguments explicitly.

        from os.path import join, exists

        sysroot = macosx_sdk_root() # path to the SDK or '/'

        if macosx_sdk_specified():
            # Use case #2: an SDK other than '/' was specified.
            # Only search there.
            framework_dirs = [
                join(sysroot, 'Library', 'Frameworks'),
                join(sysroot, 'System', 'Library', 'Frameworks'),
            ]
        else:
            # Use case #1: no explicit SDK selected.
            # Search the local system-wide /Library/Frameworks,
            # not the one in the default SDK, otherwise fall back to
            # /System/Library/Frameworks whose header files may be in
            # the default SDK or, on older systems, actually installed.
            framework_dirs = [
                join('/', 'Library', 'Frameworks'),
                join(sysroot, 'System', 'Library', 'Frameworks'),
            ]

        # Find the directory that contains the Tcl.framework and
        # Tk.framework bundles.
        for F in framework_dirs:
            # both Tcl.framework and Tk.framework should be present
            for fw in 'Tcl', 'Tk':
                if not exists(join(F, fw + '.framework')):
                    break
            else:
                # ok, F is now directory with both frameworks. Continue
                # building
                break
        else:
            # Tk and Tcl frameworks not found. Normal "unix" tkinter search
            # will now resume.
            return False

        include_dirs = [
            join(F, fw + '.framework', H)
            for fw in ('Tcl', 'Tk')
            for H in ('Headers',)
        ]

        # Add the base framework directory as well
        compile_args = ['-F', F]

        # Do not build tkinter for archs that this Tk was not built with.
        cflags = sysconfig.get_config_vars('CFLAGS')[0]
        archs = re.findall(r'-arch\s+(\w+)', cflags)

        tmpfile = os.path.join(self.build_temp, 'tk.arch')
        if not os.path.exists(self.build_temp):
            os.makedirs(self.build_temp)

        run_command(
            "file {}/Tk.framework/Tk | grep 'for architecture' > {}".format(F, tmpfile)
        )
        with open(tmpfile) as fp:
            detected_archs = []
            for ln in fp:
                a = ln.split()[-1]
                if a in archs:
                    detected_archs.append(ln.split()[-1])
        os.unlink(tmpfile)

        arch_args = []
        for a in detected_archs:
            arch_args.append('-arch')
            arch_args.append(a)

        compile_args += arch_args
        link_args = [','.join(['-Wl', '-F', F, '-framework', 'Tcl', '-framework', 'Tk']), *arch_args]

        # The X11/xlib.h file bundled in the Tk sources can cause function
        # prototype warnings from the compiler. Since we cannot easily fix
        # that, suppress the warnings here instead.
        if '-Wstrict-prototypes' in cflags.split():
            compile_args.append('-Wno-strict-prototypes')

        self.add(Extension('_tkinter', ['_tkinter.c', 'tkappinit.c'],
                           define_macros=[('WITH_APPINIT', 1)],
                           include_dirs=include_dirs,
                           libraries=[],
                           extra_compile_args=compile_args,
                           extra_link_args=link_args))
        return True

    def detect_tkinter(self):
        # The _tkinter module.

        # Check whether --with-tcltk-includes and --with-tcltk-libs were
        # configured or passed into the make target.  If so, use these values
        # to build tkinter and bypass the searches for Tcl and TK in standard
        # locations.
        if self.detect_tkinter_explicitly():
            return True

        # Rather than complicate the code below, detecting and building
        # AquaTk is a separate method. Only one Tkinter will be built on
        # Darwin - either AquaTk, if it is found, or X11 based Tk.
        if (MACOS and self.detect_tkinter_darwin()):
            return True

        # Assume we haven't found any of the libraries or include files
        # The versions with dots are used on Unix, and the versions without
        # dots on Windows, for detection by cygwin.
        tcllib = tklib = tcl_includes = tk_includes = None
        for version in ['8.6', '86', '8.5', '85', '8.4', '84', '8.3', '83',
                        '8.2', '82', '8.1', '81', '8.0', '80']:
            tklib = self.compiler.find_library_file(self.lib_dirs,
                                                        'tk' + version)
            tcllib = self.compiler.find_library_file(self.lib_dirs,
                                                         'tcl' + version)
            if tklib and tcllib:
                # Exit the loop when we've found the Tcl/Tk libraries
                break

        # Now check for the header files
        if tklib and tcllib:
            # Check for the include files on Debian and {Free,Open}BSD, where
            # they're put in /usr/include/{tcl,tk}X.Y
            dotversion = version
            if '.' not in dotversion and "bsd" in HOST_PLATFORM.lower():
                # OpenBSD and FreeBSD use Tcl/Tk library names like libtcl83.a,
                # but the include subdirs are named like .../include/tcl8.3.
                dotversion = dotversion[:-1] + '.' + dotversion[-1]
            tcl_include_sub = []
            tk_include_sub = []
            for dir in self.inc_dirs:
                tcl_include_sub += [dir + os.sep + "tcl" + dotversion]
                tk_include_sub += [dir + os.sep + "tk" + dotversion]
            tk_include_sub += tcl_include_sub
            tcl_includes = find_file('tcl.h', self.inc_dirs, tcl_include_sub)
            tk_includes = find_file('tk.h', self.inc_dirs, tk_include_sub)

        if (tcllib is None or tklib is None or
            tcl_includes is None or tk_includes is None):
            self.announce("INFO: Can't locate Tcl/Tk libs and/or headers", 2)
            return False

        # OK... everything seems to be present for Tcl/Tk.

        include_dirs = []
        libs = []
        defs = []
        added_lib_dirs = []
        for dir in tcl_includes + tk_includes:
            if dir not in include_dirs:
                include_dirs.append(dir)

        # Check for various platform-specific directories
        if HOST_PLATFORM == 'sunos5':
            include_dirs.append('/usr/openwin/include')
            added_lib_dirs.append('/usr/openwin/lib')
        elif os.path.exists('/usr/X11R6/include'):
            include_dirs.append('/usr/X11R6/include')
            added_lib_dirs.append('/usr/X11R6/lib64')
            added_lib_dirs.append('/usr/X11R6/lib')
        elif os.path.exists('/usr/X11R5/include'):
            include_dirs.append('/usr/X11R5/include')
            added_lib_dirs.append('/usr/X11R5/lib')
        else:
            # Assume default location for X11
            include_dirs.append('/usr/X11/include')
            added_lib_dirs.append('/usr/X11/lib')

        # If Cygwin, then verify that X is installed before proceeding
        if CYGWIN:
            x11_inc = find_file('X11/Xlib.h', [], include_dirs)
            if x11_inc is None:
                return False

        # Check for BLT extension
        if self.compiler.find_library_file(self.lib_dirs + added_lib_dirs,
                                               'BLT8.0'):
            defs.append( ('WITH_BLT', 1) )
            libs.append('BLT8.0')
        elif self.compiler.find_library_file(self.lib_dirs + added_lib_dirs,
                                                'BLT'):
            defs.append( ('WITH_BLT', 1) )
            libs.append('BLT')

        # Add the Tcl/Tk libraries
        libs.append('tk'+ version)
        libs.append('tcl'+ version)

        # Finally, link with the X11 libraries (not appropriate on cygwin)
        if not CYGWIN:
            libs.append('X11')

        # XXX handle these, but how to detect?
        # *** Uncomment and edit for PIL (TkImaging) extension only:
        #       -DWITH_PIL -I../Extensions/Imaging/libImaging  tkImaging.c \
        # *** Uncomment and edit for TOGL extension only:
        #       -DWITH_TOGL togl.c \
        # *** Uncomment these for TOGL extension only:
        #       -lGL -lGLU -lXext -lXmu \

        self.add(Extension('_tkinter', ['_tkinter.c', 'tkappinit.c'],
                           define_macros=[('WITH_APPINIT', 1)] + defs,
                           include_dirs=include_dirs,
                           libraries=libs,
                           library_dirs=added_lib_dirs))
        return True

    def configure_ctypes(self, ext):
        return True

    def detect_ctypes(self):
        # Thomas Heller's _ctypes module

        if (not sysconfig.get_config_var("LIBFFI_INCLUDEDIR") and MACOS):
            self.use_system_libffi = True
        else:
            self.use_system_libffi = '--with-system-ffi' in sysconfig.get_config_var("CONFIG_ARGS")

        include_dirs = []
        extra_compile_args = ['-DPy_BUILD_CORE_MODULE']
        extra_link_args = []
        sources = ['_ctypes/_ctypes.c',
                   '_ctypes/callbacks.c',
                   '_ctypes/callproc.c',
                   '_ctypes/stgdict.c',
                   '_ctypes/cfield.c']
        depends = ['_ctypes/ctypes.h']

        if MACOS:
            sources.append('_ctypes/malloc_closure.c')
            extra_compile_args.append('-DUSING_MALLOC_CLOSURE_DOT_C=1')
            extra_compile_args.append('-DMACOSX')
            include_dirs.append('_ctypes/darwin')

        elif HOST_PLATFORM == 'sunos5':
            # XXX This shouldn't be necessary; it appears that some
            # of the assembler code is non-PIC (i.e. it has relocations
            # when it shouldn't. The proper fix would be to rewrite
            # the assembler code to be PIC.
            # This only works with GCC; the Sun compiler likely refuses
            # this option. If you want to compile ctypes with the Sun
            # compiler, please research a proper solution, instead of
            # finding some -z option for the Sun compiler.
            extra_link_args.append('-mimpure-text')

        elif HOST_PLATFORM.startswith('hp-ux'):
            extra_link_args.append('-fPIC')

        ext = Extension('_ctypes',
                        include_dirs=include_dirs,
                        extra_compile_args=extra_compile_args,
                        extra_link_args=extra_link_args,
                        libraries=[],
                        sources=sources,
                        depends=depends)
        self.add(ext)
        if TEST_EXTENSIONS:
            # function my_sqrt() needs libm for sqrt()
            self.add(Extension('_ctypes_test',
                               sources=['_ctypes/_ctypes_test.c'],
                               libraries=['m']))

        ffi_inc = sysconfig.get_config_var("LIBFFI_INCLUDEDIR")
        ffi_lib = None

        ffi_inc_dirs = self.inc_dirs.copy()
        if MACOS:
            ffi_in_sdk = os.path.join(macosx_sdk_root(), "usr/include/ffi")

            if not ffi_inc:
                if os.path.exists(ffi_in_sdk):
                    ext.extra_compile_args.append("-DUSING_APPLE_OS_LIBFFI=1")
                    ffi_inc = ffi_in_sdk
                    ffi_lib = 'ffi'
                else:
                    # OS X 10.5 comes with libffi.dylib; the include files are
                    # in /usr/include/ffi
                    ffi_inc_dirs.append('/usr/include/ffi')

        if not ffi_inc:
            found = find_file('ffi.h', [], ffi_inc_dirs)
            if found:
                ffi_inc = found[0]
        if ffi_inc:
            ffi_h = ffi_inc + '/ffi.h'
            if not os.path.exists(ffi_h):
                ffi_inc = None
                print('Header file {} does not exist'.format(ffi_h))
        if ffi_lib is None and ffi_inc:
            for lib_name in ('ffi', 'ffi_pic'):
                if (self.compiler.find_library_file(self.lib_dirs, lib_name)):
                    ffi_lib = lib_name
                    break

        if ffi_inc and ffi_lib:
            ffi_headers = glob(os.path.join(ffi_inc, '*.h'))
            if grep_headers_for('ffi_prep_cif_var', ffi_headers):
                ext.extra_compile_args.append("-DHAVE_FFI_PREP_CIF_VAR=1")
            if grep_headers_for('ffi_prep_closure_loc', ffi_headers):
                ext.extra_compile_args.append("-DHAVE_FFI_PREP_CLOSURE_LOC=1")
            if grep_headers_for('ffi_closure_alloc', ffi_headers):
                ext.extra_compile_args.append("-DHAVE_FFI_CLOSURE_ALLOC=1")

            ext.include_dirs.append(ffi_inc)
            ext.libraries.append(ffi_lib)
            self.use_system_libffi = True

        if sysconfig.get_config_var('HAVE_LIBDL'):
            # for dlopen, see bpo-32647
            ext.libraries.append('dl')

    def detect_decimal(self):
        # Stefan Krah's _decimal module
        extra_compile_args = []
        undef_macros = []
        if '--with-system-libmpdec' in sysconfig.get_config_var("CONFIG_ARGS"):
            include_dirs = []
            libraries = [':libmpdec.so.2']
            sources = ['_decimal/_decimal.c']
            depends = ['_decimal/docstrings.h']
        else:
            include_dirs = [os.path.abspath(os.path.join(self.srcdir,
                                                         'Modules',
                                                         '_decimal',
                                                         'libmpdec'))]
            libraries = ['m']
            sources = [
              '_decimal/_decimal.c',
              '_decimal/libmpdec/basearith.c',
              '_decimal/libmpdec/constants.c',
              '_decimal/libmpdec/context.c',
              '_decimal/libmpdec/convolute.c',
              '_decimal/libmpdec/crt.c',
              '_decimal/libmpdec/difradix2.c',
              '_decimal/libmpdec/fnt.c',
              '_decimal/libmpdec/fourstep.c',
              '_decimal/libmpdec/io.c',
              '_decimal/libmpdec/mpalloc.c',
              '_decimal/libmpdec/mpdecimal.c',
              '_decimal/libmpdec/numbertheory.c',
              '_decimal/libmpdec/sixstep.c',
              '_decimal/libmpdec/transpose.c',
              ]
            depends = [
              '_decimal/docstrings.h',
              '_decimal/libmpdec/basearith.h',
              '_decimal/libmpdec/bits.h',
              '_decimal/libmpdec/constants.h',
              '_decimal/libmpdec/convolute.h',
              '_decimal/libmpdec/crt.h',
              '_decimal/libmpdec/difradix2.h',
              '_decimal/libmpdec/fnt.h',
              '_decimal/libmpdec/fourstep.h',
              '_decimal/libmpdec/io.h',
              '_decimal/libmpdec/mpalloc.h',
              '_decimal/libmpdec/mpdecimal.h',
              '_decimal/libmpdec/numbertheory.h',
              '_decimal/libmpdec/sixstep.h',
              '_decimal/libmpdec/transpose.h',
              '_decimal/libmpdec/typearith.h',
              '_decimal/libmpdec/umodarith.h',
              ]

        config = {
          'x64':     [('CONFIG_64','1'), ('ASM','1')],
          'uint128': [('CONFIG_64','1'), ('ANSI','1'), ('HAVE_UINT128_T','1')],
          'ansi64':  [('CONFIG_64','1'), ('ANSI','1')],
          'ppro':    [('CONFIG_32','1'), ('PPRO','1'), ('ASM','1')],
          'ansi32':  [('CONFIG_32','1'), ('ANSI','1')],
          'ansi-legacy': [('CONFIG_32','1'), ('ANSI','1'),
                          ('LEGACY_COMPILER','1')],
          'universal':   [('UNIVERSAL','1')]
        }

        cc = sysconfig.get_config_var('CC')
        sizeof_size_t = sysconfig.get_config_var('SIZEOF_SIZE_T')
        machine = os.environ.get('PYTHON_DECIMAL_WITH_MACHINE')

        if machine:
            # Override automatic configuration to facilitate testing.
            define_macros = config[machine]
        elif MACOS:
            # Universal here means: build with the same options Python
            # was built with.
            define_macros = config['universal']
        elif sizeof_size_t == 8:
            if sysconfig.get_config_var('HAVE_GCC_ASM_FOR_X64'):
                define_macros = config['x64']
            elif sysconfig.get_config_var('HAVE_GCC_UINT128_T'):
                define_macros = config['uint128']
            else:
                define_macros = config['ansi64']
        elif sizeof_size_t == 4:
            ppro = sysconfig.get_config_var('HAVE_GCC_ASM_FOR_X87')
            if ppro and ('gcc' in cc or 'clang' in cc) and \
               not 'sunos' in HOST_PLATFORM:
                # solaris: problems with register allocation.
                # icc >= 11.0 works as well.
                define_macros = config['ppro']
                extra_compile_args.append('-Wno-unknown-pragmas')
            else:
                define_macros = config['ansi32']
        else:
            raise DistutilsError("_decimal: unsupported architecture")

        # Workarounds for toolchain bugs:
        if sysconfig.get_config_var('HAVE_IPA_PURE_CONST_BUG'):
            # Some versions of gcc miscompile inline asm:
            # http://gcc.gnu.org/bugzilla/show_bug.cgi?id=46491
            # http://gcc.gnu.org/ml/gcc/2010-11/msg00366.html
            extra_compile_args.append('-fno-ipa-pure-const')
        if sysconfig.get_config_var('HAVE_GLIBC_MEMMOVE_BUG'):
            # _FORTIFY_SOURCE wrappers for memmove and bcopy are incorrect:
            # http://sourceware.org/ml/libc-alpha/2010-12/msg00009.html
            undef_macros.append('_FORTIFY_SOURCE')

        # Uncomment for extra functionality:
        #define_macros.append(('EXTRA_FUNCTIONALITY', 1))
        self.add(Extension('_decimal',
                           include_dirs=include_dirs,
                           libraries=libraries,
                           define_macros=define_macros,
                           undef_macros=undef_macros,
                           extra_compile_args=extra_compile_args,
                           sources=sources,
                           depends=depends))

    def detect_openssl_hashlib(self):
        # Detect SSL support for the socket module (via _ssl)
        config_vars = sysconfig.get_config_vars()

        def split_var(name, sep):
            # poor man's shlex, the re module is not available yet.
            value = config_vars.get(name)
            if not value:
                return ()
            # This trick works because ax_check_openssl uses --libs-only-L,
            # --libs-only-l, and --cflags-only-I.
            value = ' ' + value
            sep = ' ' + sep
            return [v.strip() for v in value.split(sep) if v.strip()]

        openssl_includes = split_var('OPENSSL_INCLUDES', '-I')
        openssl_libdirs = split_var('OPENSSL_LDFLAGS', '-L')
        openssl_libs = split_var('OPENSSL_LIBS', '-l')
        if not openssl_libs:
            # libssl and libcrypto not found
            self.missing.extend(['_ssl', '_hashlib'])
            return None, None

        # Find OpenSSL includes
        ssl_incs = find_file(
            'openssl/ssl.h', self.inc_dirs, openssl_includes
        )
        if ssl_incs is None:
            self.missing.extend(['_ssl', '_hashlib'])
            return None, None

        # OpenSSL 1.0.2 uses Kerberos for KRB5 ciphers
        krb5_h = find_file(
            'krb5.h', self.inc_dirs,
            ['/usr/kerberos/include']
        )
        if krb5_h:
            ssl_incs.extend(krb5_h)

        if config_vars.get("HAVE_X509_VERIFY_PARAM_SET1_HOST"):
            self.add(Extension(
                '_ssl', ['_ssl.c'],
                include_dirs=openssl_includes,
                library_dirs=openssl_libdirs,
                libraries=openssl_libs,
                depends=['socketmodule.h', '_ssl/debughelpers.c'])
            )
        else:
            self.missing.append('_ssl')

        self.add(Extension('_hashlib', ['_hashopenssl.c'],
                           depends=['hashlib.h'],
                           include_dirs=openssl_includes,
                           library_dirs=openssl_libdirs,
                           libraries=openssl_libs))

    def detect_hash_builtins(self):
        # By default we always compile these even when OpenSSL is available
        # (issue #14693). It's harmless and the object code is tiny
        # (40-50 KiB per module, only loaded when actually used).  Modules can
        # be disabled via the --with-builtin-hashlib-hashes configure flag.
        supported = {"md5", "sha1", "sha256", "sha512", "sha3", "blake2"}

        configured = sysconfig.get_config_var("PY_BUILTIN_HASHLIB_HASHES")
        configured = configured.strip('"').lower()
        configured = {
            m.strip() for m in configured.split(",")
        }

        self.disabled_configure.extend(
            sorted(supported.difference(configured))
        )

        if "sha256" in configured:
            self.add(Extension(
                '_sha256', ['sha256module.c'],
                extra_compile_args=['-DPy_BUILD_CORE_MODULE'],
                depends=['hashlib.h']
            ))

        if "sha512" in configured:
            self.add(Extension(
                '_sha512', ['sha512module.c'],
                extra_compile_args=['-DPy_BUILD_CORE_MODULE'],
                depends=['hashlib.h']
            ))

        if "md5" in configured:
            self.add(Extension(
                '_md5', ['md5module.c'],
                depends=['hashlib.h']
            ))

        if "sha1" in configured:
            self.add(Extension(
                '_sha1', ['sha1module.c'],
                depends=['hashlib.h']
            ))

        if "blake2" in configured:
            blake2_deps = glob(
                os.path.join(escape(self.srcdir), 'Modules/_blake2/impl/*')
            )
            blake2_deps.append('hashlib.h')
            self.add(Extension(
                '_blake2',
                [
                    '_blake2/blake2module.c',
                    '_blake2/blake2b_impl.c',
                    '_blake2/blake2s_impl.c'
                ],
                depends=blake2_deps
            ))

        if "sha3" in configured:
            sha3_deps = glob(
                os.path.join(escape(self.srcdir), 'Modules/_sha3/kcp/*')
            )
            sha3_deps.append('hashlib.h')
            self.add(Extension(
                '_sha3',
                ['_sha3/sha3module.c'],
                depends=sha3_deps
            ))

    def detect_nis(self):
        if MS_WINDOWS or CYGWIN or HOST_PLATFORM == 'qnx6':
            self.missing.append('nis')
            return

        libs = []
        library_dirs = []
        includes_dirs = []

        # bpo-32521: glibc has deprecated Sun RPC for some time. Fedora 28
        # moved headers and libraries to libtirpc and libnsl. The headers
        # are in tircp and nsl sub directories.
        rpcsvc_inc = find_file(
            'rpcsvc/yp_prot.h', self.inc_dirs,
            [os.path.join(inc_dir, 'nsl') for inc_dir in self.inc_dirs]
        )
        rpc_inc = find_file(
            'rpc/rpc.h', self.inc_dirs,
            [os.path.join(inc_dir, 'tirpc') for inc_dir in self.inc_dirs]
        )
        if rpcsvc_inc is None or rpc_inc is None:
            # not found
            self.missing.append('nis')
            return
        includes_dirs.extend(rpcsvc_inc)
        includes_dirs.extend(rpc_inc)

        if self.compiler.find_library_file(self.lib_dirs, 'nsl'):
            libs.append('nsl')
        else:
            # libnsl-devel: check for libnsl in nsl/ subdirectory
            nsl_dirs = [os.path.join(lib_dir, 'nsl') for lib_dir in self.lib_dirs]
            libnsl = self.compiler.find_library_file(nsl_dirs, 'nsl')
            if libnsl is not None:
                library_dirs.append(os.path.dirname(libnsl))
                libs.append('nsl')

        if self.compiler.find_library_file(self.lib_dirs, 'tirpc'):
            libs.append('tirpc')

        self.add(Extension('nis', ['nismodule.c'],
                           libraries=libs,
                           library_dirs=library_dirs,
                           include_dirs=includes_dirs))


class PyBuildInstall(install):
    # Suppress the warning about installation into the lib_dynload
    # directory, which is not in sys.path when running Python during
    # installation:
    def initialize_options (self):
        install.initialize_options(self)
        self.warn_dir=0

    # Customize subcommands to not install an egg-info file for Python
    sub_commands = [('install_lib', install.has_lib),
                    ('install_headers', install.has_headers),
                    ('install_scripts', install.has_scripts),
                    ('install_data', install.has_data)]


class PyBuildInstallLib(install_lib):
    # Do exactly what install_lib does but make sure correct access modes get
    # set on installed directories and files. All installed files with get
    # mode 644 unless they are a shared library in which case they will get
    # mode 755. All installed directories will get mode 755.

    # this is works for EXT_SUFFIX too, which ends with SHLIB_SUFFIX
    shlib_suffix = sysconfig.get_config_var("SHLIB_SUFFIX")

    def install(self):
        outfiles = install_lib.install(self)
        self.set_file_modes(outfiles, 0o644, 0o755)
        self.set_dir_modes(self.install_dir, 0o755)
        return outfiles

    def set_file_modes(self, files, defaultMode, sharedLibMode):
        if not files: return

        for filename in files:
            if os.path.islink(filename): continue
            mode = defaultMode
            if filename.endswith(self.shlib_suffix): mode = sharedLibMode
            log.info("changing mode of %s to %o", filename, mode)
            if not self.dry_run: os.chmod(filename, mode)

    def set_dir_modes(self, dirname, mode):
        for dirpath, dirnames, fnames in os.walk(dirname):
            if os.path.islink(dirpath):
                continue
            log.info("changing mode of %s to %o", dirpath, mode)
            if not self.dry_run: os.chmod(dirpath, mode)


class PyBuildScripts(build_scripts):
    def copy_scripts(self):
        outfiles, updated_files = build_scripts.copy_scripts(self)
        fullversion = '-{0[0]}.{0[1]}'.format(sys.version_info)
        minoronly = '.{0[1]}'.format(sys.version_info)
        newoutfiles = []
        newupdated_files = []
        for filename in outfiles:
            if filename.endswith('2to3'):
                newfilename = filename + fullversion
            else:
                newfilename = filename + minoronly
            log.info('renaming %s to %s', filename, newfilename)
            os.rename(filename, newfilename)
            newoutfiles.append(newfilename)
            if filename in updated_files:
                newupdated_files.append(newfilename)
        return newoutfiles, newupdated_files


def main():
    set_compiler_flags('CFLAGS', 'PY_CFLAGS_NODIST')
    set_compiler_flags('LDFLAGS', 'PY_LDFLAGS_NODIST')

    class DummyProcess:
        """Hack for parallel build"""
        ProcessPoolExecutor = None

    sys.modules['concurrent.futures.process'] = DummyProcess
    validate_tzpath()

    # turn off warnings when deprecated modules are imported
    import warnings
    warnings.filterwarnings("ignore",category=DeprecationWarning)
    setup(# PyPI Metadata (PEP 301)
          name = "Python",
          version = sys.version.split()[0],
          url = "http://www.python.org/%d.%d" % sys.version_info[:2],
          maintainer = "Guido van Rossum and the Python community",
          maintainer_email = "python-dev@python.org",
          description = "A high-level object-oriented programming language",
          long_description = SUMMARY.strip(),
          license = "PSF license",
          classifiers = [x for x in CLASSIFIERS.split("\n") if x],
          platforms = ["Many"],

          # Build info
          cmdclass = {'build_ext': PyBuildExt,
                      'build_scripts': PyBuildScripts,
                      'install': PyBuildInstall,
                      'install_lib': PyBuildInstallLib},
          # The struct module is defined here, because build_ext won't be
          # called unless there's at least one extension module defined.
          ext_modules=[Extension('_struct', ['_struct.c'])],

          # If you change the scripts installed here, you also need to
          # check the PyBuildScripts command above, and change the links
          # created by the bininstall target in Makefile.pre.in
          scripts = ["Tools/scripts/pydoc3", "Tools/scripts/idle3",
                     "Tools/scripts/2to3"]
        )

# --install-platlib
if __name__ == '__main__':
    main()

external/cpython_master/setup.py
DragonHPC/dragon
import json

# We first need to run: dragon-network-config --wlm slurm --output-to-json
# this script reads from the slurm.json the ip addresses of the allocation 
# and prints a list that can be then used from the manual_launch.sh to set up a ray cluster

def get_ip_addresses(json_data):
    ip_addresses = []

    for key, value in json_data.items():
        ip_addrs = value["ip_addrs"]
        for ip_addr in ip_addrs:
            ip = ip_addr.split(':')[0]
            ip_addresses.append(ip)

    return ip_addresses

def main():
    with open('slurm.json', 'r') as json_file:
        json_data = json.load(json_file)

    ip_addresses = get_ip_addresses(json_data)
    print(' '.join(ip_addresses))

if __name__ == "__main__":
    main()


related_work/parallel_python_ray_numerical_computation/get_ips_from_json.py
DragonHPC/dragon
import argparse
import numpy as np
import psutil
import ray
import scipy.signal
import time


def get_args():
    parser = argparse.ArgumentParser(description='Basic SciPy test with Ray')

    parser.add_argument('--num_workers', type=int, default=4,
                        help='number of workers')

    parser.add_argument('--iterations', type=int, default=10,
                        help='number of iterations to do')

    parser.add_argument('--burns', type=int, default=2,
                        help='number of iterations to burn/ignore in order to warm up')

    parser.add_argument('--size', type=int, default=1000,
                        help='size of the array')

    parser.add_argument('--mem', type=int, default=(512 * 1024 * 1024),
                        help='overall footprint of image dataset to process')

    parser.add_argument('--work_time', type=float, default=1.0,
                        help='how many seconds of compute per image')

    parser.add_argument('--multinode', action='store_true',
                        help='run on multiple nodes')

    parser.add_argument('--ipaddress', type=str,
                        help='for multinode only - ip address of the head node of the Ray cluster')

    my_args = parser.parse_args()

    return my_args


@ray.remote
def f(image, random_filter, work_time):
    # Do some image processing
    elapsed = 0.
    start = time.perf_counter()

    # Explicitly control compute time per image
    while elapsed < work_time:
        scipy.signal.convolve2d(image, random_filter)[::5, ::5]
        elapsed = time.perf_counter() - start
    return scipy.signal.convolve2d(image, random_filter)[::5, ::5]


if __name__ == '__main__':
    args = get_args()

    # Image and filter generation
    image = np.zeros((args.size, args.size))
    nimages = int(float(args.mem) / float(image.size))
    print(f'Number of images: {nimages}')
    images = []
    for j in range(nimages):
        images.append(np.zeros((args.size, args.size)))
    filters = [np.random.normal(size=(4, 4)) for _ in range(nimages)]

    # Ray initialization
    # when using multiple nodes, the number of workers
    # is defined during the setup of the Ray cluster
    if args.multinode is False: # single node case
        num_cpus = args.num_workers
        print(f'Number of workers: {num_cpus}')
        ray.init(num_cpus=num_cpus)
    else:
        ray.init(address="auto", _node_ip_address=args.ipaddress)

    # Main body of the computation
    times = []
    for i in range(args.iterations + args.burns):
        start = time.perf_counter()
        remotes = []
        for j in range(nimages):
            remotes.append(f.remote(ray.put(images[j]), filters[j], float(args.work_time)))
        res = ray.get(remotes)
        del(res)
        del(remotes)
        times.append(time.perf_counter() - start)
        print(f"Time for iteration {i} is {times[i]} seconds", flush=True)

    # Print out the mean and standard deviation, excluding the burns iterations
    print(f"Average time: {round(np.mean(times[args.burns:]), 2)} second(s)")
    print(f"Standard deviation: {round(np.std(times[args.burns:]), 2)} second(s)")

related_work/parallel_python_ray_numerical_computation/parallel_python_ray_numerical_computation.py
DragonHPC/dragon
""" This code is first executed when the `import dragon` statement is encountered
in the user code.

If the environment variable `DRAGON_PATCH_MP` is set, we assume the user wants
to use Dragon below the standard Multiprocessing module. We then monkeypatch the
multiprocessing module to hook in our own context into the context hierarchy:

    BaseContext <|- DefaultContext <|- AugmentedDefaultContext

Our `AugmentedDefaultContext` replaces the `DefaultContext`, and replaces the
`get_start_method` and `set_start_method` to make our context available.

We collect the public API of the Multiprocessing API and point it towards
our context, that contains simple functions returning our Dragon classes.

Finally, we need to make sure that imports like this one:
```
import multiprocessing.context.wait
```
or inheritance like this:
```
class Foo(multiprocessing.synchronize.Event):
```
work with our Dragon classes. To do so, we save all Multiprocessing classes in
the variable `original_multiprocessing` that can be imported from the
dragon.mpbridge.monkeypatching module and replace them with our own Dragon
classes.

Note that this final step cannot be completely undone when switching contexts.

I.e. switching from start method "dragon" back to a start method "spawn" will
break user code that has inherited or imported directly from Multiprocessing.

General advice to users is to use the context to get the Multiprocessing
functionality.
"""

import os
from distutils.util import strtobool


def _patch_multiprocessing():
    # Set DRAGON_PATCH_MP so multiprocessing is automatically patched when
    # imported in a subprocess.
    os.environ["DRAGON_PATCH_MP"] = str(True)

    from .mpbridge.monkeypatching import patch_multiprocessing
    patch_multiprocessing()


if bool(strtobool(os.environ.get("DRAGON_PATCH_MP", "False"))):
    _patch_multiprocessing()

src/dragon/__init__.py
DragonHPC/dragon
def main():
    import sys

    from .cli.__main__ import main as _main
    from .infrastructure.facts import PROCNAME_LA_FE

    _main([PROCNAME_LA_FE] + sys.argv[1:])


if __name__ == '__main__':
    main()

src/dragon/__main__.py
DragonHPC/dragon
import sys
from typing import Any

try:
    # Python 3.10+ provides EntryPoints._from_text()
    from importlib.metadata import EntryPoints
except ImportError:
    # Python 3.9 provides EntryPoint._from_text()
    from importlib.metadata import EntryPoint
    # Emulate Python 3.10 interface
    class EntryPoints:
        @staticmethod
        def _from_text(text):
            return EntryPoint._from_text(text)


PROCNAME_GS = 'dragon-globalservices'
PROCNAME_LA_FE = 'dragon'
PROCNAME_LA_BE = 'dragon-backend'
PROCNAME_LS = 'dragon-localservices'
PROCNAME_TCP_TA = 'dragon-tcp'
PROCNAME_OVERLAY_TA = 'dragon-overlay-tcp'
PROCNAME_NETWORK_CONFIG = 'dragon-network-config'
PROCNAME_NETWORK_CONFIG_LAUNCH_HELPER = 'dragon-network-config-launch-helper'
PROCNAME_NETWORK_CONFIG_SHUTDOWN_HELPER = 'dragon-network-config-shutdown-helper'
PROCNAME_NETWORK_IFADDRS='dragon-network-ifaddrs'

# TODO Refactor frontend entry point. See ../__main__.py and
# TODO ../launcher/launch_selector.py.

# TODO Refactor local services entry point. Need CLI option to select single
# TODO vs multi mode.

# TODO Refactor global services entry point. Need to replace facts
# TODO GS_SINGLE_LAUNCH_CMD and GS_MULTI_LAUNCH_CMD, which are used in
# TODO ../launcher/dragon_single.py and ../launcher/launch_multi_ls.py,
# TODO respectively. Combine into single CLI entry point with option to
# TODO configure mode: single vs multi.

entry_points = f"""
[console_scripts]
{PROCNAME_LA_FE} = dragon.launcher.launch_selector:main
{PROCNAME_LA_BE} = dragon.launcher.dragon_multi_be:main
{PROCNAME_LS} = dragon.launcher.launch_multi_ls:main
{PROCNAME_GS} = dragon.globalservices.server:multi
{PROCNAME_TCP_TA} = dragon.transport.tcp.__main__:main
{PROCNAME_NETWORK_CONFIG} = dragon.launcher.network_config:main
{PROCNAME_NETWORK_CONFIG_LAUNCH_HELPER} = dragon.launcher.network_config:deliver_backend_node_descriptor
{PROCNAME_NETWORK_CONFIG_SHUTDOWN_HELPER} = dragon.launcher.network_config:shutdown_backend_node_deliverer
{PROCNAME_NETWORK_IFADDRS} = dragon.transport.ifaddrs:main
"""


# Mapping of script name to EntryPoint object
console_scripts = {
    ep.name: ep
    for ep in EntryPoints._from_text(entry_points)
    if ep.group == 'console_scripts'
}

# Sanity check
if not console_scripts:
    raise RuntimeError(f'No console scripts found')


def console_script_args(name: str, *args: Any) -> list:
    """Returns list of command line arguments for the specified console script."""
    if name not in console_scripts:
        raise ValueError(f'Unknown console script: {name}')
    _args = [sys.executable, '-m', 'dragon.cli', name]
    _args.extend(args)
    return _args

src/dragon/cli/__init__.py
DragonHPC/dragon
def main(args=None):
    import argparse
    import sys

    from . import console_scripts

    name_padding = min(24, max(map(len, console_scripts.keys())))

    parser = argparse.ArgumentParser(
        prog='dragon',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        usage='SCRIPT ...',
        description='Run dragon console script',
        epilog='console scripts:\n' + '\n'.join(
            f'  {{:<{name_padding}}}  {{:<}}'.format(ep.name, ep.value)
            for ep in console_scripts.values()
        ),
        add_help=False,
    )

    parser.add_argument(
        'script',
        nargs='?',
        choices=console_scripts.keys(),
        help=argparse.SUPPRESS,
    )

    args, remaining = parser.parse_known_args(args)

    # Print help by default
    if not args.script:
        parser.print_help()
        parser.exit(1)

    # Patch sys.argv so the console script believes it was called correctly
    sys.argv.clear()
    sys.argv.append(args.script)
    sys.argv.extend(remaining)

    # Load and exec the console script
    _main = console_scripts[args.script].load()
    _main()


if __name__ == '__main__':
    main()

src/dragon/cli/__main__.py
DragonHPC/dragon
"""
Dragon dictionary managers
"""
import os
import sys
import enum
import pickle
import logging


from ... import managed_memory
from ...globalservices import channel
from ...globalservices import pool
from ...infrastructure import facts
from ...infrastructure import parameters
from ...channels import Channel, Message, Many2ManyWritingChannelFile, Many2ManyReadingChannelFile

LOG = logging.getLogger(__name__)

# Define the byteorder as the native system byteorder
_BYTEORDER = sys.byteorder

@enum.unique
class DictOp(enum.Enum):
    """Action to be performed on the dictionary item"""
    INIT_CLIENT = enum.auto()
    CLOSE_CLIENT = enum.auto()
    STOP_DICT = enum.auto()
    SET_ITEM = enum.auto()
    GET_ITEM = enum.auto()
    DEL_ITEM = enum.auto()
    POP_ITEM = enum.auto()
    CONT_ITEM = enum.auto()
    GET_LEN = enum.auto()
    GET_KEYS = enum.auto()
    GET_VALS = enum.auto()
    GET_ITEMS = enum.auto()

@enum.unique
class ReturnCode(enum.Enum):
    """Code returned by the dictionary manager"""
    SUCCESS_RET_VAL = enum.auto()
    SUCCESS_RET_TRUE = enum.auto()
    SUCCESS_RET_FALSE = enum.auto()
    FAILED = enum.auto()

def send_message(client_chnl, manager_pool_mem, ret_code, value=None):
    """Send message to the client channel from the manager process.

    :param client_chnl: Serialized channel descriptor of the client to attach to
    :type client_chnl: bytes-like object
    :param manager_pool_mem: Memory pool for the writing adapter
    :type manager_pool_mem: MemoryPool object
    :param ret_code: Return value of the exectued operation to the client
    :type ret_code: int
    :param value: Value fetched from the dictionary and to be returned to the client
    :type value: bytes-like object
    """
    try:
        ret_adapter = Many2ManyWritingChannelFile(client_chnl, buffer_pool=manager_pool_mem)
        ret_adapter.open()
        ret_header = ret_code.value.to_bytes(4, byteorder=_BYTEORDER, signed=True)
        ret_adapter.write(ret_header)

        if value:
            value_size = sys.getsizeof(value)
            ret_adapter.write(value_size.to_bytes(4, byteorder=_BYTEORDER, signed=True))
            ret_adapter.write(value)
    except Exception as e:
        raise ConnectionError(f"Could not complete send bytes operation to the client channel: {e}")
    finally:
        ret_adapter.close()

def manager_proc(pool_size: int, mowner_chnl_name: str) -> None:
    """Manager process creates the associated channel and pool.
    Communicates its information to the manager owner process.
    Waits for the client messages on the created manager channel.

    :param pool_size: Size of the manager pool
    :type pool_size: int
    :param mowner_chnl_name: channel name that acts an identifier(c_uid) to the channel
    :type mowner_chnl_name: string
    """
    # Collect the puid of the manager process
    manager_puid = parameters.this_process.my_puid

    # Allocate pool for the manager process
    manager_pool_name = f'{facts.DEFAULT_DICT_POOL_NAME_BASE}{manager_puid}'
    manager_pool_desc = pool.create(pool_size, user_name=manager_pool_name)
    manager_pool_mem = managed_memory.MemoryPool.attach(manager_pool_desc._sdesc)
    LOG.debug(f"Dictionary manager process with puid:{parameters.this_process.my_puid} \
              created on node-index:{parameters.this_process.index} with pool-name: {manager_pool_name}.")

    # Create the manager channel that receives the requests from the client
    manager_req_chnl = channel.create(manager_pool_desc.m_uid)
    manager_chnl_desc = channel.query(manager_req_chnl.name)
    manager_chnl = Channel.attach(manager_chnl_desc.sdesc)

    # Information of the manager process
    manager_info = {'p_uid': manager_puid,
                    'chnl_desc': manager_chnl_desc,
                    'pool_desc': manager_pool_desc}

    # Send the manager information to the owner channel
    mowner_chnl_desc = channel.query(mowner_chnl_name)
    mowner_chnl = Channel.attach(mowner_chnl_desc.sdesc)
    with Many2ManyWritingChannelFile(mowner_chnl, buffer_pool=manager_pool_mem) as adapter:
        pickle.dump(manager_info, file=adapter, protocol=5)
    mowner_chnl.detach()

    # Manager dictionary that stores key and memory descriptor of the value
    manager_kv = dict()
    manager_client_map = dict()

    serving = True
    while serving:
        try:
            read_adapter = Many2ManyReadingChannelFile(manager_chnl)
            read_adapter.open()
            key_info = pickle.load(read_adapter)
            value = None
            ret_code = ReturnCode.SUCCESS_RET_VAL

            # Identify the client channel from the manager-client channel map
            if key_info['dict_op'] == DictOp.INIT_CLIENT:
                try:
                    client_chnl = Channel.attach(key_info['client_chnl_desc'].sdesc)
                    manager_client_map[key_info['client_puid']] = client_chnl
                    LOG.debug(f"Client with puid:{key_info['client_puid']} initialized by manager process:{parameters.this_process.my_puid} \
                                on node-index:{parameters.this_process.index} with pool-name: {manager_pool_name}.")
                except Exception as e:
                        raise ConnectionError(f"Could not find client channel desc: {e}")
                ret_code = ReturnCode.SUCCESS_RET_TRUE
                send_message(client_chnl, manager_pool_mem, ret_code)
                continue
            elif key_info['dict_op'] == DictOp.CLOSE_CLIENT:
                client_puid = key_info['client_puid']
                if manager_client_map.get(client_puid):
                    try:
                        client_chnl = manager_client_map[client_puid]
                        ret_code = ReturnCode.SUCCESS_RET_TRUE
                        send_message(client_chnl, manager_pool_mem, ret_code)
                        client_chnl.detach()
                        LOG.debug(f"Detached the manager process with puid:{parameters.this_process.my_puid} from client channel {client_chnl}.")
                    except Exception as e:
                        # Client channel must have already deleted
                        pass
                    del manager_client_map[client_puid]
                    continue
            else:
                try:
                    client_chnl = manager_client_map[key_info['client_puid']]
                except Exception as e:
                    raise ConnectionError(f"Could not find client channel desc: {e}")

            if manager_client_map.get(key_info['client_puid']):
                client_chnl = manager_client_map[key_info['client_puid']]
            else:
                try:
                    client_chnl = Channel.attach(key_info['client_chnl_desc'].sdesc)
                    manager_client_map['client_puid'] = client_chnl
                except Exception as e:
                    raise ConnectionError(f"Could not find client channel desc: {e}")

            if key_info['dict_op'] == DictOp.GET_LEN:
                manager_len = len(manager_kv)
                length = pickle.dumps(manager_len)
                send_message(client_chnl, manager_pool_mem, ret_code, value=length)
            elif key_info['dict_op'] == DictOp.SET_ITEM:
                try:
                    mem_alloc_obj, offset = read_adapter.get_remaining_data()
                    mem_alloc_clone_obj = mem_alloc_obj.clone(offset)
                    manager_kv[key_info['key']] = mem_alloc_clone_obj
                    ret_code = ReturnCode.SUCCESS_RET_TRUE
                    send_message(client_chnl, manager_pool_mem, ret_code)
                except Exception as e:
                    raise ConnectionError(f"Could not complete SET {key_info['key']} operation: {e}")
            elif key_info['dict_op'] == DictOp.GET_ITEM:
                if manager_kv.get(key_info['key']):
                    mem_alloc_obj = manager_kv[key_info['key']]
                    msg = Message.create_from_mem(mem_alloc_obj)
                    value = msg.tobytes()
                    send_message(client_chnl, manager_pool_mem, ret_code, value=value)
                else:
                    ret_code = ReturnCode.FAILED
                    send_message(client_chnl, manager_pool_mem, ret_code)
            elif key_info['dict_op'] == DictOp.DEL_ITEM:
                if manager_kv.get(key_info['key']):
                    del manager_kv[key_info['key']]
                    ret_code = ReturnCode.SUCCESS_RET_TRUE
                else:
                    ret_code = ReturnCode.SUCCESS_RET_FALSE
                send_message(client_chnl, manager_pool_mem, ret_code)
            elif key_info['dict_op'] == DictOp.POP_ITEM:
                if manager_kv.get(key_info['key']):
                    mem_alloc_obj = manager_kv[key_info['key']]
                    msg = Message.create_from_mem(mem_alloc_obj)
                    value = msg.tobytes()
                    del manager_kv[key_info['key']]
                    send_message(client_chnl, manager_pool_mem, ret_code, value=value)
                else:
                    ret_code = ReturnCode.SUCCESS_RET_FALSE
                    send_message(client_chnl, manager_pool_mem, ret_code)
            elif key_info['dict_op'] == DictOp.CONT_ITEM:
                if manager_kv.get(key_info['key']):
                    ret_code = ReturnCode.SUCCESS_RET_TRUE
                else:
                    ret_code = ReturnCode.SUCCESS_RET_FALSE
                send_message(client_chnl, manager_pool_mem, ret_code)
            elif key_info['dict_op'] == DictOp.GET_KEYS:
                value = pickle.dumps(list(manager_kv.keys()))
                send_message(client_chnl, manager_pool_mem, ret_code, value=value)
            elif key_info['dict_op'] == DictOp.GET_VALS:
                values = []
                for value in manager_kv.values():
                    msg = Message.create_from_mem(value)
                    values.append(pickle.loads(msg.tobytes()))
                value = pickle.dumps(values)
                send_message(client_chnl, manager_pool_mem, ret_code, value=value)
            elif key_info['dict_op'] == DictOp.GET_ITEMS:
                items = []
                for key, value in manager_kv.items():
                    msg = Message.create_from_mem(value)
                    items.append((key, pickle.loads(msg.tobytes())))
                value = pickle.dumps(items)
                send_message(client_chnl, manager_pool_mem, ret_code, value=value)
            elif key_info['dict_op'] == DictOp.STOP_DICT:
                serving = False
            else:
                continue
        except Exception as e:
            raise ConnectionError(f"Could not complete receive bytes operation on the manager channel: {e}")
        finally:
            read_adapter.close()

    manager_pool_mem.destroy()
src/dragon/data/distdictionary/dict_managers.py
DragonHPC/dragon
""" The Dragon Dictionary is a performant and distributed solution that is
available to the applications and workflows written for the Dragon ecosystem.

This is Dragon's specialized implementation based on Dragon channels and memory,
that works across single and multi node settings. The Dragon Dictionary is designed
to be used as a standard Python Dictionary.

The Dragon Dictionary creates a manager owner process, along with manager processes,
and their associated dragon channels, memory pools. The communication with the client
is established by hashing and redirecting the client to the correct manager process.
"""

import os
import sys
import enum
import uuid
import dragon
import pickle
import logging
import multiprocessing

from ...utils import B64
from ... import managed_memory
from ...infrastructure.process_desc import ProcessOptions
from ...globalservices import channel
from ...globalservices import pool
from ...globalservices import process
from ...infrastructure import facts
from ...infrastructure import parameters
from ...channels import Channel, Message, Many2ManyWritingChannelFile, Many2ManyReadingChannelFile
from ...native.process import TemplateProcess, Process
from ...native.process_group import ProcessGroup
from .dict_managers import manager_proc

LOG = logging.getLogger(__name__)

class managerPool(ProcessGroup):
    """Inherited from the native pool workers. The manager pool is intended to maintain
    specific variables and functionality requirements to maintain the pool of managers.

    :param nworkers: number of manager processes to start
    :type nworkers: int
    :param psize: size of the manager pool
    :type psize: int
    :param cmd: command
    :type cmd: str
    :param args: arguments to cmd
    :type args: tuple
    :param options: options will be used to pass along arbitrary extra info for things like placement
    :type options: tuple
    """

class ManagerOwnerProc:
    """Instantiate a manager owner process object that takes care of handling the manager processes.

    :param managers_per_node: Number of manager processes to be created on each node.
    :type managers_per_node: int
    :param n_nodes: Number of nodes specified for creating the manager processes.
    :type n_nodes: int
    :param total_mem: Total memory for the allocated pools of all the manager processes.
    This parameter along with managers_per_node, n_nodes helps to decide the pool size of each manager.
    This parameter expects to be specified in GB.
    :type total_mem: int
    """
    def __init__(self, managers_per_node, n_nodes, total_mem):
        LOG.debug(f"Initializing manager owner proc")
        self.managers_per_node = managers_per_node
        self.n_nodes = n_nodes
        self.total_mem = total_mem
        # Store the owner process id, channel, and manager processes
        self.mowner_puid = None
        self.mowner_chnl_desc = None
        self.mowner_chnl = None
        self.manager_processes = []
        # Manager pool workers object
        self.manager_pool = None
        # Store the handle to the dictionay object
        self.dict_chnl_desc = None
        self.dict_chnl = None


    def start(self, dict_chnl_name):
        n_workers = (self.n_nodes * self.managers_per_node)
        mpool_size = self.total_mem // n_workers
        LOG.debug(f"Manager pool size is {mpool_size} bytes")

        # Attach the manager owner process to the dictionary channel
        self.dict_chnl_desc = channel.query(dict_chnl_name)
        self.dict_chnl = Channel.attach(self.dict_chnl_desc.sdesc)

        # Create the manager owner channel
        default_muid = facts.default_pool_muid_from_index(parameters.this_process.index)
        mowner_chnl_uuid = str(uuid.uuid4())
        mowner_chnl_create = channel.create(default_muid, user_name="mowner-proc-"+mowner_chnl_uuid)
        self.mowner_chnl_desc = channel.join(mowner_chnl_create.name)
        self.mowner_chnl = Channel.attach(self.mowner_chnl_desc.sdesc)

        # Create the manager owner process
        mownerproc = Process(self.manager_owner_proc)
        mownerproc.start()
        self.mowner_puid = mownerproc.puid

        # Initialize the manager pool of workers
        # Pool workers are manager processes that serve the data to the client
        args = (mpool_size, mowner_chnl_create.name)
        template = TemplateProcess(manager_proc, args=args)
        self.manager_pool = managerPool()
        self.manager_pool.add_process(n_workers, template)
        self.manager_pool.init()
        self.manager_pool.start()

        # Wait for the manager process information from the owner process
        wait = True
        while(wait):
            with Many2ManyReadingChannelFile(self.dict_chnl) as adapter:
                resp_msg = pickle.load(adapter)
                if("manager-info" in resp_msg):
                    manager_proc_info = resp_msg.split(':')
                    bytes = B64.str_to_bytes(manager_proc_info[0])
                    self.manager_processes = pickle.loads(bytes)
                    wait = False

        LOG.debug(f"Manager owner received information from {len(self.manager_processes)} manager workers")


    def manager_owner_proc(self):
        """Manager owner process that holds information about the manager workers.
        Join and wait for the messages from the manager processes on the owner channel
        """
        worker_count = 0
        mproc_puids = []
        manager_processes = []
        n_workers = (self.n_nodes * self.managers_per_node)

        # Worker count for all the manager workers to be created
        while (worker_count < n_workers):
            with Many2ManyReadingChannelFile(self.mowner_chnl) as adapter:
                manager_info = pickle.load(adapter)
                worker_count += 1
                mproc_puids.append(manager_info['p_uid'])
                manager_process = {'p_uid': manager_info['p_uid'],
                                   'chnl_desc': manager_info['chnl_desc'],
                                   'pool_desc': manager_info['pool_desc']}
                manager_processes.append(manager_process)

        # Store the manager processes information onto the client
        msg_str = B64.bytes_to_str(pickle.dumps(manager_processes)) + ":" + "manager-info"
        # Utilize the default pool memory of the node of the manager owner process
        pool_mem = managed_memory.MemoryPool.attach(B64.str_to_bytes(parameters.this_process.default_pd))

        with Many2ManyWritingChannelFile(self.dict_chnl, buffer_pool=pool_mem) as adapter:
            pickle.dump(msg_str, file=adapter, protocol=5)
        pool_mem.detach()

        process.multi_join(mproc_puids, join_all=True)
        LOG.debug(f"Manager owner exiting the dictionary process")


    def free_manager_resources(self):
        """
        Free the dictionary resources. Destroy the manager owner channels and pools.
        """
        # Detach the owner process with the dictionary channel
        LOG.debug(f"Destroying the dictionary channel")
        self.dict_chnl.detach()
        channel.destroy(self.dict_chnl_desc.c_uid)

        # Free up the manager owner channel
        LOG.debug(f"Destroying the manager owner channel")
        self.mowner_chnl.detach()
        channel.destroy(self.mowner_chnl_desc.c_uid)

        # Free up the channel and pool associated with manager processes
        for i in range(0, len(self.manager_processes)):
            manager_chnl_desc = self.manager_processes[i]['chnl_desc']
            channel.destroy(manager_chnl_desc.c_uid)
            manager_pool_desc = self.manager_processes[i]['pool_desc']
            pool.destroy(manager_pool_desc.m_uid)


    def stop(self):
        """
        Kill the manager owner process
        """
        LOG.debug(f"Killing the manager owner process")
        process.kill(self.mowner_puid, sig=process.signal.SIGKILL)


    def __del__(self):
        # Destroy the manager pool object
        del self.manager_pool


class DistributedDict:
    def __init__(self, managers_per_node, n_nodes, total_mem):
        """Instantiate a distributed dictionary creating the manager owner and the processes.
        :param managers_per_node: Number of manager processes to be created on each node.
        :type managers_per_node: int
        :param n_nodes: Number of nodes specified for creating the manager processes.
        :type n_nodes: int
        :param total_mem: Total memory for the allocated pools of all the manager processes.
        This parameter along with managers_per_node, n_nodes helps to decide the pool size
        of each manager. This parameter expects to be specified in GB.
        :type total_mem: int
        """
        self._managers_per_node = managers_per_node
        self._nodes = n_nodes
        self._total_mem = total_mem
        # Manager owner process created as part of the distributed dictionary
        self.manager_owner = None


    def start(self):
        """
        Creates the client channel using the Global Services and the default dragon memory pool.
        Instantiates the manager owner process, that inturn creates all the manager processes and channels.
        """
        # Using default memory pool for the client channel
        dpool_muid = facts.default_pool_muid_from_index(parameters.this_process.index)

        # Ask Global services to create a main channel which listens the manager owner information
        dict_chnl_uuid = str(uuid.uuid4())
        dict_chnl_create = channel.create(dpool_muid, user_name="ddict-chnl-"+dict_chnl_uuid)

        self.manager_owner = ManagerOwnerProc(self.managers_per_node, self.nodes, self.total_mem)
        self.manager_owner.start(dict_chnl_create.name)


    @property
    def nodes(self):
        """
        Return the nodes the dictionary spanned across
        """
        return self._nodes


    @property
    def managers_per_node(self):
        """
        Return the managers per node of the dictionary
        """
        return self._managers_per_node


    @property
    def total_mem(self):
        """
        Total allocated memory for the dictionary
        """
        return self._total_mem


    @property
    def manager_processes(self):
        """
        Return the manager processes of the dictionary
        """
        return self.manager_owner.manager_processes


    def retrieve_channel_desc(self, request_str):
        """
        Return the manager channel descriptor to the client
        """
        manager_processes = self.manager_processes
        chnl_num = hash(request_str) % len(manager_processes)
        return manager_processes[chnl_num]['chnl_desc']


    def stop(self):
        """
        Stop the distributed dictionary. Free all the resources including processes,
        associated channels, handles etc.
        """
        LOG.debug(f"Executing the STOP call on the dictionary")
        # Kill all the manager processes
        self.manager_owner.manager_pool.stop()

        # Free all the manager process resources
        self.manager_owner.free_manager_resources()


    def __del__(self):
        # Destroy the manager owner object
        del self.manager_owner


if __name__ == "__main__":
    multiprocessing.set_start_method("dragon")

    # Disable the non-deterministic behavior of the python hashing algorithm behavior
    hashseed = os.getenv('PYTHONHASHSEED')
    if not hashseed:
        os.environ['PYTHONHASHSEED'] = '0'
        os.execv(sys.executable, [sys.executable] + sys.argv)

    try:
        managers_per_node = int(sys.argv[1])
    except IndexError:
        # By default, create only a single manager on each node
        managers_per_node = 1

    try:
        n_nodes = int(sys.argv[2])
    except IndexError:
        # By defualt, consider only the client node
        n_nodes = 1

    try:
        total_mem = int(sys.argv[3])
    except IndexError:
        # By default, allocate 1GB of memory for the dictionary
        total_mem = 1

    # Convert total_mem to bytes
    total_mem *= (1024*1024*1024)

    # Executing the script directly. Just start and stop the distributed dictionary
    distDict = DistributedDict(managers_per_node, n_nodes, total_mem)
    distDict.start()
    distDict.stop()

src/dragon/data/distdictionary/distributed_dict.py
DragonHPC/dragon
﻿"""
Dragon dictionary
"""

import os
import uuid
import pickle
import logging
import multiprocessing as mp

from . import distributed_dict
from .dict_managers import DictOp, ReturnCode, _BYTEORDER
from ...infrastructure import parameters
from ...globalservices import channel
from ...infrastructure import facts
from ... import managed_memory
from ...channels import Channel, Many2ManyWritingChannelFile, Many2ManyReadingChannelFile
from ...utils import B64

LOG = logging.getLogger(__name__)

class DragonDict:
    """
    A dictionary like object created over Dragon channels
    """

    def __init__(self, managers_per_node: int, n_nodes: int, total_mem: int) -> None:
        """Instantiate a distributed dictionary object that acts as a (key,value) store

        :param managers_per_node: Number of manager processes to be created on each node.
        :type managers_per_node: int
        :param n_nodes: Number of nodes specified for creating the manager processes.
        :type n_nodes: int
        :param total_mem: Total memory for the allocated pools of all the manager processes.
        This parameter along with managers_per_node, n_nodes helps to decide the pool size of each manager.
        This parameter expects to be specified in GB.
        :type total_mem: int
        """
        self._dist_dict = distributed_dict.DistributedDict(
            managers_per_node=managers_per_node,
            n_nodes=n_nodes,
            total_mem=total_mem)
        self._dist_dict.start()
        self._client_map = dict()
        self._reset()

    def __setstate__(self, state):
        self._dist_dict = state['_dist_dict']
        self._client_map = state['_client_map']
        self._reset()

    def __getstate__(self):
        dragon_dict = dict()
        dragon_dict['_dist_dict'] = self._dist_dict
        dragon_dict['_client_map'] = self._client_map
        return dragon_dict

    def _send_and_read_message_from_all_managers(self, key_info, ret_val=0):
        """Message will be sent to all the dictionary managers.
        Collect the results from the managers to return the cumulative result.

        :param key_info: Information about the operation to the manager
        :type key_info: dict
        :param ret_val: Cumulative value collected from the managers
        :type ret_val: int/list based on the caller. Defaults to zero.
        """
        manager_processes = self._dist_dict.manager_processes

        # Prepare the dictionary message and send over to the manager
        message = DragonDictMessage(key_info, buffer_pool=self._buffer_pool,
                                    client_chnl=self.client_chnl)
        for i in range(0, len(manager_processes)):
            manager_chnl_desc = manager_processes[i]['chnl_desc']
            message.send_message(manager_chnl_desc=manager_chnl_desc)

        manager_count = 0
        n_managers = (self._dist_dict.managers_per_node * self._dist_dict.nodes)
        for _ in range(n_managers):
            # Wait on the client channel to receive the message from the manager
            value = message.read_message()
            if value and not isinstance(value, bool):
                ret_val += value
        return ret_val

    def _reset(self):
        """Collect the client information from the client map to return the response
        from the dictionary operation. If the client is accessing the dictionary for
        first time, add it's information to the client map. Each manager will hold
        the client information and return the information to the client channel.
        """
        self.default_muid = facts.default_pool_muid_from_index(parameters.this_process.index)
        self._buffer_pool = managed_memory.MemoryPool.attach(B64.str_to_bytes(parameters.this_process.default_pd))
        self.client_puid = parameters.this_process.my_puid

        if self._client_map.get(self.client_puid):
            client_chnl = self._client_map[self.client_puid]
            self.client_chnl_create = client_chnl[0]
            self.client_chnl_desc = client_chnl[1]
        else:
            self.client_chnl_uuid = str(uuid.uuid4())
            self.client_chnl_create = channel.create(self.default_muid, user_name="client-"+self.client_chnl_uuid)
            self.client_chnl_desc = channel.query(self.client_chnl_create.name)
            self.client_chnl = Channel.attach(self.client_chnl_desc.sdesc)
            self._client_map[self.client_puid] = (self.client_chnl_create, self.client_chnl_desc)

            # Construct the information needed for serving the request
            key_info = {'dict_op': DictOp.INIT_CLIENT,
                        'client_puid': self.client_puid,
                        'client_chnl_desc': self.client_chnl_desc}
            self._send_and_read_message_from_all_managers(key_info)

    def close(self):
        """Close the client. This typically implies that client does not want to access
        the dictionary no more. The dragon dictionary will erase the client information
        across it and all the managers. Closing typically implies to the clients, while
        stop call implies closing the entire dragon dictionary.
        """
        # Construct the information needed for closing the client request
        key_info = {'dict_op': DictOp.CLOSE_CLIENT,
                    'client_puid': self.client_puid}

        self._send_and_read_message_from_all_managers(key_info)

        # Remove the client entry from the dictionary
        del self._client_map[self.client_puid]

        # Destroy the client channel
        self.client_chnl.detach()
        channel.destroy(self.client_chnl_desc.c_uid)

    def stop(self):
        """Stop the dragon dictionary. Destroys all the manager processes, associated
        channels and memory pools. Implies no further access to the dictionary object.
        """
        self._dist_dict.stop()

    def __len__(self):
        # Construct the information needed for serving the request
        key_info = {'dict_op': DictOp.GET_LEN,
                    'client_puid': self.client_puid}
        length = self._send_and_read_message_from_all_managers(key_info)
        return length

    def length(self):
        """Length of the dictionary. Collects the number of all the key,value
        pairs across all the managers and returns the cumulative result.
        """
        len = self.__len__()
        return len

    def __getitem__(self, key):
        """GET item in the dictionary. The request will be routed to the
        corresponding manager in the dictionary based on the provided key.

        :param key: Information about the GET operation to the manager
        :type key: Any hashable structure
        :return: Value for the stored key in the dictionary
        :rtype: Any object that can be serialized
        """
        # Retrieve the manager channel descriptor for the request
        manager_chnl_desc = self._dist_dict.retrieve_channel_desc(key)

        # Construct the information needed for serving the request
        key_info = {'key': key, 'dict_op': DictOp.GET_ITEM,
                    'client_puid': self.client_puid}

        # Prepare the dictionary message and send over to the manager
        message = DragonDictMessage(key_info, buffer_pool=self._buffer_pool,
                                    manager_chnl_desc=manager_chnl_desc,
                                    client_chnl=self.client_chnl)
        message.send_message()

        # Wait on the client channel to receive the message from the manager
        client_val = message.read_message()
        LOG.debug(f"GET ITEM: Retrieved {key} successfully from the dictionary.")
        return client_val

    def __setitem__(self, key, value):
        """SET item in the dictionary. The request will be routed to the
        corresponding manager in the dictionary based on the provided key.

        :param key: Information about the SET operation to the manager
        :type key: Any hashable structure
        :param value: Value to be stored inside the manager
        :type value: Any structure that can be serialized
        """
        # Retrieve the manager channel descriptor for the request
        manager_chnl_desc = self._dist_dict.retrieve_channel_desc(key)

        # Construct the information needed for serving the request
        key_info = {'key': key,
                    'dict_op': DictOp.SET_ITEM,
                    'client_puid': self.client_puid}

        # Prepare the dictionary message and send over to the manager
        message = DragonDictMessage(key_info, value=value,
                                    buffer_pool=self._buffer_pool,
                                    manager_chnl_desc=manager_chnl_desc,
                                    client_chnl=self.client_chnl)
        message.send_message()
        message.read_message()
        LOG.debug(f"SET ITEM: Stored {key} successfully in the dictionary.")

    def __delitem__(self, key):
        """DEL item in the dictionary. The request will be routed to the
        corresponding manager in the dictionary based on the provided key.

        :param key: Information about the DEL operation to the manager
        :type key: Any hashable structure
        """
        # Retrieve the manager channel descriptor for the request
        manager_chnl_desc = self._dist_dict.retrieve_channel_desc(key)

        # Construct the information needed for serving the request
        key_info = {'key': key,
                    'dict_op': DictOp.DEL_ITEM,
                    'client_puid': self.client_puid}

        # Prepare the dictionary message and send over to the manager
        message = DragonDictMessage(key_info, buffer_pool=self._buffer_pool,
                                    manager_chnl_desc=manager_chnl_desc,
                                    client_chnl=self.client_chnl)
        message.send_message()
        if(message.read_message()):
            LOG.debug(f"Deleted key: {key} successfully from the dictionary.")
        else:
            LOG.debug(f"Key: {key} does not exist. Delete {key} is not successful.")

    def __iter__(self):
        """
        Return an iterator to the list of keys in the dictionary
        """
        return iter(self.keys())

    def get(self, key: object) -> object:
        """GET item from the dictionary. The request will be routed to the
        corresponding manager in the dictionary based on the provided key.

        :param key: Information about the GET operation to the manager
        :type key: Any hashable structure
        :return: Value for the stored key in the dictionary
        :rtype: Any object that can be serialized
        """
        return self[key]

    def keys(self) -> list:
        """Retrieve the keys from the dictionary. Collects all the keys
        across the managers and returns the result as a list.
        :return: List of all the keys in the dictionary
        :rtype: list
        """
        op_info = {'dict_op': DictOp.GET_KEYS,
                   'client_puid': self.client_puid}
        keys = self._send_and_read_message_from_all_managers(op_info, [])
        return keys

    def values(self) -> list:
        """Retrieve the values from the dictionary. Collects all the values
        across the managers and returns the result as a list.
        :return: List of all the values in the dictionary
        :rtype: list
        """
        op_info = {'dict_op': DictOp.GET_VALS,
                   'client_puid': self.client_puid}
        values = self._send_and_read_message_from_all_managers(op_info, [])
        return values

    def items(self) -> list[tuple]:
        """Retrieve the items from the dictionary. Collects all the keys,values
        across the managers and returns the result as a list, with each key,value
        pair as a tuple.
        :return: List of all the key, value pairs as tuples in the dictionary
        :rtype: list
        """
        op_info = {'dict_op': DictOp.GET_ITEMS,
                   'client_puid': self.client_puid}
        items = self._send_and_read_message_from_all_managers(op_info, [])
        return items

    def __contains__(self, key):
        """Check if the key is present in the dictionary. The request will be
        routed to the corresponding manager based on the provided key.

        :param key: Information about the key to the manager
        :type key: Any hashable structure
        """
        # Retrieve the manager channel descriptor for the request
        manager_chnl_desc = self._dist_dict.retrieve_channel_desc(key)

        # Construct the information needed for checking the key
        key_info = {'key': key, 'dict_op': DictOp.CONT_ITEM,
                    'client_puid': self.client_puid}

        # Prepare the dictionary message and send over to the manager
        message = DragonDictMessage(key_info, buffer_pool=self._buffer_pool,
                                    manager_chnl_desc=manager_chnl_desc,
                                    client_chnl=self.client_chnl)
        message.send_message()

        # Wait on the client channel to receive the message from the manager
        ret_val = message.read_message()
        if not ret_val:
            LOG.debug(f"Failed: Dictionay does not contain the key: {key}.")
        return ret_val

    def pop(self, key: object) -> object:
        """Pop an item in the dictionary. The request will be routed to the
        corresponding manager in the dictionary based on the provided key. The
        key will be deleted, collecting the value from it.

        :param key: Information about the key to be popped
        :type new_key: Any hashable structure
        :return: Value of the removed key from the dictionary
        :rtype: Any object that can be serialized
        """
        # Retrieve the manager channel descriptor for the request
        manager_chnl_desc = self._dist_dict.retrieve_channel_desc(key)

        # Construct the information needed for serving the request
        key_info = {'key': key,
                    'dict_op': DictOp.POP_ITEM,
                    'client_puid': self.client_puid}

        # Prepare the dictionary message and send over to the manager
        message = DragonDictMessage(key_info, buffer_pool=self._buffer_pool,
                                    manager_chnl_desc=manager_chnl_desc,
                                    client_chnl=self.client_chnl)
        message.send_message()
        ret_val = message.read_message()
        if not ret_val:
            LOG.debug(f"Key: {key} does not exist. Pop {key} is not successful.")
            raise KeyError(f"on request for {key}")
        else:
            LOG.debug(f"Popped key: {key} successfully from the dictionary.")
            return ret_val

    def rename(self, old_key: object, new_key: object) -> None:
        """Rename an item in the dictionary. The request will be routed to the
        corresponding manager in the dictionary based on the provided key. The
        old key will be deleted, collecting the value from it. A new set operation
        will be executed on the dictionary with a new key and collected value.

        :param old_key: Information about the key to be replaced
        :type old_key: Any hashable structure
        :param new_key: Information about the new key to be added
        :type new_key: Any hashable structure
        :return: None
        :rtype None
        """
        ret_val = self.pop(old_key)
        if ret_val:
            self.__setitem__(new_key, ret_val)
        else:
            LOG.debug(f"Key: {old_key} does not exist. Rename {old_key} is not successful.")


class DragonDictMessage():
    """
    A message object created to carry the information over the Dragon channel.
    Contains the information about the operation and the (key, value) to be stored
    in the dragon dictionary.
    :param key_info: Information about the operation to be performed on the manager.
    :type key_info: dict
    :param value: Value to be stored inside the dictionary
    :type value: Any structure that can be serialized
    :param buffer_pool: Memory pool that will be used while transferring the data
    information to the remote node.
    :type buffer_pool: Dragon Memory Pool
    :param manager_chnl_desc: Manager channel descriptor that is used to be send the message
    :type manager_chnl_desc: Dragon Channel Descriptor
    :param client_chnl: Client channel used to receive the response
    :type client_chnl: Dragon channel object
    """

    def __init__(self, key_info, value=None, buffer_pool=None,
                 manager_chnl_desc=None, client_chnl=None):
        self.key_info = key_info
        self.value = value

        if buffer_pool:
            self.buffer_pool = buffer_pool
        else:
            self.buffer_pool = None

        if manager_chnl_desc:
            self.manager_chnl = Channel.attach(manager_chnl_desc.sdesc)
            self.write_adapter = Many2ManyWritingChannelFile(self.manager_chnl, buffer_pool=self.buffer_pool)
        else:
            self.manager_chnl = None
            self.write_adapter = None

        if client_chnl:
            self.read_adapter = Many2ManyReadingChannelFile(client_chnl)
        else:
            self.read_adapter = None

    def send_message(self, manager_chnl_desc: object = None) -> None:
        """Send the message to the dictionary manager over it's corresponding channel.
        Use the memory pool if the channel is on the remote node.
        :param manager_chnl_desc: Manager channel descriptor that is used to be send the message
        :type manager_chnl_desc: Dragon channel descriptor
        """
        if manager_chnl_desc:
            self.manager_chnl = Channel.attach(manager_chnl_desc.sdesc)
            self.write_adapter = Many2ManyWritingChannelFile(self.manager_chnl, buffer_pool=self.buffer_pool)

        try:
            self.write_adapter.open()
            pickle.dump(self.key_info, file=self.write_adapter, protocol=5)
            if self.value is not None:
                pickle.dump(self.value, file=self.write_adapter, protocol=5)
        except Exception as e:
            raise ConnectionError(f"Could not complete send operation: {e}")
        finally:
            self.write_adapter.close()

        if manager_chnl_desc:
            self.manager_chnl.detach()

    def read_message(self, client_chnl: object = None) -> None:
        """Read the message from the dictionary manager over the client channel.
        :param client_chnl: Client channel that is used to read the message
        :type client_chnl: Dragon channel object
        :return: Value retrieved from the dictionary for the requested message
        :rtype: Boolean or any value that is serializable
        """
        value = None

        if client_chnl:
            self.read_adapter = Many2ManyReadingChannelFile(self.client_chnl)

        try:
            self.read_adapter.open()
            buf = bytearray(4)
            self.read_adapter.readinto(memoryview(buf))
            ret_val = ReturnCode(int.from_bytes(buf, byteorder=_BYTEORDER, signed=True))

            if (ret_val == ReturnCode.FAILED):
                raise KeyError(f"on request for {self.key_info['key']}")
            elif (ret_val == ReturnCode.SUCCESS_RET_VAL):
                self.read_adapter.readinto(memoryview(buf))
                value_size = int.from_bytes(buf, byteorder=_BYTEORDER, signed=True)
                buf = bytearray(value_size)
                self.read_adapter.readinto(memoryview(buf))
                value = pickle.loads(buf)
            elif (ret_val == ReturnCode.SUCCESS_RET_TRUE):
                value = True
            elif (ret_val == ReturnCode.SUCCESS_RET_FALSE):
                value = False
        except EOFError:
            raise EOFError
        except Exception as e:
            raise ConnectionError(f"Could not complete receive operation: {e}")
        finally:
            self.read_adapter.close()

        return value

src/dragon/data/distdictionary/dragon_dict.py
DragonHPC/dragon
"""Logging setup for Dragon runtime infrastructure

This is to bring in the logging package and allow the infrastructure
to log in a consistent way.
"""

import logging
import os
import time


def setup_logging(*, basename, basedir=None, the_level=logging.INFO, force=False, add_console=False):
    """Global logging setup.

        Sets up logging with specified logging level.

        The logfile has the timestamp in the name so it's easy to disambiguate
        between multiple test runs, etc.
    """

    if basedir is None:
        basedir = os.getcwd()

    now = time.strftime('%m_%d_%H_%M_%S', time.localtime())
    fn = '{}_{}.log'.format(basename, now)
    full_fn = os.path.join(basedir, fn)
    fmt = '%(asctime)-15s %(levelname)-6s  %(name)s :: %(message)s'
    logging.basicConfig(format=fmt, filename=full_fn, level=the_level, force=force)

    if add_console:
        console = logging.StreamHandler()
        console.setLevel(logging.DEBUG)
        formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')
        console.setFormatter(formatter)
        logging.getLogger('').addHandler(console)


src/dragon/dlogging/log_setup.py
DragonHPC/dragon
"""Logging setup for Dragon runtime infrastructure"""
import logging
import os
import time
import socket
import enum
from typing import Tuple, Union
import argparse

import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.messages as dmsg
from dragon.utils import B64
from dragon.managed_memory import MemoryPool
from dragon.dlogging.logger import DragonLogger, DragonLoggingError


# This is random and exists to conform to the messaging infrastructure.
# It is discarded once logs are written out
LOGGING_TAG = 2**8

LOGGING_OUTPUT_DEVICE_STDERR = 'stderr'
LOGGING_OUTPUT_DEVICE_DRAGON_FILE = 'dragon_file'
LOGGING_OUTPUT_DEVICE_ACTOR_FILE = 'actor_file'

LOGGING_OUTPUT_DEVICES = [
    LOGGING_OUTPUT_DEVICE_STDERR,
    LOGGING_OUTPUT_DEVICE_DRAGON_FILE,
    LOGGING_OUTPUT_DEVICE_ACTOR_FILE
]

LOGGING_LEVEL_NONE = 'NONE'
LOGGING_LEVEL_NAMES_MAPPING = {
    logging.getLevelName(level): level for level in
        [logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR, logging.CRITICAL]
}

LOGGING_STDERR_MAX_LOG_LEVEL = logging.INFO
LOGGING_DRAGON_FILE_MAX_LOG_LEVEL = logging.INFO

LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING = {
    f'DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_STDERR.upper()}': LOGGING_LEVEL_NONE,
    f'DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_DRAGON_FILE.upper()}': LOGGING_LEVEL_NONE,
    f'DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_ACTOR_FILE.upper()}': LOGGING_LEVEL_NONE
}

# Default formatting for logging
default_single_fmt = '%(asctime)-15s %(levelname)-8s %(name)s :: %(message)s'
default_FE_fmt = '%(asctime)-15s %(levelname)-8s %(service)s:%(name)s (%(hostname)s) :: %(message)s'
default_services_fmt = '%(time)-15s %(levelname)-8s %(service)s:%(name)s (%(hostname)s) :: %(message)s'


class LoggingValue(argparse.Action):
    def __init__(self,
                 option_strings,
                 dest,
                 nargs=None,
                 const=None,
                 default=None,
                 type=None,
                 choices=None,
                 required=False,
                 help=None,
                 metavar=None):

        if default == dict():
            default = LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING

        if choices == None:
            choices = []
            choices.append(LOGGING_LEVEL_NONE)
            choices.extend(LOGGING_LEVEL_NAMES_MAPPING.keys())
            for device in LOGGING_OUTPUT_DEVICES:
                choices.append(f'{device}={LOGGING_LEVEL_NONE}')
                for level in LOGGING_LEVEL_NAMES_MAPPING.keys():
                    choices.append(f'{device}={level}')

        super(LoggingValue, self).__init__(
            option_strings=option_strings,
            dest=dest,
            nargs=nargs,
            const=const,
            default=default,
            type=type,
            choices=choices,
            required=required,
            help=help,
            metavar=metavar)

    def __call__(self , parser, namespace, values, option_string = None):
        level_names = LOGGING_LEVEL_NAMES_MAPPING.keys()

        log_device_level_map = getattr(namespace, self.dest)

        for value in values:
            if value.upper() == LOGGING_LEVEL_NONE:
                setattr(namespace, self.dest, LOGGING_DEFAULT_DEVICE_LEVEL_MAPPING)

            elif value.upper() in level_names:
                level_name = value.upper()
                log_device_level_map[f'DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_STDERR.upper()}'] = level_name
                log_device_level_map[f'DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_DRAGON_FILE.upper()}'] = level_name
                log_device_level_map[f'DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_ACTOR_FILE.upper()}'] = level_name

            elif '=' in value:
                k, v = value.split('=')
                output_device = k.lower()
                level_name = v.upper()

                if output_device not in LOGGING_OUTPUT_DEVICES:
                    raise ValueError(f"Unsupported log output device '{k}' in parameter '{value}'")

                if level_name not in level_names:
                    raise ValueError(f"Unsupported log level '{v}' in parameter '{value}'")

                log_device_level_map[f'DRAGON_LOG_DEVICE_{output_device.upper()}'] = level_name

            else:
                raise ValueError(f"Unsupported parameter {value}.")


class DragonLoggingServices(str, enum.Enum):
    """Enumerator for logging names of backend services

    Attributes:
        :LA_FE: 'LA_FE'
        :LA_BE: 'LA_BE'
        :GS: 'GS'
        :TA: 'TA'
        :ON: 'ON'
        :LS: 'LS'
        :TEST: 'TEST'
        :PERF: 'PERF'
    """
    LA_FE = 'LA_FE'
    LA_BE = 'LA_BE'
    GS = 'GS'
    TA = 'TA'
    ON = "ON"
    LS = 'LS'
    TEST = 'TEST'
    PERF = 'PERF'

    def __str__(self):
        return str(self.value)


def next_tag():
    global LOGGING_TAG
    tmp = LOGGING_TAG
    LOGGING_TAG += 1
    return tmp


class DragonFEFilter(logging.Filter):

    def __init__(self) -> None:
        """Create logging filter to add record attributes to FE logger

        Adds hostname, IP address, port, and service name to all FE logging
        entries. These are needed to satisfy the ``dragon.dlogging.util.default_FE_fmt``
        format string
        """

        super().__init__()

    def filter(self, record: logging.LogRecord) -> bool:
        """Filter out BE messages and let through all FE messages while defining record attributes

        Args:
            record (logging.LogRecord): record that will be populated

        Returns:
            bool: Log message if True. Don't if False
        """

        # time is only defined for our backend services. If the attr exists, we will
        # filter out the message as another handler will log the backend logs
        if hasattr(record, "time"):
            return False
        else:
            record.hostname = socket.gethostname()
            record.ip_address = socket.gethostbyname(record.hostname)
            record.port = '      '
            record.service = DragonLoggingServices.LA_FE
        return True


class DragonLoggingHandler(logging.StreamHandler):

    def __init__(self, serialized_log_descr: B64, mpool: MemoryPool = None, hostname: str = None, ip_address: str = None,
                 port: str = None, service: Union[DragonLoggingServices, str] = None) -> None:
        """Logging handler to enable python logging to be sent to the Dragon Logging channel on BE

        Ultimately a `logging.StreamHandler` that sends stream output to ``/dev/null`` and also
        emits log records into a  :ref:`dragon.infrastructure.messages.LoggingMsg <loggingmsg>` object that
        is serialized into a :class:`dragon.dlogging.logger.DragonLogger`

        :param serialized_log_descr: serialized descriptor from
            :class:`dragon.dlogging.logger.DragonLogger`
        :type serialized_log_descr: B64
        :param mpool: Memory pool to use for attaching to logging channel, defaults to None
        :type mpool: MemoryPool, optional
        :param hostname: hostname of logging service, defaults to None
        :type hostname: str, optional
        :param ip_address: IP address of logging service, defaults to None
        :type ip_address: str, optional
        :param port: Port of logging service, defaults to None
        :type port: str, optional
        :param service: Name of logging service. Acceptable names are defined
            in :class:`DragonLoggingServices`
        :type service: Union[DragonLoggingServices, str], optional
        """
        with open(os.devnull, 'w') as f:
            super().__init__(f)
        self._dlog = DragonLogger.attach(serialized_log_descr.decode(), mpool=mpool)
        self.hostname = hostname
        self.ip_address = ip_address
        self.port = port
        self.service = service

    def emit(self, record: logging.LogRecord) -> None:
        """Emit a dragon.insfrauction.messages.LoggingMsg message

        Args:
            record (logging.LogRecord): Record populated by classes format method

        Raises:
            err: DragonLoggingError if message could not be emitted into logging channel
        """
        try:
            self.format(record)
            msg = dmsg.LoggingMsg(tag=next_tag(),
                                  name=record.name,
                                  msg=record.msg,
                                  time=record.asctime,
                                  func=record.funcName,
                                  hostname=self.hostname,
                                  ip_address=self.ip_address,
                                  port=self.port,
                                  service=self.service,
                                  level=record.levelno)
            self._dlog.put(msg.serialize(), record.levelno)
        except (Exception, DragonLoggingError) as err:
            if isinstance(err, DragonLoggingError):
                raise err
            else:
                self.handleError(record)


def _clear_root_log_handlers():
    log = logging.getLogger()
    for handler in log.handlers:
        handler.close()
        log.removeHandler(handler)


def _setup_actor_logging(service: Union[DragonLoggingServices, str], level: int, fname: str = None) -> Tuple[int, str]:
    """Get the logging level and setup file logging if in DEBUG mode

    Args:
        service (str, optional): logging service, options defined in :class:`DragonLoggingServices`. Defaults to 'dragon'.

        fname (str, optional): Filename to log to if debug logging is enabled. Defaults to None.

    Returns:
        int: logging level

        str: logging filename if created or None if not
    """
    if service is None or service is DragonLoggingServices.LA_FE:
        service_name = 'dragon'
    else:
        service_name = service

    full_fn = None
    if service is not DragonLoggingServices.LA_FE:
        try:
            basedir = os.environ.get('DRAGON_LA_LOG_DIR')
        except AttributeError:
            basedir = None
        if fname is None:
            full_fn = _setup_file_logging(basename=f'{service_name}_{socket.gethostname()}', basedir=basedir, level=level)
        else:
            full_fn = _setup_file_logging(fname=os.path.basename(fname), basedir=basedir, level=level)

    return full_fn


def _configure_console_logging(level, fmt=None):
    '''Add logging to stderr handler'''

    console = logging.StreamHandler()
    console.setLevel(level)

    if fmt is None:
        formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')
    else:
        formatter = logging.Formatter(fmt)

    console.setFormatter(formatter)
    logging.getLogger().addHandler(console)


def setup_logging(*, basename: str = None, basedir: str = None, level: int = logging.INFO, force: bool = False,
                  add_console: bool = False):
    """Original dragon logging function to configure root logger.

    Args:
        basename (str, optional): Logging file basename. Date and time will be appended. Defaults to 'dragon'. Defaults to None.

        basedir (str, optional): Directory to save log file to os.getcwd(). Defaults to None.

        level (int, optional): Logging level. Defaults to logging.INFO.

        force (bool, optional): Remove any handlers attached to the root logger. Defaults to False.

        add_console (bool, optional): Enable loggng to stderr. Defaults to False.
    """
    if basedir is None:
        basedir = os.getcwd()
    else:
        # Make sure the path exists. If not, just put to /tmp,
        # so we have a defined place for it to go and look for
        # it
        if not os.path.exists(basedir):
            basedir = '/tmp'

    fmt = default_single_fmt

    now = time.strftime('%m_%d_%H_%M_%S', time.localtime())
    fn = '{}_{}.log'.format(basename, now)
    full_fn = os.path.join(basedir, fn)

    logging.basicConfig(format=fmt, filename=full_fn, level=level, force=force)

    if add_console:
        _configure_console_logging(level=level)


def _setup_file_logging(basename: str = None, fname: str = None, basedir: str = None, level: int = logging.INFO,
                        force: bool = False, fmt: str = None) -> str:
    """Configure file-based logging

    Sets up the root file logging with specified logging level.
    The logfile has the timestamp in the name so it's easy to disambiguate
    between multiple test runs, etc.

    Args:
        basename (str, optional): Logging file basename. Date and time will be appended. Defaults to 'dragon'

        fname (str, optional): Full path and file name of log to create. Overrides fname and basedir. Defaults to None.

        basedir (str, optional): Directory to save log file to. Defaults to pwd

        level (int, optional): Logging level. Defaults to logging.INFO.

        force (bool, optional): Remove all existing handlers from root logger. Defaults to False.

        fmt (str, optional): Formatting string for file log messages. Defaults to None.

    Returns:
        str: file and path of log file
    """
    if basedir is None:
        basedir = os.getcwd()
    else:
        # Make sure the path exists. If not, just put to /tmp,
        # so we have a defined place for it to go and look for
        # it
        if not os.path.exists(basedir):
            basedir = '/tmp'

    if fmt is None:
        fmt = default_single_fmt

    if fname is not None:
        full_fn = os.path.join(basedir, fname)
    else:
        now = time.strftime('%m_%d_%H_%M_%S', time.localtime())
        fn = '{}_{}.log'.format(basename, now)
        full_fn = os.path.join(basedir, fn)

    file_handler = logging.FileHandler(filename=full_fn, mode='a')
    formatter = logging.Formatter(fmt=fmt)
    file_handler.setFormatter(formatter)
    file_handler.setLevel(level)
    logging.getLogger().addHandler(file_handler)
    logging.getLogger().setLevel(level=level)

    return full_fn


def setup_FE_logging(log_device_level_map: dict = {}, basename: str = None, basedir: str = None, force: bool = False,
                     add_console: bool = False, test_BE_dir: str = None, multi_node_mode: bool = False) -> Tuple[int, str]:
    """Set-up launcher frontend logging infrastructure. Will log to file.

    Configures the individual service loggers the FE will use to log all messages
    sent to it from the backend services of the MRNet overlay and creates the FileHandler
    they will all write to

    Args:
        basename (str, optional): Logging file basename. Date and time will be appended. Defaults to 'dragon'.

        basedir (str, optional): Directory to save log file to. Defaults to pwd

        force (bool, optional): Remove all existing handlers from root logger. Defaults to False.

        add_console (bool, optional): Turn on logging to stderr. Defaults to False.

        test_BE_dir (str, optional): Directory to save backend log files to. Used for testing. Defaults to None.

        multi_node_mode (bool, optional): Enable multi-node filtering and formatting.
    Returns:
        int: logging level

        str: full path and filename of logging file
    """

    # Make sure the environment is updated with the directory to log to:
    if test_BE_dir is not None:
        os.environ.update({'DRAGON_LA_LOG_DIR': test_BE_dir})
    elif os.environ.get('DRAGON_LA_LOG_DIR') is None:
        os.environ.update({'DRAGON_LA_LOG_DIR': os.getcwd()})
    for log_device, log_level in log_device_level_map.items():
        os.environ.update({log_device: log_level})

    log_to_stderr, stderr_log_level = _get_dragon_log_device_level(log_device_level_map, LOGGING_OUTPUT_DEVICE_STDERR)
    log_to_dragon_file, dragon_file_log_level = _get_dragon_log_device_level(log_device_level_map, LOGGING_OUTPUT_DEVICE_DRAGON_FILE)

    _clear_root_log_handlers()

    logger = logging.getLogger()

    if any([log_to_stderr, log_to_dragon_file]):

        root_log_level = max(stderr_log_level, dragon_file_log_level)
        logger.setLevel(root_log_level)

        if log_to_stderr:
            _configure_console_logging(level=stderr_log_level, fmt=default_single_fmt)

        if log_to_dragon_file:
            file_formatter = default_FE_fmt if multi_node_mode else default_single_fmt
            full_fn = _setup_file_logging(basename=basename, basedir=basedir, level=dragon_file_log_level, force=force, fmt=file_formatter)

        if multi_node_mode:
            # Add a filter to all the root handlers to only log on the root if it's a FE service
            for handler in logging.getLogger().handlers:
                handler.addFilter(DragonFEFilter())
    else:
        logger.addHandler(logging.NullHandler())


def close_FE_logging():
    """Close individual service loggers used for logging on frontend"""
    for service in list(DragonLoggingServices):
        log = logging.getLogger(service)
        for handle in log.handlers[:]:
            handle.close()
            log.removeHandler(handle)


def _get_logging_mpool(node_index: int):
    '''Create memory pool for logging infrastructure

    Args:
        node_index (int): node ID of caller
    '''
    _user = os.environ.get('USER', str(os.getuid()))
    lps = int(dfacts.DEFAULT_SINGLE_DEF_SEG_SZ)
    lpn = f'{_user}_{os.getpid()}_{dfacts.LOGGING_POOL_SUFFIX}'
    lp_muid = dfacts.logging_pool_muid_from_index(node_index)
    logging_mpool = MemoryPool(lps, lpn, lp_muid, None)
    return logging_mpool


def setup_dragon_logging(node_index: int) -> DragonLogger:
    """Create a :class:`dragon.dlogging.logger.DragonLogger` instance

    Args:
        node_index (int): node index logger is being created on

    Returns:
        DragonLogger: created DragonLogger instance
    """
    logging_mpool = _get_logging_mpool(node_index)
    dragon_logger = DragonLogger(logging_mpool)

    return dragon_logger


def detach_from_dragon_handler(service: Union[DragonLoggingServices, str]):
    """Detach logger from dragon handler

    This only needs to be called for processes that will outlive the
    MRNet backend service, which "owns" the logging memory pool
    and relevant backend infrastructure

    Args:
        service (str): logging service, options defined in :class:`DragonLoggingServices`
    """
    log = logging.getLogger(service)
    root_log = logging.getLogger()

    for handle in log.handlers[:]:
        if isinstance(handle, DragonLoggingHandler):
            log.debug("removing dragonlogginghandler")
            log.removeHandler(handle)

    for handle in root_log.handlers[:]:
        if isinstance(handle, DragonLoggingHandler):
            log.debug("removing dragonlogginghandler from root")
            root_log.removeHandler(handle)


def _get_dragon_log_device_level(env_map: dict(), device_name: str) -> Tuple[bool, int]:
    log_level_value =  0
    got_dragon_log_level = False
    dragon_log_device_env = f'DRAGON_LOG_DEVICE_{device_name.upper()}'

    if dragon_log_device_env in env_map:
        try:
            log_level_name = env_map[dragon_log_device_env]
            log_level_value = LOGGING_LEVEL_NAMES_MAPPING[log_level_name]
            got_dragon_log_level = log_level_value != LOGGING_LEVEL_NONE
        except KeyError:
            # Unsupported level was set, don't enable logging
            pass

    return got_dragon_log_level, log_level_value


def setup_BE_logging(service: Union[DragonLoggingServices, str], logger_sdesc: B64 = None,
                     mpool: MemoryPool = None,
                     hostname: str = None, ip_address: str = None, port: str = None,
                     fname: str = None) -> Tuple[int, str]:
    """Configure backend logging for requested backend service

    Create a handler using the :class:`dragon.dlogging.logger.DragonLogger` serialized
    descriptor and add it to the given service's logger
    :param service: logging service, options defined in :class:`DragonLoggingServices`
    :type service: Union[DragonLoggingServices, str]
    :param logger_sdesc: serialized logging descriptor, defaults to None
    :type logger_sdesc: B64, optional
    :param mpool: memory pool to use for attaching to channel, defaults to None
    :type mpool: MemoryPool, optional
    :param hostname:  Hostname of logging service, defaults to None
    :type hostname: str, optional
    :param ip_address: IP address of logging service, defaults to None
    :type ip_address: str, optional
    :param port: Port of logging service, defaults to None
    :type port: str, optional
    :param fname: Filename for file logging, defaults to None
    :type fname: str, optional
    :return: logging level, filename logging will be saved to
    :rtype: Tuple[int, str]
    """
    full_fn = ''
    log_level = 0

    # Make sure there isn't a duplicate handler already attached to this root instance
    _clear_root_log_handlers()

    log_to_stderr, stderr_log_level = _get_dragon_log_device_level(os.environ, LOGGING_OUTPUT_DEVICE_STDERR)
    log_to_dragon_file, dragon_file_log_level = _get_dragon_log_device_level(os.environ, LOGGING_OUTPUT_DEVICE_DRAGON_FILE)
    log_to_actor_file, actor_file_log_level = _get_dragon_log_device_level(os.environ, LOGGING_OUTPUT_DEVICE_ACTOR_FILE)

    logger = logging.getLogger()

    if any([log_to_stderr, log_to_dragon_file, log_to_actor_file]):

        # determine the maximum log level that we want to the logger to handle
        if log_to_stderr:
            log_level = stderr_log_level
        if log_to_dragon_file:
            log_level = max(log_level, dragon_file_log_level)
        if log_to_actor_file:
            log_level = max(log_level, actor_file_log_level)

        # Set the loglevel on the logger
        logger.setLevel(log_level)

        # determine the maximum log level that we want to send over MRNet
        if all([log_to_stderr, log_to_dragon_file]):
            front_end_log_level = max(stderr_log_level, dragon_file_log_level)
        elif log_to_stderr:
            front_end_log_level = stderr_log_level
        else:
            front_end_log_level = dragon_file_log_level

        # If we're logging to either/both stderr or the combined dragon file,
        # we then need to send appropriate log messages to the FrontEnd via
        # the DragonLoggingHandler.
        if logger_sdesc and any([log_to_stderr, log_to_dragon_file]):

            # Determine our hostname if we weren't told it
            if hostname is None:
                hostname = socket.gethostname()

            # Create the Dragon Logging Handler
            handler = DragonLoggingHandler(logger_sdesc,
                                           mpool=mpool,
                                           hostname=hostname,
                                           ip_address=ip_address,
                                           port=port,
                                           service=service)

            # Give it a formatter with asctime so time gets populated in the log record
            handler.setFormatter(logging.Formatter(fmt=default_single_fmt + '%(funcName)s'))

            # The intent of this next block is to prevent sending DEBUG or higher log
            # messages over the MRNet connection if those log messages will be logged
            # to the actor file.
            if log_to_actor_file and (front_end_log_level <= actor_file_log_level):
                if logging.INFO > front_end_log_level:
                    front_end_log_level = logging.INFO

            # Set the loglevel on the handler
            handler.setLevel(front_end_log_level)

            # FIXME: This needs fixing.
            # The logger currently fills up the channel because no back pressure is applied
            # and instead of handling that internally, returns a CHANNEL_FULL error which
            # causes some tests to fail. It should instead, either block on sending or
            # discard log messages when full. We should have a way of configuring which
            # behavior we want.

            # Add the handler to the logger object
            logger.addHandler(handler)

        if log_to_actor_file:
            full_fn = _setup_actor_logging(service=service, level=actor_file_log_level, fname=fname)

    else:
        logger.addHandler(logging.NullHandler())

    return log_level, full_fn

src/dragon/dlogging/util.py
DragonHPC/dragon
"""Setup and global objects for Dragon's GS client API.

This module includes initialization code  for
getting handles to and from the local shepherd and global services
through which API generated messages are sent.

Python thread locks are provided here so that threaded Python clients
can have more than one transaction outstanding to GS at once.
"""

import logging
import sys
import threading
import gc
import os

import dragon.channels as dch
import dragon.globalservices.process as dproc
import dragon.infrastructure.process_desc as dpdesc
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.parameters as dp
import dragon.infrastructure.debug_support as ddbg
import dragon.infrastructure.facts as dfacts
import dragon.utils as du

LOG = logging.getLogger('api_setup')

_GS_API_LOCK = threading.RLock()
_GS_SEND_LOCK = threading.RLock()
_GS_RECV_LOCK = threading.RLock()

_MSG_TAG_CTR = 0


def next_tag():
    global _MSG_TAG_CTR
    with _GS_API_LOCK:
        tmp = _MSG_TAG_CTR
        _MSG_TAG_CTR += 1
    return tmp


_ARG_PAYLOAD = None

_GS_INPUT_CHANNEL = None
_GS_RETURN_CHANNEL = None
_LOCAL_SHEP_INPUT_CHANNEL = None
_LOCAL_SHEP_RETURN_CHANNEL = None

_GS_INPUT = None
_GS_RETURN = None
_SHEP_INPUT = None
_SHEP_RETURN = None

_GS_RETURN_CUID = None
_SHEP_RETURN_CUID = None


def get_gs_ret_cuid():
    if not _INFRASTRUCTURE_CONNECTED:
        connect_to_infrastructure()

    return _GS_RETURN_CUID


def load_launch_parameter(name):
    param = getattr(dp.this_process, name.lower())
    assert param, f'Launch parameter not initialized: {name}'
    return du.B64.str_to_bytes(param)


def _connect_gs_input():
    global _GS_INPUT_CHANNEL

    channel_descriptor = load_launch_parameter(dfacts.GS_CD)
    _GS_INPUT_CHANNEL = dch.Channel.attach(channel_descriptor)
    # size chosen to be safely above ARG_IMMEDIATE_LIMIT
    # TODO PE-38745
    return dconn.Connection(outbound_initializer=_GS_INPUT_CHANNEL,
                            options=dconn.ConnectionOptions(min_block_size=2 ** 21, large_block_size=2**22, huge_block_size=2**23), policy=dp.POLICY_INFRASTRUCTURE)


def _connect_gs_return():
    global _GS_RETURN_CHANNEL
    global _GS_RETURN_CUID

    channel_descriptor = load_launch_parameter(dfacts.GS_RET_CD)
    _GS_RETURN_CHANNEL = dch.Channel.attach(channel_descriptor)
    _GS_RETURN_CUID = _GS_RETURN_CHANNEL.cuid
    return dconn.Connection(inbound_initializer=_GS_RETURN_CHANNEL, policy=dp.POLICY_INFRASTRUCTURE)


def _connect_shep_input():
    global _LOCAL_SHEP_INPUT_CHANNEL

    channel_descriptor = load_launch_parameter(dfacts.LOCAL_SHEP_CD)
    _LOCAL_SHEP_INPUT_CHANNEL = dch.Channel.attach(channel_descriptor)
    return dconn.Connection(outbound_initializer=_LOCAL_SHEP_INPUT_CHANNEL, policy=dp.POLICY_INFRASTRUCTURE)


def _connect_shep_return():
    global _LOCAL_SHEP_RETURN_CHANNEL
    global _SHEP_RETURN_CUID

    # currently we have no use for the shep return cd (although we might in the future)
    # it isn't currently being constructed by GS so here we will see if it
    # is present and return None if it isn't.

    try:
        channel_descriptor = load_launch_parameter(dfacts.SHEP_RET_CD)
    except AssertionError:
        return None
    _LOCAL_SHEP_RETURN_CHANNEL = dch.Channel.attach(channel_descriptor)
    _SHEP_RETURN_CUID = _LOCAL_SHEP_RETURN_CHANNEL.cuid
    return dconn.Connection(inbound_initializer=_LOCAL_SHEP_RETURN_CHANNEL, policy=dp.POLICY_INFRASTRUCTURE)


def _close_connections():
    with _GS_API_LOCK:
        _GS_INPUT.close()
        _GS_RETURN.close()
        _SHEP_INPUT.close()
        _SHEP_RETURN.close()


def _detach_infrastructure():
    with _GS_API_LOCK:
        if _GS_INPUT_CHANNEL is not None:
            _GS_INPUT_CHANNEL.detach()

        if _GS_RETURN_CHANNEL is not None:
            _GS_RETURN_CHANNEL.detach()

        if _LOCAL_SHEP_INPUT_CHANNEL is not None:
            _LOCAL_SHEP_INPUT_CHANNEL.detach()

        if _LOCAL_SHEP_RETURN_CHANNEL is not None:
            _LOCAL_SHEP_RETURN_CHANNEL.detach()


# global because of threading.

# contains condition variables corresponding to threads awaiting
# a response from global services.  Key = tag, value = cv
_WAKEUPS = dict()

# contains messages gotten back from global services, to be picked up by
# waiting threads.
_RESULTS = dict()

_INFRASTRUCTURE_CONNECTED = False

def gs_request(req_msg, *, expecting_response=True):
    """Posts a message to GS and gets the response in a thread safe way

    This is an attempt to make global services transactions blocking
    and thread safe while avoiding needless head of queue blocking - we
    don't want there to be only one transaction outstanding at a time
    because it might take a long time to finish and block other threads'
    transactions in the meantime, but without starting a separate dedicated
    thread simply to deal and service gs transactions.  The implementation
    here tries to solve this without involving a new thread at all, but
    that might not be the best way to do it; a separate service thread
    would probably be simpler.

    The overhead of doing this stuff if only one thread is ever involved
    means acquiring a couple locks that have no contention on them
    and constructing an event variable which means another lock.


    Arguments:
        :param req_msg: request message to send.
        :param expecting_response: Bool default True, is a response message expected?

    Returns:
        the reply to the request message
    """

    global _WAKEUPS
    global _RESULTS

    with _GS_API_LOCK:
        if not _INFRASTRUCTURE_CONNECTED:
            connect_to_infrastructure()

    # PE-44998, do not allow the GC to run during the life of this function as it may result in a thread re-entering
    # this function, which can lead to a hang. GC cleanup of Dragon-native objects for multiprocessing
    # include communication with GS, and that is where the re-entering can happen.
    gc.disable()

    # Serialize request message
    req_msg_bytes = req_msg.serialize()

    # Ensure wake-up event is available prior to sending
    if expecting_response:
        ready = _WAKEUPS[req_msg.tag] = threading.Event()

    # Send request
    with _GS_SEND_LOCK:
        try:
            _GS_INPUT.send(req_msg_bytes)
        except Exception as e:
            gc.enable()
            raise e

    if not expecting_response:
        gc.enable()
        return

    # Wait for response
    while req_msg.tag not in _RESULTS:
        if _GS_RECV_LOCK.acquire(blocking=False):
            # Read responses until expected response is received
            try:
                while req_msg.tag not in _RESULTS:
                    # Receive and parse response
                    resp = dmsg.parse(_GS_RETURN.recv())

                    try:
                        wakeup = _WAKEUPS.pop(resp.ref)
                    except KeyError:
                        logging.warning(f"Received unexpected message {resp.ref}")
                    else:
                        # Save response message
                        _RESULTS[resp.ref] = resp
                        # Alert owner response is available
                        wakeup.set()

                # Wake up last sender to ensure someone takes over receiving
                try:
                    # Avoid iteration for thread-safety using CPython
                    _k, _v = _WAKEUPS.popitem()
                except KeyError:
                    # No one else is waiting for a response
                    pass
                else:
                    # Re-add the event to _WAKEUPS and set it
                    _WAKEUPS[_k] = _v
                    _v.set()

                # I suppose we could return here too
                gc.enable()
                return _RESULTS.pop(resp.ref)
            except Exception as e:
                gc.enable()
                raise e
            finally:
                _GS_RECV_LOCK.release()

        # Wait for response or to be woken up to take over receiving.
        if ready.wait():
            # Reset event in case response is not ready, i.e., we were woken up
            # but not alerted
            ready.clear()

    gc.enable()
    return _RESULTS.pop(req_msg.tag)


def test_connection_override(test_gs_input=None, test_gs_return=None,
                             test_gs_return_cuid=None,
                             test_shep_input=None, test_shep_return=None,
                             test_shep_return_cuid=None):
    global _GS_INPUT
    global _GS_RETURN
    global _SHEP_INPUT
    global _SHEP_RETURN
    global _INFRASTRUCTURE_CONNECTED
    global _GS_RETURN_CUID
    global _SHEP_RETURN_CUID

    LOG.debug('dragon connection override')

    with _GS_API_LOCK:

        _INFRASTRUCTURE_CONNECTED = True

        if test_gs_input is not None:
            _GS_INPUT = test_gs_input
        else:
            _GS_INPUT = _connect_gs_input()

        if test_gs_return is not None:
            _GS_RETURN = test_gs_return
            _GS_RETURN_CUID = test_gs_return_cuid
        else:
            _GS_RETURN = _connect_gs_return()

        if test_shep_input is not None:
            _SHEP_INPUT = test_shep_input
        else:
            _SHEP_INPUT = _connect_shep_input()

        if test_shep_return is not None:
            _SHEP_RETURN = test_shep_return
            _SHEP_RETURN_CUID = test_shep_return_cuid
        else:
            _SHEP_RETURN = _connect_shep_return()


def connect_to_infrastructure():
    global _GS_INPUT
    global _GS_RETURN
    global _SHEP_INPUT
    global _SHEP_RETURN
    global _INFRASTRUCTURE_CONNECTED

    # Here we register the gateway channels
    # provided to this process as environment
    # variables. This is done first to enable
    # off-node communication from this process
    # as soon as possible.


    with _GS_API_LOCK:
        if _INFRASTRUCTURE_CONNECTED:
            return

        LOG.info(f'We are registering gateways for this process. {dp.this_process.num_gw_channels_per_node=}')
        dch.register_gateways_from_env()

        LOG.info(f'connecting to infrastructure from {os.getpid()}')

        _GS_INPUT = _connect_gs_input()
        _GS_RETURN = _connect_gs_return()
        _SHEP_INPUT = _connect_shep_input()
        _SHEP_RETURN = _connect_shep_return()
        _INFRASTRUCTURE_CONNECTED = True

        LOG.debug('waiting for handshake')

        global _ARG_PAYLOAD

        # recv something from gs_return
        handshake = dmsg.parse(_GS_RETURN.recv())
        LOG.debug(f'Got response {handshake}')
        assert isinstance(handshake, dmsg.GSPingProc)

        if handshake.mode == dpdesc.ArgMode.NONE:
            _ARG_PAYLOAD = None
        elif handshake.mode == dpdesc.ArgMode.PYTHON_IMMEDIATE:
            _ARG_PAYLOAD = dproc.get_argdata_from_immediate_handshake(handshake.argdata)
        elif handshake.mode == dpdesc.ArgMode.PYTHON_CHANNEL:
            # receive as bytes through the gs ret connection.
            # This one lump of data comes from the managed process
            # that caused this process to be started and not GS.
            # There is no race with what GS might do with the channel
            # because this process can't send any responses
            # until this message is received
            _ARG_PAYLOAD = _GS_RETURN.recv_bytes()
        else:
            raise NotImplementedError('close case')

        LOG.debug('got handshake')

        sys.breakpointhook = ddbg.dragon_debug_hook
        LOG.info('debug entry hooked')

        # Would like to do this, but the bare except:
        # on line 327 of lib/python3.9/multiprocessing/process.py prevents
        # any variation in handling.      So the question arises:
        # if you want this, you have to do something about _bootstrap, and if
        # you are doing something about _bootstrap, why not everything?
        # sys.excepthook = ddbg.dragon_exception_hook
        # LOG.info('exception entry hooked')
src/dragon/globalservices/api_setup.py
DragonHPC/dragon
"""Global Services API for the Channels primary resource component.
"""
import logging
from math import inf
import time

import dragon.globalservices.api_setup as das

import dragon.infrastructure.messages as dmsg

from dragon.infrastructure.parameters import this_process
from dragon.infrastructure.channel_desc import ChannelDescriptor

log = logging.getLogger('channel_api')


class ChannelError(Exception):
    pass


def create(m_uid, user_name='', options=None, soft=False, puid_override=None):
    """Asks Global Services to create a new channel

        :param m_uid: m_uid of pool to create the channel in
        :param user_name: Requested user specified reference name
        :param options: ChannelOptions object, what options to apply to creation
        :param soft: Default False.
                     If channel already exists with given name, do not create and
                     return descriptor instead.
        :param puid_override: Only needed when the channel is being created on behalf
                     of some other process. Normally this is not needed.
        :return: ChannelDescriptor object
    """

    if options is None:
        options = {}

    if soft and not user_name:
        raise ChannelError('soft create requires a user supplied channel name')

    puid = puid_override

    if puid is None:
        puid = this_process.my_puid

    req_msg = dmsg.GSChannelCreate(tag=das.next_tag(), p_uid=puid,
                                   r_c_uid=das.get_gs_ret_cuid(), m_uid=m_uid,
                                   user_name=user_name, options=options)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSChannelCreateResponse)

    ec = dmsg.GSChannelCreateResponse.Errors

    if ec.SUCCESS == reply_msg.err:
        the_desc = reply_msg.desc
    else:
        if soft and ec.ALREADY == reply_msg.err:
            the_desc = reply_msg.desc
        else:
            raise ChannelError(f'channel create {req_msg} failed: {reply_msg!r}')

    return the_desc


def get_list():
    """Asks Global Services for a list of the c_uids of all currently active channels

        TODO: add some options to this, and to the message itself, to have a finer
        control over which channels you get back.  Deferring until we have a use case,
        but one that springs to mind is channels in a particular pool by m_uid.

        :return: list of the  c_uids of all currently existing channels
        """
    req_msg = dmsg.GSChannelList(tag=das.next_tag(), p_uid=this_process.my_puid,
                                 r_c_uid=das.get_gs_ret_cuid())

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSChannelListResponse)

    return reply_msg.clist


def query(identifier, *, inc_refcnt=False):
    """Asks Global Services for the ChannelDescriptor by c_uid or name

    Note you can only query currently existing channels

    :param identifier: string indicating channel name or integer indicating a c_uid
    :param inc_refcnt: bool indicating whether this query is also to inc the refcnt on the channel
    :return: ChannelDescriptor object corresponding to specified channel
    :raises: ChannelError if there is no such channel
    """
    if isinstance(identifier, str):
        req_msg = dmsg.GSChannelQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
                                      r_c_uid=das.get_gs_ret_cuid(),
                                      inc_refcnt=inc_refcnt, user_name=identifier)
    else:
        req_msg = dmsg.GSChannelQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
                                      r_c_uid=das.get_gs_ret_cuid(), inc_refcnt=inc_refcnt,
                                      c_uid=int(identifier))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSChannelQueryResponse)

    if dmsg.GSChannelQueryResponse.Errors.SUCCESS == reply_msg.err:
        the_desc = reply_msg.desc
    else:
        raise ChannelError(f'channel query {req_msg} failed: {reply_msg.err_info}')

    return the_desc


def wait_for(identifier, predicate, timeout=None, tries=1, interval=0.1):
    """Wait for predicate to return True. predicate must be a callable that
    takes a ChannelDescriptor instance. A timeout may be provided giving the
    maximum time to wait after a minimum number of tries for predicate to return
    True are completed.

    :param identifier: String indicating channel name or integer indicating a c_uid
    :param predicate: Callable that takes a ChannelDescriptor instance and returns a boolean
    :param timeout: Number of seconds before timing out (default: infinite)
    :param tries: Minimum number of query attempts before timing out (default: 1)
    :param interval: Number of seconds to sleep in between query attempts (default: 0.1)
    :return: ChannelDescriptor object corresponding to the specified channel
    :raises: TypeError if predicate is not callable
    :raises: ChannelError if there is no such channel
    :raises: TimeoutError if the number of tries is exhausted and the timeout expires
    """
    if not callable(predicate):
        raise TypeError('predicate is not callable')
    if timeout is None:
        timeout = inf
    end = time.monotonic() + timeout
    while time.monotonic() < end or tries > 0:
        desc = query(identifier)
        if predicate(desc):
            return desc
        if tries > 0:
            tries -= 1
        time.sleep(interval)
    raise TimeoutError(f'waiting for channel {identifier}')


def wait(identifier, state=ChannelDescriptor.State.DEAD, **kwds):
    """Wait for channel to be in a specified state. Other keyword arguments are
    passed through to wait_for().

    :param identifier: String indicating channel name or integer indicating a c_uid
    :param state: Desired ChannelDescriptor.State (default: DEAD)
    :return: ChannelDescriptor object corresponding to the specified channel
    :raises: ChannelError if there is no such channel
    :raises: TimeoutError if the number of tries is exhausted and the timeout expires
    """
    return wait_for(identifier, lambda desc: desc.state is state, **kwds)


def join(identifier, timeout=None):
    """Waits for a channel with a particular name to get made.

        :param identifier: string indicating channel name
        :param timeout: Time in seconds to timeout.  if absent, block indefinitely.
        :return: ChannelDescriptor for the channel specified
        :raises: TimeoutError if the timeout expires
    """

    assert isinstance(identifier, str)
    if timeout is None:
        timeout = -1
    elif timeout < 0:
        timeout = 0
    else:
        timeout = int(1000000 * timeout)

    req = dmsg.GSChannelJoin(tag=das.next_tag(), p_uid=this_process.my_puid,
                             r_c_uid=das.get_gs_ret_cuid(), name=identifier,
                             timeout=timeout)

    reply = das.gs_request(req)
    assert isinstance(reply, dmsg.GSChannelJoinResponse)

    if dmsg.GSChannelJoinResponse.Errors.TIMEOUT == reply.err:
        raise TimeoutError(f'join request to {identifier} timed out')

    return reply.desc


def destroy(identifier, *, just_decref=False):
    """Asks Global Services to destroy a specified channel.

    :param just_decref: bool default False, just decrement refcount, not hard destroy.
    :param identifier: string indicating channel name or integer indicating a c_uid
    :return: Nothing if successful
    :raises: ChannelError if there is no such channel
    """

    rreq = not just_decref

    if isinstance(identifier, str):
        req_msg = dmsg.GSChannelDestroy(tag=das.next_tag(), p_uid=this_process.my_puid,
                                        r_c_uid=das.get_gs_ret_cuid(), dec_ref=just_decref,
                                        reply_req=rreq,  user_name=identifier)
    else:
        req_msg = dmsg.GSChannelDestroy(tag=das.next_tag(), p_uid=this_process.my_puid,
                                        r_c_uid=das.get_gs_ret_cuid(),dec_ref=just_decref,
                                        reply_req=rreq, c_uid=int(identifier))

    reply_msg = das.gs_request(req_msg, expecting_response=rreq)
    if not rreq:
        return

    assert isinstance(reply_msg, dmsg.GSChannelDestroyResponse)

    ec = dmsg.GSChannelDestroyResponse.Errors

    if ec.SUCCESS == reply_msg.err:
        return
    else:
        raise ChannelError(f'channel destroy {req_msg} failed: {reply_msg.err_info}')


def get_refcnt(identifier):
    return query(identifier, inc_refcnt=True)


def release_refcnt(identifier):
    return destroy(identifier, just_decref=True)

src/dragon/globalservices/channel.py
DragonHPC/dragon
"""Channel related objects internal to Global Services.

This file includes the ChannelContext object, and code that
talks to a shepherd to manipulate channel life cycle.
"""

import dragon.infrastructure.channel_desc as channel_desc
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.messages as dmsg
import dragon.channels as dch

import logging

from dragon.dtypes import DragonError

LOG = logging.getLogger('channel:')


class ChannelContext:
    """Everything to do with a single channel in global services

        This class includes the code that does the lifecycle of the
        channel to the shepherd hosting that channel - sending and receiving
        shepherd command messages.
    """

    def __init__(self, server, request, reply_channel, c_uid, node):
        self.server = server
        self.request = request
        self.destroy_request = None
        self.reply_channel = reply_channel  # TODO: really is a Connection, should change var name
        self.attached_connection = None
        self.is_reading = False
        self.is_refcounted = self.request.options.ref_count
        self.refcnt = 0

        self._descriptor = channel_desc.ChannelDescriptor(node=node, m_uid=request.m_uid,
                                                          c_uid=c_uid,
                                                          name=request.user_name)

    def __str__(self):
        return f"[[{self.__class__.__name__}]] desc:{self.descriptor!r} req:{self.request!r}"

    def incref(self):
        self.refcnt += 1

    @property
    def descriptor(self):
        return self._descriptor

    def _mk_sh_channel_create(self, the_tag):
        self.shchannelcreate_msg =  dmsg.SHChannelCreate(tag=the_tag,
                                    p_uid=dfacts.GS_PUID,
                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                    m_uid=self.descriptor.m_uid,
                                    c_uid=self.descriptor.c_uid,
                                    options=self.request.options.local_opts)
        return self.shchannelcreate_msg

    def mk_sh_channel_destroy(self, the_tag):
        return dmsg.SHChannelDestroy(tag=the_tag,
                                     p_uid=dfacts.GS_PUID,
                                     r_c_uid=dfacts.GS_INPUT_CUID,
                                     c_uid=self.descriptor.c_uid)

    @classmethod
    def construct(cls, server, msg, reply_channel, node_override=None, send_msg=True):
        """Constructs a new channel context and sends out a message
            to a shepherd to create the channel.

            :return: a tuple:
                First element True if construction message issued, False otherwise
                Second element the tag if message issued, None otherwise
                Third element The newly constructed context object if issued, None otherwise
        """

        if msg.user_name in server.channel_names:
            LOG.info(f'channel {msg.user_name} already created')
            if send_msg:
                existing_ctx = server.channel_table[server.channel_names[msg.user_name]]
                rm = dmsg.GSChannelCreateResponse(tag=server.tag_inc(),
                                            ref=msg.tag,
                                            err=dmsg.GSChannelCreateResponse.Errors.ALREADY,
                                            desc=existing_ctx.descriptor)
                reply_channel.send(rm.serialize())
            return False, None, None

        # check if the m_uid is ok
        if not (dfacts.is_pre_defined_pool(msg.m_uid) or msg.m_uid in server.pool_table):
            errmsg = f'm_uid {msg.m_uid} unknown'
            LOG.info(errmsg)
            if send_msg:
                rm = dmsg.GSChannelCreateResponse(tag=server.tag_inc(), ref=msg.tag,
                                            err=dmsg.GSChannelCreateResponse.Errors.FAIL,
                                            err_info=errmsg)
                reply_channel.send(rm.serialize())
            return False, None, None

        this_cuid, auto_name = server.new_cuid_and_default_name()

        if not msg.user_name:
            msg.user_name = auto_name

        outbound_tag = server.tag_inc()

        try:
            which_node = node_override
            if which_node is None:
                which_node = server.choose_channel_node(msg)
        except ValueError:
            # the requested pool to create the channel on does not exist
            # and we need to log this
            errmsg = f'm_uid {msg.m_uid} does not exist - it is neither a predefined pool nor a user defined pool'
            LOG.info(errmsg)
            if send_msg:
                rm = dmsg.GSChannelCreateResponse(tag=server.tag_inc(), ref=msg.tag,
                                                err=dmsg.GSChannelCreateResponse.Errors.FAIL,
                                                err_info=errmsg)
                reply_channel.send(rm.serialize())
            return False, None, None

        context = cls(server, msg, reply_channel, this_cuid, which_node)
        server.channel_names[msg.user_name] = this_cuid
        server.channel_table[this_cuid] = context

        # send request to local services, remember pending process.
        shep_create_msg = context._mk_sh_channel_create(outbound_tag)
        if send_msg:
            LOG.debug(f'shep_inputs: {server.shep_inputs}')
            shep_hdl = server.shep_inputs[which_node]
            LOG.debug(f'sending {shep_create_msg} with muid {shep_create_msg.m_uid} to shep via {shep_hdl} on node {which_node}')

            # In cases where we are sending a large amount of
            # messages, such as with the GSGroupCreate handler,
            # we can fill the GS Input Queue with responses and
            # basically cause the GS / TA / LS to be unable to
            # send/receive any messages. To prevent this, we'll
            # enqueue pending sends and interleave sending and
            # receiving messages to allow us to process responses
            # on the input queue.

            server.pending_sends.put((shep_hdl, shep_create_msg.serialize()))
            LOG.debug(f'request {context} to shep')

        return True, outbound_tag, context

    def complete_construction(self, msg):
        """Completes channel construction, waiting on message from shep

            :return: True or False, according to if construction succeeded
        """

        if not isinstance(msg, dmsg.SHChannelCreateResponse):
            print(f'Failure: The message was: {msg}')
        assert isinstance(msg, dmsg.SHChannelCreateResponse)

        if dmsg.SHChannelCreateResponse.Errors.SUCCESS == msg.err:
            self.descriptor.state = channel_desc.ChannelDescriptor.State.ACTIVE
            self.descriptor.sdesc = msg.desc
            response = dmsg.GSChannelCreateResponse(tag=self.server.tag_inc(),
                                                    ref=self.request.tag,
                                                    err=dmsg.GSChannelCreateResponse.Errors.SUCCESS,
                                                    desc=self.descriptor)

            self.reply_channel.send(response.serialize())
            create_succeeded = True
            if self.is_refcounted:
                self.refcnt = 1  # original requester.
                LOG.debug(f'channel {self.descriptor.c_uid}: refcnt={self.refcnt}')

            # only check pending joins if the create succeeded.  The entity creating
            # the channel will get the error.  There could be a message back to
            # the entity trying to join explaining that the create failed, but
            # there isn't anything that entity can do about it, only the one making
            # the original create
            if self.descriptor.name in self.server.pending_channel_joins:
                for join_req in self.server.pending_channel_joins.get(self.descriptor.name):
                    reply_channel = self.server.get_reply_handle(join_req)
                    rm = dmsg.GSChannelJoinResponse(tag=self.server.tag_inc(),
                                                    ref=join_req.tag,
                                                    err=dmsg.GSChannelJoinResponse.Errors.SUCCESS,
                                                    desc=self.descriptor)
                    reply_channel.send(rm.serialize())
                    LOG.debug(f'join response to {join_req}: {rm}')
        elif dmsg.SHChannelCreateResponse.Errors.FAIL == msg.err:
            self.descriptor.state = channel_desc.ChannelDescriptor.State.DEAD
            response = dmsg.GSChannelCreateResponse(tag=self.server.tag_inc(),
                                                    ref=self.request.tag,
                                                    err=dmsg.GSChannelCreateResponse.Errors.FAIL,
                                                    err_info=msg.err_info)

            # we don't keep descriptor for things that were never alive
            del self.server.channel_table[self.descriptor.c_uid]
            del self.server.channel_names[self.request.user_name]

            # log a warning if joins are orphaned.
            if self.descriptor.name in self.server.pending_channel_joins:
                LOG.warn(f'channel create name {self.descriptor.name} failed with joins pending')

            self.reply_channel.send(response.serialize())
            create_succeeded = False
        else:
            raise RuntimeError(f'got {str(msg)} err {msg.err} unknown')

        self.reply_channel = None  # deliberately drop reference to reply channel

        return create_succeeded

    @staticmethod
    def destroy(server, msg, reply_channel):
        """Action to destroy this channel

        :param server: global server context
        :param msg: destroy request message
        :param reply_channel: response channel
        :return: True if a destroy message was issued and False otherwise
        """
        target_uid, found, errmsg = server.resolve_cuid(msg.user_name, msg.c_uid)
        gscdr = dmsg.GSChannelDestroyResponse

        if not found:
            rm = gscdr(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gscdr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown channel: response to {msg}: {rm}')
            return False
        else:
            channelctx = server.channel_table[target_uid]
            channeldesc = channelctx.descriptor
            cds = channel_desc.ChannelDescriptor.State
            if cds.DEAD == channeldesc.state:
                rm = gscdr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gscdr.Errors.BUSY)
                reply_channel.send(rm.serialize())
                LOG.debug(f'gone; response to {msg}: {rm}')
                return False
            elif cds.PENDING == channeldesc.state:
                rm = gscdr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gscdr.Errors.BUSY,
                           err_info=f'{target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'pending; response to {msg}: {rm}')
                return False
            elif cds.ACTIVE == channeldesc.state:
                channelctx.destroy_request = msg
                channelctx.reply_channel = reply_channel

                if channelctx.is_refcounted and msg.dec_ref:
                    channelctx.refcnt -= 1
                    LOG.debug(f'channel {channelctx.descriptor.c_uid}: refcnt={channelctx.refcnt}')
                    if channelctx.refcnt == 0:
                        really_destroying = True
                    elif channelctx.refcnt < 0:
                        really_destroying = False
                        LOG.warning(f'refcnt overrun from msg={msg!r}')
                    else:
                        really_destroying = False
                else:
                    really_destroying = True

                if not really_destroying:  # just a refcount op that didn't cause destruction.
                    rm = gscdr(tag=server.tag_inc(), ref=msg.tag, err=gscdr.Errors.SUCCESS)
                    reply_channel.send(rm.serialize())
                    LOG.debug(f'channel {channelctx.descriptor.c_uid}: Not destroying, refcnt={channelctx.refcnt}')
                    return False
                else:
                    channelctx.descriptor.state = cds.PENDING  # do we want a pending destruction state?
                    the_tag = server.tag_inc()
                    shep_destroy_msg = channelctx.mk_sh_channel_destroy(the_tag)
                    target_node = channeldesc.node
                    server.pending[the_tag] = channelctx.complete_destruction
                    server.pending_sends.put((server.shep_inputs[target_node], shep_destroy_msg.serialize()))

                    LOG.debug(f'forwarded destroy {msg} as {shep_destroy_msg} to {target_node}')
                    return True
            else:
                raise NotImplementedError('close case')

    def complete_destruction(self, msg):
        gscdr = dmsg.GSChannelDestroyResponse
        shcdr = dmsg.SHChannelDestroyResponse

        if shcdr.Errors.FAIL == msg.err:
            rm = gscdr(tag=self.server.tag_inc(),
                       ref=self.destroy_request.tag,
                       err=gscdr.Errors.UNKNOWN,
                       err_info=msg.err_info)
            destroy_succeeded = False
        elif shcdr.Errors.SUCCESS == msg.err:
            self.descriptor.state = channel_desc.ChannelDescriptor.State.DEAD
            rm = gscdr(tag=self.server.tag_inc(), ref=self.destroy_request.tag,
                       err=gscdr.Errors.SUCCESS)
            destroy_succeeded = True
        else:
            raise NotImplementedError('no case')

        LOG.debug(f'channel destroy response to {self.destroy_request}: {rm}')
        self.reply_channel.send(rm.serialize())

        return destroy_succeeded

src/dragon/globalservices/channel_int.py
DragonHPC/dragon
"""Global Services API for the Group of resources (processes, channels, pools, etc.)
primary resource component. Only process groups are supported at this time.
"""

import signal
import logging
from ..infrastructure import messages as dmsg
from ..infrastructure.parameters import this_process
from ..globalservices import api_setup as das

log = logging.getLogger('group_api')


class GroupError(Exception):
    pass


def _check_msg_types(items):
    for _, msg in items:
        if not isinstance(dmsg.parse(msg), (dmsg.GSChannelCreate, dmsg.GSProcessCreate, dmsg.GSPoolCreate)):
            raise GroupError('An invalid message type was provided.')


def create(items, policy, user_name='', soft=False):
    """Asks Global Services to create a new group of specified resources.

    :param items: list of tuples where each tuple contains a replication factor `n` and the Dragon create message. Dragon create messages do not have to be of the same primary resources, i.e. ChannelCreate and ProcessCreate messages can be mixed to into groups to represent composite resources.
    :type items: list[tuple[int, dragon.infrastructure.messages.Message]]
    :param policy: determines the placement of the group resources
    :type policy: dragon.infrastructure.policy.Policy
    :param user_name: Requested user specified reference name, defaults to ''
    :type user_name: str, optional
    :param soft: If group already exists with given name, do not create and return descriptor instead, defaults to False
    :type soft: bool, optional
    :raises GroupError: if soft is True and not given a group name
    :raises GroupError: if the action fails
    :return: the descriptor of the group object
    :rtype: dragon.infrastructure.group_desc.GroupDescriptor
    """

    if soft and not user_name:
        raise GroupError('soft create requires a user supplied group name')

    # check that we have valid message types on each tuple
    _check_msg_types(items)

    req_msg = dmsg.GSGroupCreate(tag=das.next_tag(), p_uid=this_process.my_puid,
                                 r_c_uid=das.get_gs_ret_cuid(), items=items,
                                 policy=policy, user_name=user_name)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSGroupCreateResponse)

    if reply_msg.desc:
        the_desc = reply_msg.desc
        # TODO: add check that the items tuples correspond to the
        # correct tuples in the response message
    else:
        raise GroupError(f'group create {req_msg} failed, no descriptor returned')

    return the_desc


def kill(identifier, sig=signal.SIGKILL):
    """Asks Global Services to send the processes belonging to a specified group
    a specified signal.

    Note that this is like the unix 'kill' command - the signal given to the process
    might not necessarily be intended to cause it to terminate.
    If the group contains members other than processes, such as channels, then
    this operation has no effect to these members.

    :param identifier: string indicating group name or integer indicating a g_uid
    :type identifier: str|int
    :param sig: signal to use to kill the process, default=signal.SIGKILL
    :type sig: int
    :raises GroupError: if there is no such group
    :raises GroupError: if the group has not yet started
    :raises NotImplementedError: if any other case not implemented
    :return: Nothing if successful
    """
    if isinstance(identifier, str):
        req_msg = dmsg.GSGroupKill(tag=das.next_tag(), p_uid=this_process.my_puid,
                                   r_c_uid=das.get_gs_ret_cuid(), sig=int(sig),
                                   user_name=identifier)
    else:
        req_msg = dmsg.GSGroupKill(tag=das.next_tag(), p_uid=this_process.my_puid,
                                   r_c_uid=das.get_gs_ret_cuid(), sig=int(sig),
                                   g_uid=int(identifier))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSGroupKillResponse)

    ec = dmsg.GSGroupKillResponse.Errors

    if reply_msg.err in {ec.SUCCESS, ec.ALREADY}:
        # if a successful GSGroupKillResponse message was sent or all the
        # processes were already dead
        return reply_msg.desc
    elif ec.UNKNOWN == reply_msg.err:
        raise GroupError(f'group kill {req_msg} failed: {reply_msg.err_info}')
    elif ec.PENDING == reply_msg.err:
        raise GroupError(f'group kill {req_msg} failed pending: {reply_msg.err_info}')
    else:
        raise NotImplementedError('close case')


def destroy(identifier):
    """Asks Global Services to destroy a specified group of resources.
    It will destroy all the members of the group as well as the group itself.

    :param identifier: string indicating group name or integer indicating a g_uid
    :type identifier: str|int
    :raises GroupError: if there is no such group
    :raises GroupError: if the group has not yet started
    :raises NotImplementedError: if any other case not implemented
    :return: the group descriptor in case of a success
    :rtype: dragon.infrastructure.group_desc.GroupDescriptor
    """
    if isinstance(identifier, str):
        req_msg = dmsg.GSGroupDestroy(tag=das.next_tag(), p_uid=this_process.my_puid,
                                      r_c_uid=das.get_gs_ret_cuid(),
                                      user_name=identifier)
    else:
        req_msg = dmsg.GSGroupDestroy(tag=das.next_tag(), p_uid=this_process.my_puid,
                                      r_c_uid=das.get_gs_ret_cuid(),
                                      g_uid=int(identifier))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSGroupDestroyResponse)

    ec = dmsg.GSGroupDestroyResponse.Errors

    if ec.SUCCESS == reply_msg.err:
        return reply_msg.desc
    elif ec.UNKNOWN == reply_msg.err:
        raise GroupError(f'group destroy {req_msg} failed: {reply_msg.err_info}')
    elif ec.PENDING == reply_msg.err:
        raise GroupError(f'group destroy {req_msg} failed pending: {reply_msg.err_info}')
    else:
        raise NotImplementedError('close case')


def add_to(identifier, items):
    """Add existing resources to an existing group.

    :param identifier: string indicating group name or integer indicating a g_uid
    :type identifier: str|int
    :param items: list of int or str corresponding to resource id or name respectively
    :type items: List[ids : int|str]
    :raises GroupError: if the addition of resources failed
    :raises GroupError: if the identifier corresponds to an unknown group
    """

    if isinstance(identifier, str):
        req_msg = dmsg.GSGroupAddTo(tag=das.next_tag(), p_uid=this_process.my_puid,
                                    r_c_uid=das.get_gs_ret_cuid(),
                                    user_name=identifier, items=items)
    else:
        req_msg = dmsg.GSGroupAddTo(tag=das.next_tag(), p_uid=this_process.my_puid,
                                    r_c_uid=das.get_gs_ret_cuid(),
                                    g_uid=int(identifier), items=items)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSGroupAddToResponse)

    if reply_msg.desc:
        the_desc = reply_msg.desc
    elif dmsg.GSGroupAddToResponse.Errors.UNKNOWN == reply_msg.err:
        raise GroupError(f'group addition {req_msg} failed due to unknown group: {reply_msg.err_info}')
    else:
        raise GroupError(f'group addition {req_msg} failed: {reply_msg.err_info}')

    return the_desc


def create_add_to(identifier, items, policy):
    """Create resources and then add them to an existing group.

    :param identifier: string indicating group name or integer indicating a g_uid
    :type identifier: str|int
    :param items: list of tuples where each tuple contains a replication factor `n` and the Dragon create message. Dragon create messages do not have to be of the same primary resources, i.e. ChannelCreate and ProcessCreate messages can be mixed to into groups to represent composite resources.
    :type items: list[tuple[int, dragon.infrastructure.messages.Message]]
    :param policy: determines the placement of the group resources
    :type policy: dragon.infrastructure.policy.Policy
    :raises GroupError: if the addition of resources failed
    :raises GroupError: if the identifier corresponds to an unknown group
    """

    # check that we have valid message types on each tuple
    _check_msg_types(items)

    if isinstance(identifier, str):
        req_msg = dmsg.GSGroupCreateAddTo(tag=das.next_tag(), p_uid=this_process.my_puid,
                                          r_c_uid=das.get_gs_ret_cuid(),
                                          user_name=identifier, items=items, policy=policy)
    else:
        req_msg = dmsg.GSGroupCreateAddTo(tag=das.next_tag(), p_uid=this_process.my_puid,
                                          r_c_uid=das.get_gs_ret_cuid(),
                                          g_uid=int(identifier), items=items, policy=policy)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSGroupCreateAddToResponse)

    if reply_msg.desc:
        return reply_msg.desc
    elif dmsg.GSGroupCreateAddToResponse.Errors.UNKNOWN == reply_msg.err:
        raise GroupError(f'group addition {req_msg} failed due to unknown group: {reply_msg.err_info}')
    else:
        raise GroupError(f'group addition {req_msg} failed: {reply_msg.err_info}')


def remove_from(identifier, items):
    """Remove resources from an existing group.

    :param identifier: string indicating group name or integer indicating a g_uid
    :type identifier: str|int
    :param items: list of int or str corresponding to resource id or name respectively
    :type items: List[ids : int|str]
    :raises GroupError: if the addition of resources failed
    :raises GroupError: if the identifier corresponds to an unknown group
    """

    if isinstance(identifier, str):
        req_msg = dmsg.GSGroupRemoveFrom(tag=das.next_tag(), p_uid=this_process.my_puid,
                                         r_c_uid=das.get_gs_ret_cuid(),
                                         user_name=identifier, items=items)
    else:
        req_msg = dmsg.GSGroupRemoveFrom(tag=das.next_tag(), p_uid=this_process.my_puid,
                                         r_c_uid=das.get_gs_ret_cuid(),
                                         g_uid=int(identifier), items=items)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSGroupRemoveFromResponse)

    if dmsg.GSGroupRemoveFromResponse.Errors.SUCCESS == reply_msg.err:
        return reply_msg.desc
    elif dmsg.GSGroupRemoveFromResponse.Errors.UNKNOWN == reply_msg.err:
        raise GroupError(f'removal of group resources {req_msg} failed due to unknown group: {reply_msg.err_info}')
    else:
        raise GroupError(f'removal of group resources {req_msg} failed: {reply_msg.err_info}')


def destroy_remove_from(identifier, items):
    """Remove resources from an existing group and call destroy on them.
    For processes we call process.kill (sends a SIGKILL signal) and for other
    types of resources we call destroy.

    :param identifier: string indicating group name or integer indicating a g_uid
    :type identifier: str|int
    :param items: list of int or str corresponding to resource id or name respectively
    :type items: List[ids : int|str]
    :raises GroupError: if the addition of resources failed
    :raises GroupError: if the identifier corresponds to an unknown group
    """

    if isinstance(identifier, str):
        req_msg = dmsg.GSGroupDestroyRemoveFrom(tag=das.next_tag(), p_uid=this_process.my_puid,
                                                r_c_uid=das.get_gs_ret_cuid(),
                                                user_name=identifier, items=items)
    else:
        req_msg = dmsg.GSGroupDestroyRemoveFrom(tag=das.next_tag(), p_uid=this_process.my_puid,
                                                r_c_uid=das.get_gs_ret_cuid(),
                                                g_uid=int(identifier), items=items)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSGroupDestroyRemoveFromResponse)

    if dmsg.GSGroupDestroyRemoveFromResponse.Errors.SUCCESS == reply_msg.err:
        return reply_msg.desc
    elif dmsg.GSGroupDestroyRemoveFromResponse.Errors.UNKNOWN == reply_msg.err:
        raise GroupError(f'removal of group resources {req_msg} failed due to unknown group: {reply_msg.err_info}')
    else:
        raise GroupError(f'removal of group resources {req_msg} failed: {reply_msg.err_info}')


# def get_list():
#     """Asks Global Services for a list of the g_uids of all groups.

#     :return: list of the g_uids of all currently existing groups
#     :rtype: list[g_uids]
#     """
#     req_msg = dmsg.GSGroupList(tag=das.next_tag(), p_uid=this_process.my_puid,
#                                r_c_uid=das.get_gs_ret_cuid())

#     reply_msg = das.gs_request(req_msg)
#     assert isinstance(reply_msg, dmsg.GSGroupListResponse)

#     return reply_msg.glist


# def query(identifier):
#     """Asks Global Services for the GroupDescriptor of a specified group of resources.
#     Note you can only query currently existing groups.

#     :param identifier: string indicating group name or integer indicating a g_uid
#     :type identifier: str or int
#     :raises GroupError: if there is no such group
#     :return: GroupDescriptor object corresponding to specified group
#     :rtype: GroupDescriptor
#     """
#     if isinstance(identifier, str):
#         req_msg = dmsg.GSGroupQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
#                                     r_c_uid=das.get_gs_ret_cuid(),
#                                     user_name=identifier)
#     else:
#         req_msg = dmsg.GSGroupQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
#                                     r_c_uid=das.get_gs_ret_cuid(),
#                                     g_uid=int(identifier))

#     reply_msg = das.gs_request(req_msg)
#     assert isinstance(reply_msg, dmsg.GSGroupQueryResponse)

#     if dmsg.GSGroupQueryResponse.Errors.SUCCESS == reply_msg.err:
#         the_desc = reply_msg.desc
#     else:
#         raise GroupError(f'group query {req_msg} failed: {reply_msg.err_info}')

#     return the_desc
src/dragon/globalservices/group.py
DragonHPC/dragon
"""Global services' internal Group context..
"""

from .process_int import ProcessContext
from .channel_int import ChannelContext
from .pool_int import PoolContext
from ..infrastructure import group_desc as group_desc
from ..infrastructure import process_desc as process_desc
from ..infrastructure import facts as dfacts
from ..infrastructure import messages as dmsg
from ..infrastructure.parameters import this_process
from ..infrastructure.policy import Policy
from ..infrastructure import util as dutil
from ..globalservices import api_setup as das
from .group import GroupError
from .policy_eval import PolicyEvaluator

import logging
import copy
import signal

LOG = logging.getLogger('group:')


class PMIJobHelper:
    """
    THe PMIJobHelper class calculates the per-process PMI information required
    to start each PMI rank.
    """

    _FIRST_PMI_JOB_ID = 1
    _NEXT_PMI_JOB_ID = _FIRST_PMI_JOB_ID

    _PMI_PID_DICT = {} # key: hostname; value: list of tuples representing (base, size) allocations

    def __init__(self, items, layout_list):
        self._index = 0
        self.items = items
        self.layout_list = layout_list

        # calculated values
        self.job_id = self.get_next_pmi_job_id()
        self.nranks = self.get_nranks()

        (
            self.nid_list,
            self.host_list,
            self.ppn_map,
            self.lrank_list,
        ) = self.get_nid_and_host_data()
        self.pmi_h_uid_list = list(self.ppn_map.keys())

        self.pmi_nnodes = len(self.host_list)

        self.control_port = dfacts.DEFAULT_PMI_CONTROL_PORT + self.job_id
        self.pid_base_map = {}  # key h_uid, value base value
        LOG.debug('Before allocating pid base _PMI_PID_DICT=%s', str(PMIJobHelper._PMI_PID_DICT))
        for h_uid, lranks in self.ppn_map.items():
            self.pid_base_map[h_uid] = PMIJobHelper.allocate_pmi_pid_base(h_uid, lranks)
        LOG.debug('After allocating pid base _PMI_PID_DICT=%s', str(PMIJobHelper._PMI_PID_DICT))

    @classmethod
    def is_pmi_required(cls, items:list[tuple]):
        for _, res_msg in items:
            resource_msg = dmsg.parse(res_msg)
            if isinstance(resource_msg, dmsg.GSProcessCreate) and resource_msg.pmi_required:
                return True
        return False

    @classmethod
    def get_next_pmi_job_id(cls):
        job_id = cls._NEXT_PMI_JOB_ID
        cls._NEXT_PMI_JOB_ID += 1
        return job_id

    @classmethod
    def allocate_pmi_pid_base(cls, hostname, lranks):
        if hostname in cls._PMI_PID_DICT:
            allocated_bases = cls._PMI_PID_DICT[hostname]
        else:
            allocated_bases = []
            cls._PMI_PID_DICT[hostname] = allocated_bases

        # reserve 0 for the single HSTA agent. this will not be needed
        # when we begin launching HSTA with PMOD
        min_pid = 1

        # handle the case where pid_bases is empty
        if not allocated_bases:
            allocated_bases.append((min_pid, lranks))
            return min_pid

        # loop over (base, size) pairs looking for a free segment of PIDS
        # with size at least lranks

        # (1) for a given (base, size) pair, check the segment starting at
        # the PID just past the previous segment

        prev_end = min_pid

        for base, size in allocated_bases:
            if base - prev_end >= lranks:
                allocated_bases.append((prev_end, lranks))
                allocated_bases.sort()
                return prev_end
            else:
                prev_end = base + size

        # (2) check the final segment from the end of the last (base, size)
        # pair extending to the maximum PID

        max_pid = 510

        if (max_pid + 1) - prev_end >= lranks:
            allocated_bases.append((prev_end, lranks))
            allocated_bases.sort()
            return prev_end

        return -1

    @classmethod
    def free_pmi_pid_base(cls, hostname, base, size):
        allocated_bases = cls._PMI_PID_DICT[hostname]
        allocated_bases.remove((base, size))

    def cleanup(self):
        LOG.debug('Before cleaning up pid base _PMI_PID_DICT=%s', str(PMIJobHelper._PMI_PID_DICT))
        for host_id, count in self.ppn_map.items():
            PMIJobHelper.free_pmi_pid_base(host_id, self.pid_base_map[host_id], count)
        LOG.debug('After cleaning up pid base _PMI_PID_DICT=%s', str(PMIJobHelper._PMI_PID_DICT))

    def get_nranks(self):
        total_ranks = 0
        for count, res_msg in self.items:
            resource_msg = dmsg.parse(res_msg)
            if isinstance(resource_msg, dmsg.GSProcessCreate) and resource_msg.pmi_required:
                total_ranks += count
        return total_ranks

    def get_nid_and_host_data(self):
        """
        Determine which of the resource items requires PMI. For those resources
        that do require PMI, determine number of PMI processes assigned to each h_uid.
        """
        layout_index = 0
        nid_list = []
        host_list = []
        ppn_map = {}
        lrank_list = []
        for count, res_msg in self.items:
            resource_msg = dmsg.parse(res_msg)
            if isinstance(resource_msg, dmsg.GSProcessCreate) and resource_msg.pmi_required:
                for _ in range(count):
                    h_uid = self.layout_list[layout_index].h_uid
                    lrank = ppn_map.get(h_uid, 0)
                    lrank_list.append(lrank)

                    if h_uid not in ppn_map:
                        # add this host to the host list
                        host_list.append(self.layout_list[layout_index].host_name)

                    ppn_map[h_uid] = lrank + 1
                    nid_list.append(list(ppn_map.keys()).index(h_uid))
                    layout_index += 1
            else:
                layout_index += count

        return nid_list, host_list, ppn_map, lrank_list

    def get_host_id_from_index(self, rank):
        layout_index = 0
        for count, res_msg in self.items:
            resource_msg = dmsg.parse(res_msg)
            if isinstance(resource_msg, dmsg.GSProcessCreate) and resource_msg.pmi_required:
                if rank >= count:
                    rank -= count
                    layout_index += count
                else:
                    return self.layout_list[layout_index + rank].h_uid
            else:
                layout_index += count

    def __iter__(self):
        return self

    def __next__(self):
        if self._index < self.nranks:
            lrank = self.lrank_list[self._index]
            h_uid = self.get_host_id_from_index(self._index)
            pmi_info = dmsg.PMIInfo(
                job_id=self.job_id,
                lrank=lrank,
                ppn=self.ppn_map[h_uid],
                nid=self.pmi_h_uid_list.index(h_uid),
                nnodes=self.pmi_nnodes,
                nranks=self.nranks,
                nidlist=self.nid_list,
                hostlist=self.host_list,
                control_port=self.control_port,
                pid_base=self.pid_base_map[h_uid],
            )
            self._index += 1
            return pmi_info
        else:
            raise StopIteration


class GroupContext:
    """Everything to do with a single group of resources in global services.

        This object manages all the transactions to a shepherd concerning
        the lifecycle of a group.
    """

    def __init__(self, server, request, reply_channel, g_uid, policy, pmi_job_helper):
        """_summary_

        :param server: _description_
        :type server: _type_
        :param request: _description_
        :type request: _type_
        :param reply_channel: _description_
        :type reply_channel: _type_
        :param g_uid: _description_
        :type g_uid: _type_
        :param policy: _description_
        :type policy: _type_
        :param pmi_job_id: unique job id for PMI processes
        :type pmi_job_id: int
        """
        self.server = server
        self.request = request
        self.reply_channel = reply_channel
        self.destroy_request = None
        self.pmi_job_helper = pmi_job_helper
        self.destroy_called = False
        self.destroy_remove_success_ids = None # used when destroy_remove is called to keep the items to be
                                               # destroyed after having received all the SHProcessKillResponse messages
        self._descriptor = group_desc.GroupDescriptor(g_uid=g_uid,
                                                      name=request.user_name,
                                                      policy=policy)

    def __str__(self):
        return f"[[{self.__class__.__name__}]] desc:{self.descriptor!r} req:{self.request!r}"

    @property
    def descriptor(self):
        return self._descriptor

    def _update_group_member(self, this_guid, member, lst_idx, item_idx, related_to_create=True):
        self.server.group_table[this_guid].descriptor.sets[lst_idx][item_idx] = group_desc.GroupDescriptor.GroupMember.from_sdict(member)
        # update the count that measures the number of responses we have received when we're creating the resources
        if related_to_create:
            if self.server.group_resource_count[this_guid][lst_idx] >= 1:
                self.server.group_resource_count[this_guid][lst_idx] -= 1

    def _generate_member(self, guid, proc_context, err_msg_type=None, channel_related=False):
        if err_msg_type:
            if channel_related:
                response_msg = f'Failed to create {err_msg_type} channel for {proc_context.descriptor.p_uid} that belongs to group {guid}.'
            else:
                response_msg = f'Process {proc_context.descriptor.p_uid} that belongs to group {guid} failed with error: {err_msg_type}.'
            error_code = dmsg.GSProcessCreateResponse.Errors.FAIL
        else:
            response_msg = f'Process {proc_context.descriptor.p_uid} that belongs to group {guid} successfully created.'
            error_code = dmsg.GSProcessCreateResponse.Errors.SUCCESS

        LOG.info(response_msg)

        member = {'state': proc_context.descriptor.state,
                  'uid': proc_context.descriptor.p_uid,
                  'placement': proc_context.descriptor.node,
                  'desc': proc_context.descriptor,
                  'error_code': error_code,
                  'error_info': response_msg}

        return member

    def _remove_proc_from_group(self, puid, lst_idx, item_idx):
        # assign a None value to act as a placeholder temporarily
        self.descriptor.sets[lst_idx][item_idx] = None
        del self.server.resource_to_group_map[puid]

    def _update_ds_after_removals(self):
        g_uid = self.descriptor.g_uid
        # first, clean up descriptor.sets from None values
        for lst_idx, lst in enumerate(self.descriptor.sets):
            temp = [i for i in lst if i is not None]
            if temp:
                self.descriptor.sets[lst_idx] = temp
            else:
                # avoid keeping empty lists
                del self.descriptor.sets[lst_idx]

        # next, update the helper DS
        self.server.group_resource_count[g_uid] = []
        for lst_idx, lst in enumerate(self.descriptor.sets):
            self.server.group_resource_count[g_uid].append(0)
            for item_idx, item in enumerate(lst):
                self.server.resource_to_group_map[item.uid] = (g_uid, (lst_idx, item_idx))

    def _check_if_item_in_group(self, identifier, target_guid, failed_ids, success_ids):
        if isinstance(identifier, str):
            puid, found, _ = self.server.resolve_puid(name=identifier)
        else:
            puid, found, _ = self.server.resolve_puid(p_uid=identifier)

        if not found:
            # not a valid process
            failed_ids.append(identifier)
        else:
            if puid in self.server.resource_to_group_map:
                guid, (lst_idx, item_idx) = self.server.resource_to_group_map[puid]
                if guid != target_guid:
                    LOG.info(f"Error: the target_guid {target_guid} was not the same as the guid {guid} that process {puid} belongs to.")
                    failed_ids.append(identifier)
                else:
                    success_ids.append((puid, lst_idx, item_idx))
            else:
                # already dead or not a member of the group
                failed_ids.append(identifier)

    @classmethod
    def construct(cls, server, msg, reply_channel):
        """Starts the resources that belong to the requested group.
        It also creates a GroupContext and registers it with the server.

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: GSGroupCreate message
        :type msg: dragon.infrastructure.messages.GSGroupCreate
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :raises NotImplementedError: when a channel is requested as a member resource
        :raises NotImplementedError: when a memory pool is requested as a member resource
        :raises GroupError: when the message for the requested resources is of unknown type
        :raises GroupError: if a replication factor `n` of any tuple is less than 1
        :return: True if there is a pending continuation, False otherwise
        :rtype: bool
        """
        if msg.user_name in server.group_names:
            LOG.info(f'group name {msg.user_name} in use')
            existing_ctx = server.group_table[server.group_names[msg.user_name]]
            rm = dmsg.GSGroupCreateResponse(tag=server.tag_inc(),
                                            ref=msg.tag,
                                            desc=existing_ctx.descriptor)
            reply_channel.send(rm.serialize())
            return False

        this_guid, auto_name = server.new_guid_and_default_name()

        if not msg.user_name:
            msg.user_name = auto_name

        policy = msg.policy
        assert isinstance(policy, Policy)

        # Policy Evaluator wants a list of policies, one for each member
        total_members = sum([n for n, _ in msg.items])
        policies = [policy] * total_members

        # Gather our node list and evaluate the policies against it
        layout_list = server.policy_eval.evaluate(policies=policies)
        LOG.debug('layout_list=%s', layout_list)

        pmi_job_helper = None
        if PMIJobHelper.is_pmi_required(msg.items):
            # Setup PMI Helper
            pmi_job_helper = PMIJobHelper(msg.items, layout_list)
            pmi_job_iter = iter(pmi_job_helper)

        group_context = cls(server=server, request=msg, reply_channel=reply_channel,
                            g_uid=this_guid, policy=msg.policy, pmi_job_helper=pmi_job_helper)

        server.group_names[msg.user_name] = this_guid
        server.group_table[this_guid] = group_context
        server.group_resource_count[this_guid] = [] # list of items corresponding to the multiplicity of each list in server.group_resource_list

        # msg.items is a list of tuples of the form (count:int, create_msg: dragon.infrastructure.messages.Message)
        layout_index = 0
        for tuple_idx, item in enumerate(msg.items):
            count, res_msg = item
            resource_msg = dmsg.parse(res_msg)

            LOG.debug(f'Creating msg {resource_msg}')
            if count > 0:
                # initialize a new list
                group_context.descriptor.sets.append([])
                # add a new entry to server.group_resource_count[this_guid] list
                server.group_resource_count[this_guid].append(int(count))
                for item_idx in range(count):
                    item_layout = layout_list[layout_index]
                    layout_index += 1

                    # we need a unique copy of the msg for each member
                    msg = copy.deepcopy(resource_msg)
                    msg.tag = das.next_tag()

                    # create a unique name for this resource based on the tuple user_name
                    if msg.user_name:
                        msg.user_name = f'{msg.user_name}.{this_guid}.{tuple_idx}.{item_idx}'

                    # update the layout to place the process correctly
                    msg.layout = item_layout

                    # add a new resource into this list
                    group_context.descriptor.sets[tuple_idx] += [group_desc.GroupDescriptor.GroupMember()]

                    # this is where we're actually starting the creation of the group members
                    if isinstance(resource_msg, dmsg.GSProcessCreate):

                        if msg.pmi_required:
                            # The PMIJobHelper generates the pmi_info structure
                            # from the layout_map and list of group items.
                            msg._pmi_info = next(pmi_job_iter)
                            LOG.info(f'{msg._pmi_info=}')

                        issued, outbound_tag, proc_context = ProcessContext.construct(server, msg, reply_channel,
                                                                                      head=False, send_msg=False,
                                                                                      belongs_to_group=True)

                        server.resource_to_group_map[proc_context.descriptor.p_uid] = (this_guid, (tuple_idx, item_idx))

                        if issued and outbound_tag and proc_context:
                            server.pending[outbound_tag] = group_context.complete_construction
                            server.group_to_pending_resource_map[(outbound_tag, this_guid)] = proc_context
                        else:
                            if proc_context and outbound_tag == 'already': # the process was already created
                                LOG.debug(f"The process {proc_context.descriptor.p_uid} was already created.")
                                member = {'state': proc_context.descriptor.state,
                                          'uid': proc_context.descriptor.p_uid,
                                          'placement': proc_context.descriptor.node,
                                          'error_code': dmsg.GSProcessCreateResponse.Errors.ALREADY,
                                          'desc': proc_context.descriptor}
                                # Update the group with the corresponding GroupMember
                                group_context._update_group_member(this_guid, member, tuple_idx, item_idx)
                            elif proc_context and not outbound_tag:
                            # this is the case where we needed to make inf channels and
                            # there is a pending request related to the process channels
                            # and we don't need to do anything for now. When the channel related pending will
                            # be cleared, a group pending construct_completion will be issued from process_int
                                continue

                            else: # there was a failure
                                # in this case, outbound_tag contains the error message
                                LOG.debug(f"There was a failure in the process {proc_context.descriptor.p_uid} creation: {outbound_tag}")
                                member = group_context._generate_member(this_guid, proc_context, outbound_tag)
                                # Update the group with the corresponding GroupMember
                                group_context._update_group_member(this_guid, member, tuple_idx, item_idx)

                    elif isinstance(resource_msg, dmsg.GSChannelCreate):
                        raise NotImplementedError
                        # ChannelContext.construct(server, msg, reply_channel,
                        #                          belongs_to_group=True)
                    elif isinstance(resource_msg, dmsg.GSPoolCreate):
                        raise NotImplementedError
                        # PoolContext.construct(server, msg, reply_channel,
                        #                       belongs_to_group=True)
                    else:
                        raise GroupError(f'Unknown msg type {resource_msg} for a Group member.')
            else:
                raise GroupError('The Group should include at least one member in each subgroup.')

        return True

    def _construction_helper(self, msg):
        # get the resource context corresponding to this response message
        if (msg.ref, self.descriptor.g_uid) in self.server.group_to_pending_resource_map:
            proc_context = self.server.group_to_pending_resource_map.pop((msg.ref, self.descriptor.g_uid))
        else:
            raise GroupError('Error with group pending resource.')

        # this p_uid should exist in server.resource_to_group_map dict
        this_puid = proc_context.descriptor.p_uid
        guid, (lst_idx, item_idx) = self.server.resource_to_group_map[this_puid]

        if isinstance(msg, dmsg.GSProcessCreateResponse) and msg.err == dmsg.GSProcessCreateResponse.Errors.FAIL:
            # in this case, there was a failure related to the creation of the process channel
            # and _construction_helper() was called from process_int
            # So, we just update the group and we're done with this process
            LOG.debug(f"There was a failure in the process {proc_context.descriptor.p_uid} creation: {msg.err_info}")
            member = self._generate_member(guid, proc_context, msg.err_info)
            self._update_group_member(guid, member, lst_idx, item_idx)
            return

        # complete the construction of this resource without sending any response messages to the client
        succeeded = proc_context.complete_construction(msg, send_msg=False)

        if dmsg.SHProcessCreateResponse.Errors.SUCCESS == msg.err:
            # this process is successfully created
            member = self._generate_member(guid, proc_context, None)

            # in this case, something went wrong with the creation of stdin/stdout/stderr channels
            if not succeeded:
                member = self._generate_member(guid, proc_context, 'stdin/stdout/stderr', channel_related=True)

        elif dmsg.SHProcessCreateResponse.Errors.FAIL == msg.err:
            member = self._generate_member(guid, proc_context, msg.err_info)

            del self.server.resource_to_group_map[this_puid]

        else:
            raise RuntimeError(f'got {msg!s} err {msg.err} unknown')

        # Update the group with the corresponding GroupMember
        self._update_group_member(guid, member, lst_idx, item_idx)

        return succeeded, guid, lst_idx

    def complete_construction(self, msg):
        """Completes construction of a new managed group.

        :param msg: SHProcessCreateResponse message
        :type msg: dragon.infrastructure.messages.SHProcessCreateResponse
        :raises GroupError: when the context corresponding to the pending resource cannot be found
        :raises RuntimeError: when the message error is of unknown type
        :return: True or False according to success
        :rtype: bool
        """
        LOG.info("group_complete_construction")

        # TODO: implement for other types of resources

        succeeded, guid, lst_idx = self._construction_helper(msg)

        # if we have received responses for all the members of this list
        if self.server.group_resource_count[guid][lst_idx] == 0:
            # next, check if we have responses for the other lists as well
            if all(item == 0 for item in self.server.group_resource_count[guid]):
                # Update the group descriptor state to active
                self.descriptor.state = group_desc.GroupDescriptor.State.ACTIVE

                # now we can send the GSGroupCreateResponse msg back to the client
                response = dmsg.GSGroupCreateResponse(tag=self.server.tag_inc(),
                                                      ref=self.request.tag,
                                                      desc=self.descriptor)
                self.reply_channel.send(response.serialize())
                LOG.debug(f'create response sent, tag {response.tag} ref {response.ref} pending cleared')

        # since this server.pending entry was related to a particular outbound tag
        # and a specific process creation, return the value returned from
        # proc_context.complete_construction
        return succeeded

    @staticmethod
    def add(server, msg, reply_channel):
        """Action to add existing resources to a group.

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: addTo request message
        :type msg: dragon.infrastructure.messages.GSGroupAddTo
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :raises GroupError: if the list of items to add is empty
        :raises NotImplementedError: when the group's state is unknown, i.e., other than ACTIVE, DEAD, PENDING
        :return: True if a successful addTo response message was issued and False otherwise
        :rtype: bool
        """

        if len(msg.items) == 0 or msg.items is None:
            raise GroupError('There should be at least one member to add to the group.')

        target_uid, found, errmsg = server.resolve_guid(msg.user_name, msg.g_uid)
        gsgar = dmsg.GSGroupAddToResponse

        succeeded = False

        if not found:
            rm = gsgar(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gsgar.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown group of resources: response to {msg}: {rm}')
            return succeeded
        else:
            groupctx = server.group_table[target_uid]
            groupdesc = groupctx.descriptor
            gds = group_desc.GroupDescriptor.State

            if gds.DEAD == groupdesc.state:
                rm = gsgar(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgar.Errors.DEAD)
                reply_channel.send(rm.serialize())
                LOG.debug(f'group dead, response to {msg}: {rm}')
                return succeeded

            # this is a highly unlikely case to happen
            # because the user must have received a create response back
            # before requesting adding to this group, which means that the
            # state of the group is not pending anymore
            # so, what we really want here is to log this is happening and move on
            elif gds.PENDING == groupdesc.state:
                rm = gsgar(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgar.Errors.PENDING,
                           err_info=f'group {target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'group pending while add_to request -- this should not be happening, response to {msg}: {rm}')
                return succeeded
            elif gds.ACTIVE == groupdesc.state:
                # update the group's state to PENDING as long as the addTo is pending
                groupdesc.state = gds.PENDING

                groupctx.request = msg
                groupctx.reply_channel = reply_channel
                existing_lists = len(groupdesc.sets)

                failed_ids = [] # keep potential resources that were not found or were dead
                for i, identifier in enumerate(msg.items):
                    if isinstance(identifier, str):
                        p_uid, found, errmsg = server.resolve_puid(name=identifier)
                    else:
                        p_uid, found, errmsg = server.resolve_puid(p_uid=int(identifier))

                    if not found:
                        failed_ids.append(identifier)
                    else:
                        # create a single list in groupdesc.sets for each resource we add
                        lst_idx = i + existing_lists
                        item_idx = 0 # each list will have a single item
                        groupdesc.sets.append([])

                        pctx = server.process_table[p_uid]
                        pdesc = pctx.descriptor
                        pds = process_desc.ProcessDescriptor.State
                        if pdesc.state == pds.DEAD:
                            failed_ids.append(identifier)
                        else:
                            # add each item to a single list in the groupdesc.sets list
                            groupdesc.sets[lst_idx] += [group_desc.GroupDescriptor.GroupMember()]
                            server.resource_to_group_map[p_uid] = (target_uid, (lst_idx, item_idx))
                            member = groupctx._generate_member(target_uid, pctx)
                            groupctx._update_group_member(target_uid, member, lst_idx, item_idx, related_to_create=False)

                # if any of the items to be added were not found or were dead
                # we return a FAIL response message to the client
                if failed_ids:
                    errmsg = f"The items {failed_ids} were not found or already dead."
                    response = gsgar(tag=server.tag_inc(),
                                     ref=msg.tag,
                                     err=gsgar.Errors.FAIL,
                                     err_info=errmsg)
                else:
                    response = gsgar(tag=server.tag_inc(),
                                     ref=msg.tag,
                                     err=gsgar.Errors.SUCCESS,
                                     desc=groupdesc)
                    succeeded = True

                groupdesc.state = gds.ACTIVE
                reply_channel.send(response.serialize())
                LOG.debug(f'addition response sent, tag {response.tag} ref {response.ref} pending cleared')

                return succeeded

            else:
                raise NotImplementedError('close case')

    @staticmethod
    def create_add(server, msg, reply_channel):
        """Action to first create resources and then add them to an existing group.

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: addTo request message
        :type msg: dragon.infrastructure.messages.GSGroupCreateAddTo
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :raises GroupError: when the message for the requested resources is of unknown type
        :raises GroupError: if a replication factor `n` of any tuple is less than 1
        :raises NotImplementedError: when the group's state is unknown, i.e., other than ACTIVE, DEAD, PENDING
        :return: True if a successful addTo message was issued and False otherwise
        :rtype: bool
        """

        target_uid, found, errmsg = server.resolve_guid(msg.user_name, msg.g_uid)
        gsgacr = dmsg.GSGroupCreateAddToResponse

        if not found:
            rm = gsgacr(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gsgacr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown group of resources: response to {msg}: {rm}')
            return False
        else:
            groupctx = server.group_table[target_uid]
            groupdesc = groupctx.descriptor
            gds = group_desc.GroupDescriptor.State

            if gds.DEAD == groupdesc.state:
                rm = gsgacr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgacr.Errors.DEAD)
                reply_channel.send(rm.serialize())
                LOG.debug(f'group dead, response to {msg}: {rm}')
                return False

            # this is a highly unlikely case to happen
            # because the user must have received a create response back
            # before requesting adding to this group, which means that the
            # state of the group is not pending anymore
            # so, what we really want here is to log this is happening and move on
            elif gds.PENDING == groupdesc.state:
                rm = gsgacr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgacr.Errors.PENDING,
                           err_info=f'group {target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'group pending while create_add request -- this should not be happening, response to {msg}: {rm}')
                return False
            elif gds.ACTIVE == groupdesc.state:
                # update the group's state to PENDING as long as the addTo is pending
                groupdesc.state = gds.PENDING

                groupctx.request = msg
                groupctx.reply_channel = reply_channel

                if msg.policy:
                    policy = msg.policy
                    assert isinstance(policy, Policy)
                else:
                    policy = groupdesc.policy

                # Policy Evaluator wants a list of policies, one for each member
                total_members = sum([n for n, _ in msg.items])
                policies = [policy] * total_members

                # Gather our node list and evaluate the policies against it
                layout_list = server.policy_eval.evaluate(policies=policies)
                LOG.debug('layout_list=%s', layout_list)

                # we need the number of existing lists in the group
                existing_lists = len(groupdesc.sets)

                # msg.items is a list of tuples of the form (count:int, create_msg: dragon.infrastructure.messages.Message)
                layout_index = 0
                for tuple_idx, item in enumerate(msg.items):
                    count, res_msg = item
                    resource_msg = dmsg.parse(res_msg)

                    LOG.debug(f'Creating msg {resource_msg}')
                    if count > 0:
                        # initialize a new list
                        groupdesc.sets.append([])
                        # add a new entry to server.group_resource_count[target_uid] list
                        server.group_resource_count[target_uid].append(int(count))
                        for item_idx in range(count):
                            item_layout = layout_list[layout_index]
                            layout_index += 1

                            # we need a unique copy of the msg for each member
                            msg = copy.deepcopy(resource_msg)

                            # create a unique name for this resource based on the tuple user_name
                            if msg.user_name:
                                msg.user_name = f'{msg.user_name}.{target_uid}.{tuple_idx+existing_lists}.{item_idx}'

                            # update the layout to place the process correctly
                            msg.layout = item_layout

                            # add a new resource into this list
                            groupdesc.sets[tuple_idx+existing_lists] += [group_desc.GroupDescriptor.GroupMember()]

                            # this is where we're actually starting the creation of the group members
                            if isinstance(resource_msg, dmsg.GSProcessCreate):
                                issued, outbound_tag, proc_context = ProcessContext.construct(server, msg, reply_channel,
                                                                                              head=False, send_msg=False,
                                                                                              belongs_to_group=True,
                                                                                              addition=True)

                                server.resource_to_group_map[proc_context.descriptor.p_uid] = (target_uid, (tuple_idx+existing_lists, item_idx))

                                if issued and outbound_tag and proc_context:
                                    server.pending[outbound_tag] = groupctx.complete_addition
                                    server.group_to_pending_resource_map[(outbound_tag, target_uid)] = proc_context
                                else:
                                    if proc_context and outbound_tag == 'already': # the process was already created
                                        member = {'state': proc_context.descriptor.state,
                                                  'uid': proc_context.descriptor.p_uid,
                                                  'placement': proc_context.descriptor.node,
                                                  'error_code': dmsg.GSProcessCreateResponse.Errors.ALREADY,
                                                  'desc': proc_context.descriptor}
                                        # Update the group with the corresponding GroupMember
                                        groupctx._update_group_member(target_uid, member, tuple_idx+existing_lists, item_idx)
                                    elif proc_context and not outbound_tag:
                                        # this is the case where we needed to make inf channels and
                                        # there is a pending request related to the process channels
                                        # and we don't need to do anything for now. When the channel related pending will
                                        # be cleared, a group pending construct_completion will be issued from process_int
                                        continue
                                    else: # there was a failure
                                        LOG.debug(f"There was a failure in the process {proc_context.descriptor.p_uid} creation: {outbound_tag}")
                                        member = groupctx._generate_member(target_uid, proc_context, outbound_tag)
                                        # Update the group with the corresponding GroupMember
                                        groupctx._update_group_member(target_uid, member, tuple_idx+existing_lists, item_idx)

                            elif isinstance(resource_msg, dmsg.GSChannelCreate):
                                raise NotImplementedError
                                # ChannelContext.construct(server, msg, reply_channel,
                                #                          belongs_to_group=True)
                            elif isinstance(resource_msg, dmsg.GSPoolCreate):
                                raise NotImplementedError
                                # PoolContext.construct(server, msg, reply_channel,
                                #                       belongs_to_group=True)
                            else:
                                raise GroupError(f'Unknown msg type {resource_msg} for a Group member.')
                    else:
                        raise GroupError('The Group should include at least one member in each subgroup.')
                return True

            else:
                raise NotImplementedError('close case')

    def complete_addition(self, msg):
        """Completes the addition of new resources to a managed group.

        :param msg: SHProcessCreateResponse message
        :type msg: dragon.infrastructure.messages.SHProcessCreateResponse
        :raises GroupError: when the context corresponding to the pending resource cannot be found
        :raises RuntimeError: when the message error is of unknown type
        :return: True or False according to success
        :rtype: bool
        """

        LOG.info("group_complete_addition")

        # TODO: implement for other types of resources

        succeeded, guid, lst_idx = self._construction_helper(msg)

        # if we have received responses for all the members of this list
        if self.server.group_resource_count[guid][lst_idx] == 0:
            # next, check if we have responses for the other lists as well
            if all(item == 0 for item in self.server.group_resource_count[guid]):
                # Update the group descriptor state to active
                self.descriptor.state = group_desc.GroupDescriptor.State.ACTIVE

                # now we can send the GSGroupCreateAddToResponse msg back to the client
                response = dmsg.GSGroupCreateAddToResponse(tag=self.server.tag_inc(),
                                                           ref=self.request.tag,
                                                           err=dmsg.GSGroupCreateAddToResponse.Errors.SUCCESS,
                                                           desc=self.descriptor)
                self.reply_channel.send(response.serialize())
                LOG.debug(f'addition response sent, tag {response.tag} ref {response.ref} pending cleared')

        # since this server.pending entry was related to a particular outbound tag
        # and a specific process creation, return the value returned from
        # proc_context.complete_construction
        return succeeded

    @staticmethod
    def remove(server, msg, reply_channel):
        """Action to remove resources from a group.
        If there is at least one resource that is not found in the group, then we do not proceed with
        the removal request and we return a FAIL response message back to the client.

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: removeFrom request message
        :type msg: dragon.infrastructure.messages.GSGroupRemoveFrom
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :raises GroupError: if the list of items to remove is empty
        :raises NotImplementedError: when the group's state is unknown, i.e., other than ACTIVE, DEAD, PENDING
        :return: True if a successful removeFrom response message was issued and False otherwise
        :rtype: bool
        """

        if len(msg.items) == 0 or msg.items is None:
            raise GroupError('There should be at least one member to delete from the group.')

        target_uid, found, errmsg = server.resolve_guid(msg.user_name, msg.g_uid)
        gsgrr = dmsg.GSGroupRemoveFromResponse

        succeeded = False

        if not found:
            rm = gsgrr(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gsgrr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown group of resources: response to {msg}: {rm}')
            return succeeded
        else:
            groupctx = server.group_table[target_uid]
            groupdesc = groupctx.descriptor
            gds = group_desc.GroupDescriptor.State

            if gds.DEAD == groupdesc.state:
                rm = gsgrr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgrr.Errors.DEAD)
                reply_channel.send(rm.serialize())
                LOG.debug(f'group dead, response to {msg}: {rm}')
                return succeeded

            # this is a highly unlikely case to happen
            # because the user must have received a create response back
            # before requesting adding to this group, which means that the
            # state of the group is not pending anymore
            # so, what we really want here is to log this is happening and move on
            elif gds.PENDING == groupdesc.state:
                rm = gsgrr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgrr.Errors.PENDING,
                           err_info=f'group {target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'group pending while remove request -- this should not be happening, response to {msg}: {rm}')
                return succeeded
            elif gds.ACTIVE == groupdesc.state:
                # update the group's state to PENDING as long as the removeFrom is pending
                groupdesc.state = gds.PENDING

                groupctx.request = msg
                groupctx.reply_channel = reply_channel

                failed_ids = [] # keep potential resources that were not found to be members of the group
                success_ids = []
                # TODO: take into account other types of resources except for processes
                for identifier in msg.items:
                    groupctx._check_if_item_in_group(identifier, target_uid, failed_ids, success_ids)

                # if any of the items to be removed were not found in the group
                # we return a FAIL response message to the client and do not remove any items from the group
                if failed_ids:
                    errmsg = f"The items with ids {failed_ids} are not members of the group. The removal of items from group failed."
                    response = gsgrr(tag=server.tag_inc(),
                                     ref=msg.tag,
                                     err=gsgrr.Errors.FAIL,
                                     err_info=errmsg)
                    del failed_ids
                else:
                    # proceed to removing the items only if all items were found in the group
                    for puid, lst_idx, item_idx in success_ids:
                        groupctx._remove_proc_from_group(puid, lst_idx, item_idx)
                    del success_ids
                    # Update the dict that holds the position of the members of the group
                    groupctx._update_ds_after_removals()

                    response = gsgrr(tag=server.tag_inc(),
                                     ref=msg.tag,
                                     err=gsgrr.Errors.SUCCESS,
                                     desc=groupdesc)
                    succeeded = True

                groupdesc.state = gds.ACTIVE
                reply_channel.send(response.serialize())
                LOG.debug(f'removal response sent, tag {response.tag} ref {response.ref}')

                return succeeded

            else:
                raise NotImplementedError('close case')

    @staticmethod
    def destroy_remove(server, msg, reply_channel):
        """Action to remove resources from a group and also destroy these resources.
        For processes we call process.kill (sends a SIGKILL signal) and for other
        types of resources we call destroy.
        If there is at least one resource that is not found in the group, then we do not proceed with
        the removal/destroy request and we return a FAIL response message back to the client.

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: DestroyRemoveFrom request message
        :type msg: dragon.infrastructure.messages.GSGroupDestroyRemoveFrom
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :raises GroupError: if the list of items to remove is empty
        :raises NotImplementedError: when the group's state is unknown, i.e., other than ACTIVE, DEAD, PENDING
        :return: True if a successful DestroyRemoveFrom response message was issued and False otherwise
        :rtype: bool
        """

        if len(msg.items) == 0 or msg.items is None:
            raise GroupError('There should be at least one member to delete from the group.')

        target_uid, found, errmsg = server.resolve_guid(msg.user_name, msg.g_uid)
        gsgdrr = dmsg.GSGroupDestroyRemoveFromResponse

        succeeded = False

        if not found:
            rm = gsgdrr(tag=server.tag_inc(),
                        ref=msg.tag,
                        err=gsgdrr.Errors.UNKNOWN,
                        err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown group of resources: response to {msg}: {rm}')
            return succeeded
        else:
            groupctx = server.group_table[target_uid]
            groupdesc = groupctx.descriptor
            gds = group_desc.GroupDescriptor.State

            if gds.DEAD == groupdesc.state:
                rm = gsgdrr(tag=server.tag_inc(),
                            ref=msg.tag,
                            err=gsgdrr.Errors.DEAD)
                reply_channel.send(rm.serialize())
                LOG.debug(f'group dead, response to {msg}: {rm}')
                return succeeded

            # this is a highly unlikely case to happen
            # because the user must have received a create response back
            # before requesting adding to this group, which means that the
            # state of the group is not pending anymore
            # so, what we really want here is to log this is happening and move on
            elif gds.PENDING == groupdesc.state:
                rm = gsgdrr(tag=server.tag_inc(),
                            ref=msg.tag,
                            err=gsgdrr.Errors.PENDING,
                            err_info=f'group {target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'group pending while destroy_remove request -- this should not be happening, response to {msg}: {rm}')
                return succeeded
            elif gds.ACTIVE == groupdesc.state:
                # update the group's state to PENDING as long as the DestroyRemoveFrom is pending
                groupdesc.state = gds.PENDING

                groupctx.destroy_request = msg
                groupctx.reply_channel = reply_channel

                server.group_destroy_resource_count[target_uid] = []

                failed_ids = [] # keep potential resources that were not found to be members of the group
                groupctx.destroy_remove_success_ids = []
                for identifier in msg.items:
                    groupctx._check_if_item_in_group(identifier, target_uid, failed_ids, groupctx.destroy_remove_success_ids)

                # if any of the items to be removed were not found in the group
                # we return a FAIL response message to the client and do not remove any items from the group
                if failed_ids:
                    errmsg = f"The items with ids {failed_ids} are not members of the group. The removal of items from group failed."
                    response = gsgdrr(tag=server.tag_inc(),
                                      ref=msg.tag,
                                      err=gsgdrr.Errors.FAIL,
                                      err_info=errmsg)
                    del failed_ids
                    groupdesc.state = gds.ACTIVE
                    reply_channel.send(response.serialize())
                    LOG.debug(f'DestroyRemoveFrom response sent, tag {response.tag} ref {response.ref}')
                else:
                    # proceed with destroying the resources
                    for puid, lst_idx, item_idx in groupctx.destroy_remove_success_ids:
                        item = groupdesc.sets[lst_idx][item_idx]
                        if item.state == process_desc.ProcessDescriptor.State.DEAD:
                            # the process is already dead, so we just need to remove it from the group
                            groupctx._remove_proc_from_group(puid, lst_idx, item_idx)
                        else:
                            # call kill on the process
                            # this process will be removed from the group after the completion of the pending
                            # SHKillProcess request and this happens in complete_kill()
                            resource_msg = groupctx._mk_gs_proc_kill(item.uid)
                            issued, outbound_tag = ProcessContext.kill(server, resource_msg, dutil.AbsorbingChannel())
                            pctx = server.process_table[item.uid]

                            if issued:
                                # we need to issue a pending operation related to this process
                                groupdesc.sets[lst_idx][item_idx].state = process_desc.ProcessDescriptor.State.PENDING
                                groupdesc.sets[lst_idx][item_idx].desc.state = process_desc.ProcessDescriptor.State.PENDING

                                groupctx.destroy_called = True # make sure we won't send a GSGroupKillResponse msg
                                server.pending[outbound_tag] = groupctx.complete_kill
                                server.group_to_pending_resource_map[(outbound_tag, target_uid)] = pctx
                                server.group_destroy_resource_count[target_uid].append(1)
                            else:
                                # the process is either unknown or dead or pending
                                # update the member's state accordingly
                                groupdesc.sets[lst_idx][item_idx].state = pctx.descriptor.state
                                groupdesc.sets[lst_idx][item_idx].desc.state = pctx.descriptor.state

                # in this case, all the processes were already dead or no pending continuation
                # was issued and we need to send a response to the client
                if not server.group_destroy_resource_count[target_uid]:
                    # Update the dict that holds the position of the members of the group
                    groupctx._update_ds_after_removals()

                    response = gsgdrr(tag=server.tag_inc(),
                                      ref=msg.tag,
                                      err=gsgdrr.Errors.SUCCESS,
                                      desc=groupdesc)
                    succeeded = True
                    groupdesc.state = gds.ACTIVE
                    reply_channel.send(response.serialize())
                    LOG.debug(f'destroy_remove response sent, tag {response.tag} ref {response.ref}')
                return succeeded
            else:
                raise NotImplementedError('close case')

    def _mk_gs_proc_kill(self, puid, sig=signal.SIGKILL):
        return dmsg.GSProcessKill(tag=das.next_tag(), p_uid=this_process.my_puid,
                                  r_c_uid=dfacts.GS_INPUT_CUID,
                                  t_p_uid=puid, sig=int(sig))

    @staticmethod
    def kill(server, msg, reply_channel):
        """Action to to send the processes belonging to a specified group
        a specified signal.

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: kill request message
        :type msg: dragon.infrastructure.messages.GSGroupKill
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :raises NotImplementedError: when the group's state is unknown, i.e., other than ACTIVE, DEAD, PENDING
        :return: True if a successful kill message was issued and False otherwise
        :rtype: bool
        """
        target_uid, found, errmsg = server.resolve_guid(msg.user_name, msg.g_uid)
        gsgkr = dmsg.GSGroupKillResponse

        if not found:
            rm = gsgkr(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gsgkr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown group of resources: response to {msg}: {rm}')
            return False
        else:
            groupctx = server.group_table[target_uid]
            groupdesc = groupctx.descriptor
            gds = group_desc.GroupDescriptor.State

            if gds.DEAD == groupdesc.state:
                rm = gsgkr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgkr.Errors.DEAD)
                reply_channel.send(rm.serialize())
                LOG.debug(f'group dead, response to {msg}: {rm}')
                return False

            # this is a highly unlikely case to happen
            # because the user must have received a create response back
            # before requesting killing this group, which means that the
            # state of the group is not pending anymore
            # so, what we really want here is to log this is happening and move on
            elif gds.PENDING == groupdesc.state:
                rm = gsgkr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgkr.Errors.PENDING,
                           err_info=f'group {target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'group pending while kill request -- this should not be happening, response to {msg}: {rm}')
                return False
            elif gds.ACTIVE == groupdesc.state:
                groupdesc.state = gds.PENDING
                groupctx.destroy_request = msg
                groupctx.reply_channel = reply_channel

                # init sending the signal to any processes belonging to the group
                issued_pendings = False
                for lst_idx, lst in enumerate(groupdesc.sets):
                    for item_idx, item in enumerate(lst):
                        # if this is not a process, do nothing
                        if isinstance(item.desc, process_desc.ProcessDescriptor):
                            if item.state != process_desc.ProcessDescriptor.State.DEAD:
                                resource_msg = groupctx._mk_gs_proc_kill(item.uid, sig=msg.sig)
                                issued, outbound_tag = ProcessContext.kill(server, resource_msg, dutil.AbsorbingChannel())
                                pctx = server.process_table[item.uid]

                                if issued:
                                    # we need to issue a pending operation related to this process
                                    groupdesc.sets[lst_idx][item_idx].state = process_desc.ProcessDescriptor.State.PENDING
                                    groupdesc.sets[lst_idx][item_idx].desc.state = process_desc.ProcessDescriptor.State.PENDING
                                    server.pending[outbound_tag] = groupctx.complete_kill
                                    server.group_to_pending_resource_map[(outbound_tag, target_uid)] = pctx
                                    server.group_resource_count[target_uid][lst_idx] += 1
                                    if not issued_pendings:
                                        issued_pendings = True
                                else:
                                    # the process is either unknown or dead or pending
                                    # update the member's state accordingly
                                    groupdesc.sets[lst_idx][item_idx].state = pctx.descriptor.state
                                    groupdesc.sets[lst_idx][item_idx].desc.state = pctx.descriptor.state
                if not issued_pendings:
                    # there were no pending requests issued and we need to send a response to the client
                    # We consider this case as a success; for example, all the processes were already dead
                    rm = gsgkr(tag=server.tag_inc(),
                               ref=msg.tag,
                               err=gsgkr.Errors.ALREADY,
                               desc=groupdesc)
                    LOG.debug(f'sending kill response to request {msg}: {rm} -- all processes were already dead')
                    groupdesc.state = gds.ACTIVE
                    reply_channel.send(rm.serialize())
                return issued_pendings
            else:
                raise NotImplementedError('close case')

    def complete_kill(self, msg):
        """Completes the action of a group kill.

        :param msg: response message related to destroying a resource which is member to a group
        :type msg: dragon.infrastructure.messages.SHProcessKillResponse
        :raises GroupError: when the context corresponding to the pending resource cannot be found
        :raises NotImplementedError: when we receive an unknown response error related to a group's resource/member
        :return: True if killing the process of the group succeeded False otherwise.
        :rtype: bool
        """
        # TODO: implement for other types of resources apart from processes

        gsgkr = dmsg.GSGroupKillResponse
        shpkr = dmsg.SHProcessKillResponse

        if msg.err not in shpkr.Errors:
            raise NotImplementedError('close case')

        # get the resource context corresponding to this response message
        if (msg.ref, self.descriptor.g_uid) in self.server.group_to_pending_resource_map:
            proc_context = self.server.group_to_pending_resource_map.pop((msg.ref, self.descriptor.g_uid))
        else:
            raise GroupError('Error with pending process belonging to a group.')

        # this p_uid should exist in server.resource_to_group_map dict
        this_puid = proc_context.descriptor.p_uid

        guid, (lst_idx, item_idx) = self.server.resource_to_group_map[this_puid]

        # complete the kill of this process without sending any response messages to the client
        # as we've set the reply_channel of the process context to dutil.AbsorbingChannel() in kill() above
        succeeded = proc_context.complete_kill(msg)
        # TODO: should we do anything if succeeded is False, which means that the
        # SHProcessKillResponse was a FAIL?

        if not self.destroy_remove_success_ids:
            # update the member's state accordingly
            self.descriptor.sets[lst_idx][item_idx].state = proc_context.descriptor.state

        # if destroy_remove was called
        if self.destroy_remove_success_ids:
            if self.server.group_destroy_resource_count[guid]:
                # adjust the count since we received one more response msg
                del self.server.group_destroy_resource_count[guid][-1]

            # if we have received responses for all the members, we are ready to remove
            # the resources from the group and send the response back to the client
            if not self.server.group_destroy_resource_count[guid]:
                for puid, lst_idx, item_idx in self.destroy_remove_success_ids:
                    # if this is not already removed/dead
                    if puid in self.server.resource_to_group_map:
                        self._remove_proc_from_group(puid, lst_idx, item_idx)
                # Update the dict that holds the position of the members of the group
                self._update_ds_after_removals()

                self.destroy_remove_success_ids = None
                if self.descriptor.state == self.descriptor.State.PENDING:
                    self.descriptor.state = self.descriptor.State.ACTIVE
                rm = dmsg.GSGroupDestroyRemoveFromResponse(tag=self.server.tag_inc(),
                                                           ref=self.destroy_request.tag,
                                                           err=dmsg.GSGroupDestroyRemoveFromResponse.Errors.SUCCESS,
                                                           desc=self.descriptor)
                LOG.debug(f'sending remove_destroy response to request {self.destroy_request}: {rm}')
                self.reply_channel.send(rm.serialize())

                del self.server.group_destroy_resource_count[guid]

        # if this was called from destroy(), then we don't want to send a
        # kill response to the client or counting how many kill responses
        # we have received, as this is an intermediate step of destroy()
        if not self.destroy_called:
            if self.server.group_resource_count[guid][lst_idx] >= 1:
                self.server.group_resource_count[guid][lst_idx] -= 1

            # if we have received responses for all the members, we are ready
            # to send the group kill response back to the client
            # if we have received responses for all the processess of this list
            if self.server.group_resource_count[guid][lst_idx] == 0:
                # next, check if we have responses for the other lists as well
                if all(item == 0 for item in self.server.group_resource_count[guid]):
                    # we follow the same pattern with groups as we do with processes
                    if self.descriptor.state == self.descriptor.State.PENDING:
                        self.descriptor.state = self.descriptor.State.ACTIVE

                    rm = gsgkr(tag=self.server.tag_inc(),
                               ref=self.destroy_request.tag,
                               err=gsgkr.Errors.SUCCESS,
                               desc=self.descriptor)
                    LOG.debug(f'sending kill response to request {self.destroy_request}: {rm}')
                    self.reply_channel.send(rm.serialize())

        return succeeded

    @staticmethod
    def destroy(server, msg, reply_channel):
        """Action to destroy this group of resources.
        We need to destroy all the members of the group as well as the container.

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: destroy request message
        :type msg: dragon.infrastructure.messages.GSGroupDestroy
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :raises NotImplementedError: when the group's state is unknown, i.e., other than ACTIVE, DEAD, PENDING
        :return: True if a successful destroy message was issued and False otherwise
        :rtype: bool
        """
        target_uid, found, errmsg = server.resolve_guid(msg.user_name, msg.g_uid)
        gsgdr = dmsg.GSGroupDestroyResponse

        if not found:
            rm = gsgdr(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gsgdr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown group of resources: response to {msg}: {rm}')
            return False
        else:
            groupctx = server.group_table[target_uid]
            groupdesc = groupctx.descriptor
            gds = group_desc.GroupDescriptor.State

            if gds.DEAD == groupdesc.state:
                rm = gsgdr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgdr.Errors.DEAD)
                reply_channel.send(rm.serialize())
                LOG.debug(f'group dead, response to {msg}: {rm}')
                return False

            # this is a highly unlikely case to happen
            # because the user must have received a create response back
            # before requesting destroying this group, which means that the
            # state of the group is not pending anymore
            # so, what we really want here is to log this is happening and move on
            elif gds.PENDING == groupdesc.state:
                rm = gsgdr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gsgdr.Errors.PENDING,
                           err_info=f'group {target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'group pending while destroy request -- this should not be happening, response to {msg}: {rm}')
                return False
            elif gds.ACTIVE == groupdesc.state:
                # update the group's state to PENDING as long as the destroy is pending
                groupdesc.state = gds.PENDING

                groupctx.destroy_request = msg
                groupctx.reply_channel = reply_channel

                server.group_destroy_resource_count[target_uid] = []

                # init the destruction of the group's members
                for lst_idx, lst in enumerate(groupdesc.sets):
                    server.group_destroy_resource_count[target_uid].append(0)

                    for item_idx, item in enumerate(lst):

                        # TODO: implement for other types of resources apart from processes
                        # for other types we don't need to call kill first

                        # if the process is not already dead, proceed
                        if item.state != process_desc.ProcessDescriptor.State.DEAD:
                            # first, call kill to send a kill request
                            resource_msg = groupctx._mk_gs_proc_kill(item.uid, sig=signal.SIGKILL)
                            issued, outbound_tag = ProcessContext.kill(server, resource_msg, dutil.AbsorbingChannel())

                            if issued:
                                # there has been a pending operation related to this process
                                groupdesc.sets[lst_idx][item_idx].state = process_desc.ProcessDescriptor.State.PENDING
                                groupdesc.sets[lst_idx][item_idx].desc.state = process_desc.ProcessDescriptor.State.PENDING

                                # this is needed for identification after GS having received SHProcessExit msg
                                # because the SHProcessExit msg will have a different tag
                                server.pending_group_destroy[item.uid] = groupctx.complete_destruction
                                server.group_to_pending_resource_map[(outbound_tag, target_uid)] = server.process_table[item.uid]

                                # this applies only to the case where processes are members of the group
                                groupctx.destroy_called = True
                                server.pending[outbound_tag] = groupctx.complete_kill

                                server.group_destroy_resource_count[target_uid][lst_idx] += 1
                            else:
                                # the process is either unknown or dead or pending
                                groupdesc.sets[lst_idx][item_idx].state = process_desc.ProcessDescriptor.State.DEAD
                                groupdesc.sets[lst_idx][item_idx].desc.state = process_desc.ProcessDescriptor.State.DEAD
                # in this case, all the processes were already dead or no pending continuation
                # was issued and we need to send a response to the client
                # e.g., kill was called prior to destroy and the kill request completed successully (SHProcessExit was sent)
                if all(item == 0 for item in server.group_destroy_resource_count[target_uid]):
                    groupdesc.state = gds.DEAD
                    if groupctx.pmi_job_helper:
                        groupctx.pmi_job_helper.cleanup()
                    rm = gsgdr(tag=server.tag_inc(),
                               ref=msg.tag,
                               err=gsgdr.Errors.SUCCESS,
                               desc=groupdesc)
                    reply_channel.send(rm.serialize())
                    LOG.debug(f'sending destroy response to request {msg}: {rm}')

                    # no need to keep this entry since the group is destroyed
                    del server.group_destroy_resource_count[target_uid]

                return True
            else:
                raise NotImplementedError('close case')

    def complete_destruction(self, msg):
        """Completes the destruction of a group of resources.
        When dealing with processes as members of the group, at this point, we assume that
        the kill request is being completed.

        :param msg: response message related to destroying a resource which is member to a group
        :type msg: dragon.infrastructure.messages.SHProcessExit
        :raises GroupError: when the context corresponding to the pending resource cannot be found
        :raises NotImplementedError: when we receive an unknown response error related to a group's resource/member
        """
        # TODO: implement for other types of resources apart from processes

        if not isinstance(msg, dmsg.SHProcessExit):
            raise GroupError(f'Received msg type {msg} - expected an SHProcessExit msg.')

        gsgdr = dmsg.GSGroupDestroyResponse

        # get the resource context corresponding to this response message
        if (msg.p_uid, self.descriptor.g_uid) in self.server.group_to_pending_resource_map:
            proc_context = self.server.group_to_pending_resource_map.pop((msg.p_uid, self.descriptor.g_uid))
        else:
            raise GroupError('Error with group pending resource.')

        # this p_uid should exist in server.resource_to_group_map dict
        this_puid = proc_context.descriptor.p_uid

        guid, (lst_idx, item_idx) = self.server.resource_to_group_map[this_puid]

        if self.server.group_destroy_resource_count[guid][lst_idx] >= 1:
            self.server.group_destroy_resource_count[guid][lst_idx] -= 1

        del self.server.resource_to_group_map[this_puid]

        # if we have received responses for all the members, we are ready
        # to send the group destroy response to the client
        # if we have received responses for all the members of this list
        if self.server.group_destroy_resource_count[guid][lst_idx] == 0:
            # next, check if we have responses for the other lists as well
            if all(item == 0 for item in self.server.group_destroy_resource_count[guid]):
                self.descriptor.state = self.descriptor.State.DEAD
                if self.pmi_job_helper:
                    self.pmi_job_helper.cleanup()
                rm = gsgdr(tag=self.server.tag_inc(),
                           ref=self.destroy_request.tag,
                           err=gsgdr.Errors.SUCCESS,
                           desc=self.descriptor)
                LOG.debug(f'sending destroy response to request {self.destroy_request}: {rm}')
                self.reply_channel.send(rm.serialize())

                # no need to keep this entry since the group is destroyed
                del self.server.group_destroy_resource_count[guid]

        return

src/dragon/globalservices/group_int.py
DragonHPC/dragon
"""Global Services API for the Node primary resource component.

The 'Node' component abstracts the hardware on a node level, providing
information like the number of CPUs or the amount of memory available on the
node.
"""

import enum
import logging

from .api_setup import get_gs_ret_cuid, next_tag, gs_request

from ..infrastructure.messages import \
    GSNodeList, \
    GSNodeListResponse, \
    GSNodeQuery, \
    GSNodeQueryResponse, \
    GSNodeQueryTotalCPUCount, \
    GSNodeQueryTotalCPUCountResponse

from ..infrastructure.parameters import this_process
from ..infrastructure.node_desc import NodeDescriptor

LOG = logging.getLogger(__name__)


class DragonNodeError(Exception):
    pass


def query(identifier) -> NodeDescriptor:
    """ Asks GS for the node descriptor beloging to a node identified by name or h_uid .

    :param identifier: Node id to query. This is the name if of type str, h_uid if of type int
    :return: descriptor
    :raises:
    """

    if isinstance(identifier, int):  # by h_uid
        req_msg = GSNodeQuery(tag=next_tag(), p_uid=this_process.my_puid,
                              r_c_uid=get_gs_ret_cuid(), h_uid=identifier)
    else:  # by name
        req_msg = GSNodeQuery(tag=next_tag(), p_uid=this_process.my_puid,
                              r_c_uid=get_gs_ret_cuid(), name=identifier)

    reply_msg = gs_request(req_msg)
    assert isinstance(reply_msg, GSNodeQueryResponse)

    if GSNodeQueryResponse.Errors.SUCCESS == reply_msg.err:
        the_desc = reply_msg.desc

    return the_desc


def query_total_cpus() -> int:
    """ Asks GS to return the total number of CPUS beloging to all of the registered nodes.

    :return: integer value of the number of cpus
    :raises:
    """
    total_cpus = 0

    req_msg = GSNodeQueryTotalCPUCount(tag=next_tag(), p_uid=this_process.my_puid,
                                       r_c_uid=get_gs_ret_cuid())

    reply_msg = gs_request(req_msg)
    assert isinstance(reply_msg, GSNodeQueryTotalCPUCountResponse)

    if GSNodeQueryTotalCPUCountResponse.Errors.SUCCESS == reply_msg.err:
        total_cpus = reply_msg.total_cpus

    return total_cpus


def get_list() -> list[int]:
    """ Asks Global Services for a list of the h_uids of all nodes.

    :return: list of the  h_uids of all currently existing memory pools
    """
    req_msg = GSNodeList(tag=next_tag(), p_uid=this_process.my_puid,
                         r_c_uid=get_gs_ret_cuid())

    reply_msg = gs_request(req_msg)
    assert isinstance(reply_msg, GSNodeListResponse)

    return reply_msg.hlist


# def discover():
# def add():
# def remove():

src/dragon/globalservices/node.py
DragonHPC/dragon
""" Node related objects internal to Global Services. 
"""

from ..infrastructure.node_desc import NodeDescriptor

import logging

LOG = logging.getLogger(__name__)

# NOTE: This is a stub really and will need to be extended once we are 
# targetting cloud hardware.
class NodeContext:
    """ Everything to do with a single node in global services.

        In contrast to other resources handled by GS, the NodeContext
        gets its descriptors from the Shepherd started on the node via
        SHPingGS.

        This class includes the code that does the lifecycle of the 
        node to the launcher controlling that channel.
    """ 

    def __init__(self, desc, ls_index, h_uid, name):

        self.h_uid = h_uid
        self.name = name
        self.ls_index = ls_index
        self._descriptor = desc 
        self._descriptor.h_uid = h_uid

    def __str__(self):
        return f"[[{self.__class__.__name__}]] desc:{self.descriptor!r} idx:{self.ls_index!r}"

    @property
    def descriptor(self):
        return self._descriptor

    @classmethod
    def construct(cls, msg):
        """Registers a new node with Global Services, by constructing the appropriate
        context from the message.

        :param msg: Message send to GS to create the node
        :type msg: class SHPingGS 
        :return: The new node context
        :rtype: class NodeContext
        """

        ls_index = msg.idx # index of the node according to Local Services
        desc = NodeDescriptor.from_sdict(msg.node_sdesc)

        h_uid = desc.host_id
            
        name = desc.name

        context = cls(desc, ls_index, h_uid, name)

        return context
src/dragon/globalservices/node_int.py
DragonHPC/dragon
import enum
from dataclasses import dataclass, field, fields
from ..infrastructure.policy import Policy
from ..infrastructure.node_desc import NodeDescriptor

@dataclass
class ResourceLayout:
    h_uid : int
    host_name : str
    numa_node : int # TODO
    core : [int] # List of acceptable CPU cores the layout can be applied to
    accelerator : int # TODO


# Possible TODO: Affinity offset or placement offset, to "start counting from here", just make sure it starts at a certain place
class PolicyEvaluator:
    """
    Based on a list of NodeDescriptors, evaluate policies and apply them
    """

    def __init__(self, nodes : list[NodeDescriptor], default_policy=None):
        self.nodes = nodes
        self.num_nodes = len(nodes)
        self.default_policy = default_policy # This will be used to fill in any Policy values that are DEFAULT when evaluated
        self.overprovision = False # Should make this into a parameter flag to allow/disallow it?
        self.cur_node = 0


    def _find_next_open(self, cur_idx):
        """
        Find next node with an open policy slot
        If no nodes are available, return None to indicate overprovisioning mode

        :param cur_idx: Current index in the list of nodes
        """

        idx = (cur_idx + 1) % self.num_nodes
        node = None
        if self.num_nodes == 1:
            return node # Indicate we are now overprovisioning mode, we only have one node

        # Loop through nodes until we hit our starting point
        while idx != cur_idx:
            next_node = self.nodes[idx]
            if next_node.num_policies < next_node.num_cpus:
                node = next_node
                break

            idx = (idx + 1) % self.num_nodes

        return node

    def _add_layout(self, node, affinity, layouts):
        """
        Short helper function to auto-increment num_policies value on nodes

        :param node: Node we're generating the layout for
        :param affinity: List of CPU devices local services can apply to
        :param layouts: List of layouts to append to

        """
        # NOTE: Numa node and accelerator are placeholders
        layouts.append( ResourceLayout(node.h_uid, node.host_name, 1, affinity, 0) )
        node.num_policies += 1

    def _node_by_id(self, host_id) -> NodeDescriptor:
        """
        Try to find a node with a given host ID
        Raise exception if specified host is not in the list
        """
        node = None
        for n in self.nodes:
            if n.h_uid == host_id:
                node = n
                break

        if node is None:
            raise ValueError(f"Could not find host with specified ID {host_id} in policy")
        return node

    def _node_by_name(self, host_name) -> NodeDescriptor:
        """
        Try to find a node with a given host name
        Raise exception if specified host is not in the list
        """
        node = None
        for n in self.nodes:
            if n.host_name == host_name:
                node = n
                break

        if node is None:
            raise ValueError(f"Could not find host with specified name {host_name} in policy")

        return node

    def _get_node(self, p : Policy) -> NodeDescriptor:
        """
        Find the next available node based on the provided policy
        """

        # Reset to the first node if we've cycled all the way through
        if self.cur_node >= self.num_nodes:
            self.cur_node = 0

        # If placement is specified to a host ID or name, apply to that node no questions asked
        if p.placement == Policy.Placement.HOST_ID:
            return self._node_by_id(p.host_id)
        elif p.placement == Policy.Placement.HOST_NAME:
            return self._node_by_name(p.host_name)

        else:
            # TODO: Other placements not used in any meaningful way yet
            # Default to whatever the next node based on the distribution policy is
            # if p.placement == default, local, anywhere ...
            node = self.nodes[self.cur_node]

        # Default to roundrobin for now
        distribution = p.distribution # Use a temp value instead of modifying the policy
        if distribution == Policy.Distribution.DEFAULT:
            distribution = Policy.Distribution.ROUNDROBIN

        # RoundRobin defaults to always overprovisioning, it will always just assign a policy to "next node"
        # TODO: Mimic some overprovisioning logic like in BLOCK for roundrobin to find the next empty slot?
        if distribution == Policy.Distribution.ROUNDROBIN:
            self.cur_node += 1
            return node

        elif distribution == Policy.Distribution.BLOCK:
            if self.overprovision:
                self.cur_node += 1 # Basically means we're in round-robin

            elif node.num_policies >= node.num_cpus:
                next_node = self._find_next_open(self.cur_node) # Find next open node
                if next_node is None: # All nodes full
                    self.overprovision = True # Start overprovisioning
                    next_node = node # Reassign next_node from None to current node
                    self.cur_node += 1 # Move to next node in cycle (round robin, essentially)
                node = next_node # Assign current node to "next node" (either open node, or same node to begin with)

            return node

    def _get_affinity(self, p : Policy, node : NodeDescriptor) -> list[int]:
        """
        Generate a list of available devices the policy can be applied to for the given Node
        """

        affinity = p.affinity
        if affinity == Policy.Affinity.DEFAULT:
            affinity = Policy.Affinity.ANY

        device = p.device
        if device == Policy.Device.DEFAULT:
            device = Policy.Device.CPU

        if device == Policy.Device.CPU:
            if p.specific_affinity: # List not empty, assume SPECIFIC affinity
                affinity = [x for x in node.cpu_devices if x in p.specific_affinity]
                return affinity # This covers both "ANY" and "SPECIFIC" if a specific list is given

            if affinity == Policy.Affinity.ANY:
                return node.cpu_devices

        if device == Policy.Device.GPU:
            raise RuntimeError("GPU Affinity Not Implemented")
        # TODO: Affinity ROUNDROBIN and BLOCK.  These are more load balancing options than not, so tracking usage needs to be figured out
        # TODO: GPU Affinity.  NodeDescriptor needs to pull a list of available accelerators (through an os.exec grep command maybe?)

        return []

    def evaluate(self, policies : list[Policy]=None) -> list[ResourceLayout]:
        """
        Evaluate a list of policies and return layouts of what nodes to apply them to
        """

        layouts = []

        # Iterate over policies and apply them through node descriptors
        for p in policies:
            # Merge incoming policies against the self.default_policy so any DEFAULT enums get replaced with the default policy option
            p = self.merge(self.default_policy, p)
            node = self._get_node(p) # Get a node based on policy (if requesting specific nodes, may raise exception)
            affinity = self._get_affinity(p, node) # Get affinity based on policy
            self._add_layout(node, affinity, layouts)

        return layouts

    @staticmethod
    def merge(high_policy: Policy, low_policy: Policy) -> Policy:
        """
        Merge two policies, using values from high_policy for values not assigned on init
        Returns a new policy

        :param high_policy: Non-default values take precidence
        :param low_policy: Default values will be replaced by `high_policy` values
        """

        retain = {}

        for f in fields(low_policy):
            v = getattr(low_policy, f.name)
            if v == f.default:
                continue # Default value, do not retain
            retain[f.name] = v # Value was assigned specifically, retain

        kwargs = {**high_policy.get_sdict(), **retain} # Merge retained policy values with high priority values
        return type(high_policy)(**kwargs)

src/dragon/globalservices/policy_eval.py
DragonHPC/dragon
"""Global Services API for the Managed Memory Pools and Allocations primary resource component.

The current API only manages memory pools, future versions will also contain
methods to track allocations.
"""
import logging

import dragon.infrastructure.messages as dmsg
from dragon.infrastructure.parameters import this_process
import dragon.globalservices.api_setup as das

log = logging.getLogger('pool_api')


class PoolError(Exception):
    pass


def create(size, user_name='', options=None, soft=False):
    """Asks Global Services to create a new  memory pool

        :param size: size of memory pool to create
        :param user_name: Requested user specified reference name
        :param options: PoolOptions object, what options to apply to creation
        :param soft: If pool already exists with given name, do not create and return descriptor instead.
        :return: PoolDescriptor object
    """

    if options is None:
        options = {}

    if soft and not user_name:
        raise PoolError('soft create requires a user supplied pool name')

    req_msg = dmsg.GSPoolCreate(tag=das.next_tag(), p_uid=this_process.my_puid,
                                r_c_uid=das.get_gs_ret_cuid(), size=size,
                                user_name=user_name, options=options)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSPoolCreateResponse)

    ec = dmsg.GSPoolCreateResponse.Errors

    if ec.SUCCESS == reply_msg.err:
        the_desc = reply_msg.desc
    else:
        if soft and ec.ALREADY == reply_msg.err:
            the_desc = reply_msg.desc
        else:
            raise PoolError(f'pool create {req_msg} failed: {reply_msg.err_info}')

    return the_desc


def get_list():
    """Asks Global Services for a list of the m_uids of all memory pools.

        TODO: add some options to this, and to the message itself, to have a finer
        control over which pools you get back.  Deferring until we have a use case

        :return: list of the  m_uids of all currently existing memory pools
        """
    req_msg = dmsg.GSPoolList(tag=das.next_tag(), p_uid=this_process.my_puid,
                              r_c_uid=das.get_gs_ret_cuid())

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSPoolListResponse)

    return reply_msg.mlist


def query(identifier):
    """Asks Global Services for the PoolDescriptor of a specified memory pool

    Note you can only query currently existing pools.

    :param identifier: string indicating pool name or integer indicating a m_uid
    :return: PoolDescriptor object corresponding to specified pool
    :raises: PoolError if there is no such pool
    """
    if isinstance(identifier, str):
        req_msg = dmsg.GSPoolQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
                                   r_c_uid=das.get_gs_ret_cuid(),
                                   user_name=identifier)
    else:
        req_msg = dmsg.GSPoolQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
                                   r_c_uid=das.get_gs_ret_cuid(),
                                   m_uid=int(identifier))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSPoolQueryResponse)

    if dmsg.GSPoolQueryResponse.Errors.SUCCESS == reply_msg.err:
        the_desc = reply_msg.desc
    else:
        raise PoolError(f'pool query {req_msg} failed: {reply_msg.err_info}')

    return the_desc


def destroy(identifier):
    """Asks Global Services to destroy a specified pool.

    TODO: figure out the semantics on this.  Do we want to throw an error from
    GS if there are any channels in that pool currently?

    :param identifier: string indicating pool name or integer indicating a m_uid
    :return: Nothing if successful
    :raises: PoolError if there is no such pool
    """
    if isinstance(identifier, str):
        req_msg = dmsg.GSPoolDestroy(tag=das.next_tag(), p_uid=this_process.my_puid,
                                     r_c_uid=das.get_gs_ret_cuid(),
                                     user_name=identifier)
    else:
        req_msg = dmsg.GSPoolDestroy(tag=das.next_tag(), p_uid=this_process.my_puid,
                                     r_c_uid=das.get_gs_ret_cuid(),
                                     m_uid=int(identifier))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSPoolDestroyResponse)

    ec = dmsg.GSPoolDestroyResponse.Errors

    if ec.SUCCESS == reply_msg.err or ec.GONE == reply_msg.err:
        return
    elif ec.UNKNOWN == reply_msg.err or ec.FAIL == reply_msg.err:
        raise PoolError(f'pool destroy {req_msg} failed: {reply_msg.err_info}')
    elif ec.PENDING == reply_msg.err:
        raise PoolError(f'pool destroy {req_msg} failed pending: {reply_msg.err_info}')
    else:
        raise NotImplementedError('close case')

src/dragon/globalservices/pool.py
DragonHPC/dragon
"""Global services' internal Pool context..
"""

import dragon.infrastructure.pool_desc as pool_desc
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.messages as dmsg

import logging

LOG = logging.getLogger('pool:')


class PoolContext:
    """Everything to do with a single memory pool in global services.


        This object manages all the transactions to a shepherd concerning
        the lifecycle of a pool.
    """

    def __init__(self, server, request, reply_channel, m_uid, node):
        self.server = server
        self.request = request
        self.reply_channel = reply_channel
        self.destroy_request = None

        self._descriptor = pool_desc.PoolDescriptor(node=node, m_uid=m_uid,
                                                    name=request.user_name)

    def __str__(self):
        return f"[[{self.__class__.__name__}]] desc:{self.descriptor!r} req:{self.request!r}"

    @property
    def descriptor(self):
        return self._descriptor

    def _mk_sh_pool_create(self, the_tag):
        return dmsg.SHPoolCreate(tag=the_tag,
                                 p_uid=dfacts.GS_PUID,
                                 r_c_uid=dfacts.GS_INPUT_CUID,
                                 m_uid=self.descriptor.m_uid,
                                 size=self.request.size,
                                 name=self.request.user_name,
                                 attr=self.request.options.sattr)

    def mk_sh_pool_destroy(self, the_tag):
        return dmsg.SHPoolDestroy(tag=the_tag,
                                  p_uid=dfacts.GS_PUID,
                                  r_c_uid=dfacts.GS_INPUT_CUID,
                                  m_uid=self.descriptor.m_uid)

    @classmethod
    def construct(cls, server, msg, reply_channel):
        """Creates a PoolContext and registers it with the server.

        This routine will send a message to the Shepherd and will leave
        a continuation in the server's pending messages table.

        Getting the reply channel as a parameter makes it easier
        to make this an internal construction...

        :return: a tuple:
                First element True if construction message issued, False otherwise
                Second element the tag if message issued, None otherwise
                Third element The newly constructed context object if issued, None otherwise
        """

        if msg.user_name in server.pool_names:
            LOG.info(f'pool name {msg.user_name} in use')
            existing_ctx = server.pool_table[server.pool_names[msg.user_name]]
            rm = dmsg.GSPoolCreateResponse(tag=server.tag_inc(),
                                           ref=msg.tag,
                                           err=dmsg.GSPoolCreateResponse.Errors.ALREADY,
                                           desc=existing_ctx.descriptor)
            reply_channel.send(rm.serialize())
            return False, None, None

        this_muid, auto_name = server.new_muid_and_default_name()

        if not msg.user_name:
            msg.user_name = auto_name

        which_node = server.choose_pool_node(msg)

        context = cls(server, msg, reply_channel, this_muid, which_node)

        server.pool_names[msg.user_name] = this_muid
        server.pool_table[this_muid] = context

        shep_hdl = server.shep_inputs[which_node]

        outbound_tag = server.tag_inc()
        shep_create_msg = context._mk_sh_pool_create(outbound_tag)

        # In cases where we are sending a large amount of
        # messages, such as with the GSGroupCreate handler,
        # we can fill the GS Input Queue with responses and
        # basically cause the GS / TA / LS to be unable to
        # send/receive any messages. To prevent this, we'll
        # enqueue pending sends and interleave sending and
        # receiving messages to allow us to process responses
        # on the input queue.

        server.pending_sends.put((shep_hdl, shep_create_msg.serialize()))
        return True, outbound_tag, context

    def complete_construction(self, msg):
        """Completes construction of a PoolContext

            :return: True if it succeeded, False otherwise
        """
        assert isinstance(msg, dmsg.SHPoolCreateResponse)

        if msg.Errors.SUCCESS == msg.err:
            self.descriptor.state = pool_desc.PoolDescriptor.State.ACTIVE
            self.descriptor.sdesc = msg.desc

            response = dmsg.GSPoolCreateResponse(tag=self.server.tag_inc(),
                                                 ref=self.request.tag,
                                                 err=dmsg.GSPoolCreateResponse.Errors.SUCCESS,
                                                 desc=self.descriptor)
            self.reply_channel.send(response.serialize())

            create_succeeded = True

        elif msg.Errors.FAIL == msg.err:
            self.descriptor.state = pool_desc.PoolDescriptor.State.DEAD
            response = dmsg.GSPoolCreateResponse(tag=self.server.tag_inc(),
                                                 ref=self.request.tag,
                                                 err=dmsg.GSPoolCreateResponse.Errors.FAIL,
                                                 err_info=msg.err_info)

            # we don't keep descriptor for things that were never alive
            del self.server.pool_table[self.descriptor.m_uid]
            del self.server.pool_names[self.request.user_name]

            self.reply_channel.send(response.serialize())
            create_succeeded = False
        else:
            raise RuntimeError(f'got {str(msg)} err {msg.err} unknown')

        LOG.debug(f'create response sent, tag {response.tag} ref {response.ref} pending cleared')

        self.reply_channel = None  # drop reference to reply channel obj

        return create_succeeded

    @staticmethod
    def destroy(server, msg, reply_channel):
        """Handle request to destroy a context.

        Feature we lack but that we might want:

        Since it's a pool, we might want to maintain a table of which channels
        belong to which pools and error if there are any live channels
        currently in the pool being requested to be destroyed.

        :param server:
        :param msg: a GSPoolDestroyResponse
        :param reply_channel: object one can send the response to
        :return: True if a destroy request got issued False otherwise.
        """
        target_uid, found, errmsg = server.resolve_muid(msg.user_name, msg.m_uid)
        gspdr = dmsg.GSPoolDestroyResponse

        if not found:
            rm = gspdr(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gspdr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'unknown: response to {msg}: {rm}')
            return False
        else:
            poolctx = server.pool_table[target_uid]
            pooldesc = poolctx.descriptor
            pds = pool_desc.PoolDescriptor.State
            if pds.DEAD == pooldesc.state:
                rm = gspdr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gspdr.Errors.GONE)
                reply_channel.send(rm.serialize())
                LOG.debug(f'gone; response to {msg}: {rm}')
                return False
            elif pds.PENDING == pooldesc.state:
                rm = gspdr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gspdr.Errors.PENDING,
                           err_info=f'{target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'pending; response to {msg}: {rm}')
                return False
            elif pds.ACTIVE == pooldesc.state:
                the_tag = server.tag_inc()
                poolctx.reply_channel = reply_channel
                poolctx.destroy_request = msg
                poolctx.descriptor.state = pds.PENDING  # might want a PENDING_DESTRUCTION state?
                shep_destroy_msg = poolctx.mk_sh_pool_destroy(the_tag)
                target_node = pooldesc.node
                server.pending[the_tag] = poolctx.complete_destruction
                server.pending_sends.put((server.shep_inputs[target_node], shep_destroy_msg.serialize()))

                LOG.debug(f'forwarded destroy {msg} as {shep_destroy_msg} to {target_node}')
                return True
            else:
                raise NotImplementedError('bad case')

    def complete_destruction(self, msg):
        """

        :param msg:
        :return: True if the destroy succeeded False otherwise.
        """
        gspdr = dmsg.GSPoolDestroyResponse
        shpdr = dmsg.SHPoolDestroyResponse

        if shpdr.Errors.FAIL == msg.err:
            rm = gspdr(tag=self.server.tag_inc(),
                       ref=self.destroy_request.tag,
                       err=gspdr.Errors.FAIL,
                       err_info=msg.err_info)
            destroy_succeeded = False
        elif shpdr.Errors.SUCCESS == msg.err:
            self.descriptor.state = pool_desc.PoolDescriptor.State.DEAD
            rm = gspdr(tag=self.server.tag_inc(), ref=self.destroy_request.tag,
                       err=gspdr.Errors.SUCCESS)
            destroy_succeeded = True
        else:
            raise NotImplementedError('no case')

        LOG.debug(f'pool destroy response to {self.destroy_request}: {rm}')
        self.reply_channel.send(rm.serialize())
        return destroy_succeeded

src/dragon/globalservices/pool_int.py
DragonHPC/dragon
"""Global Services API for the Process primary resource component.
"""
import signal
import logging
import enum
import sys
import threading
import atexit
from distutils.util import strtobool

from .. import channels as dch

from ..globalservices import api_setup as das
from ..globalservices import channel as dgchan

from ..infrastructure import messages as dmsg
from ..infrastructure import parameters as dparm
from ..infrastructure.parameters import this_process

from ..infrastructure import process_desc as pdesc
from ..infrastructure import facts as dfacts
from ..infrastructure import connection as dconn
from ..utils import B64
import os

log = logging.getLogger('process_api')


@enum.unique
class StreamDestination(enum.Enum):
    STDIN = 0
    STDOUT = 1
    STDERR = 2

_capture_stdout_conn = None
_capture_stderr_conn = None
_capture_stdout_chan = None
_capture_stderr_chan = None

DRAGON_CAPTURE_MP_CHILD_OUTPUT = 'DRAGON_CAPTURE_MP_CHILD_OUTPUT'
DRAGON_STOP_CAPTURING_MP_CHILD_OUTPUT = 'DRAGON_STOP_CAPTURING_MP_CHILD_OUTPUT'
_capture_shutting_down = False

# This is for user processes to use when they wish to capture the
# output of child multiprocessing processes.
def start_capturing_child_mp_output():
    os.environ[DRAGON_CAPTURE_MP_CHILD_OUTPUT] = 'True'

def stop_capturing_child_mp_output():
    global _capture_shutting_down, _capture_stdout_chan, _capture_stderr_chan

    _capture_shutting_down = True

    stdout = dconn.Connection(outbound_initializer=_capture_stdout_chan,
                options=dconn.ConnectionOptions(min_block_size=512), policy=dparm.POLICY_INFRASTRUCTURE)

    stderr = dconn.Connection(outbound_initializer=_capture_stderr_chan,
                options=dconn.ConnectionOptions(min_block_size=512), policy=dparm.POLICY_INFRASTRUCTURE)

    stdout.send(DRAGON_STOP_CAPTURING_MP_CHILD_OUTPUT)
    stderr.send(DRAGON_STOP_CAPTURING_MP_CHILD_OUTPUT)

    if _capture_stdout_chan is not None:
        try:
            log.info('destroying the MP child capturing stdout channel')
            dgchan.destroy(_capture_stdout_chan.cuid)
            log.info('stdout channel destroy complete')
            _capture_stdout_chan = None
        except:
            pass

    if _capture_stderr_chan is not None:
        try:
            log.info('destroying the MP child capturing stderr channel')
            dgchan.destroy(_capture_stderr_chan.cuid)
            log.info('stderr channel destroy complete')
            _capture_stderr_chan = None
        except:
            pass


def _capture_child_output_setup():
    global _capture_stdout_conn, _capture_stderr_conn, _capture_stdout_chan, _capture_stderr_chan

    default_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)

    # ask GS to create a dragon channel
    stdout_descriptor = dgchan.create(default_muid, user_name=f'stdout_channel_{os.getpid()}')
    _capture_stdout_chan = dch.Channel.attach(stdout_descriptor.sdesc)
    os.environ[dfacts.STDOUT_DESC] = B64.bytes_to_str(stdout_descriptor.sdesc)
    _capture_stdout_conn = dconn.Connection(inbound_initializer=_capture_stdout_chan,
                            options=dconn.ConnectionOptions(min_block_size=512), policy=dparm.POLICY_INFRASTRUCTURE)


    stderr_descriptor = dgchan.create(default_muid, user_name=f'stderr_channel_{os.getpid()}')
    _capture_stderr_chan = dch.Channel.attach(stderr_descriptor.sdesc)
    os.environ[dfacts.STDERR_DESC] = B64.bytes_to_str(stderr_descriptor.sdesc)
    _capture_stderr_conn = dconn.Connection(inbound_initializer=_capture_stderr_chan,
                            options=dconn.ConnectionOptions(min_block_size=512), policy=dparm.POLICY_INFRASTRUCTURE)

    def forward(strm_conn:dconn.Connection, strm_dest:StreamDestination):
        file_dest = sys.stdout if strm_dest == StreamDestination.STDOUT else sys.stderr

        try:
            while not _capture_shutting_down:
                data = strm_conn.recv()
                if _capture_shutting_down and data == DRAGON_STOP_CAPTURING_MP_CHILD_OUTPUT:
                    return

                print(data, file=file_dest, flush=True, end="")
        except EOFError:
            pass
        except Exception as ex:
            log.debug(f'While forwarding output to {strm_dest} an unexpected error occurred:{repr(ex)}')

    def cleanup():
        try:
            stop_capturing_child_mp_output()
        except:
            pass

    stdout_monitor = threading.Thread(name='stdout monitor', target=forward, args=(_capture_stdout_conn, StreamDestination.STDOUT))
    stdout_monitor.start()

    stderr_monitor = threading.Thread(name='stderr monitor', target=forward, args=(_capture_stderr_conn, StreamDestination.STDERR))
    stderr_monitor.start()

    atexit.register(cleanup)



PIPE = dmsg.PIPE
STDOUT = dmsg.STDOUT
DEVNULL = dmsg.DEVNULL

class ProcessError(Exception):
    pass


def _create_stdio_connections(the_desc):

    if the_desc.stdin_sdesc is not None:
        # If a stdin serialized descriptor was requested, then a channel
        # has been created and the initial refcnt was 2. This keeps it from
        # being destroyed until the user has a chance to write to it even
        # if the process exits. Then once we create a Connection, the refcnt
        # gets incremented to three, so we decrement the refcnt to 2 again,
        # one for when this connection is closed, and one for when the process
        # exits and local services closes its end of the connection.
        ch = dch.Channel.attach(B64.from_str(the_desc.stdin_sdesc).decode())
        conn = dconn.Connection(outbound_initializer=ch, options=dconn.ConnectionOptions(min_block_size=512,
                            creation_policy=dconn.ConnectionOptions.CreationPolicy.PRE_CREATED), policy=dparm.POLICY_INFRASTRUCTURE)
        the_desc.stdin_conn = conn
        dgchan.release_refcnt(ch.cuid)

    if the_desc.stdout_sdesc is not None:
        # If a stdout serialized descriptor was requested, then a channel
        # has been created and the initial refcnt was 2 as was done for stdin
        # (please read the comment above). The analagous behaviour is coded
        # here.
        ch = dch.Channel.attach(B64.from_str(the_desc.stdout_sdesc).decode())
        conn = dconn.Connection(inbound_initializer=ch, options=dconn.ConnectionOptions(min_block_size=512,
                            creation_policy=dconn.ConnectionOptions.CreationPolicy.PRE_CREATED), policy=dparm.POLICY_INFRASTRUCTURE)
        the_desc.stdout_conn = conn
        dgchan.release_refcnt(ch.cuid)

    if the_desc.stderr_sdesc is not None:
        # If a stderr serialized descriptor was requested, then a channel
        # has been created and the initial refcnt was 2 as was done for stdin
        # (please read the comment above). The analagous behaviour is coded
        # here.
        ch = dch.Channel.attach(B64.from_str(the_desc.stderr_sdesc).decode())
        conn = dconn.Connection(inbound_initializer=ch, options=dconn.ConnectionOptions(min_block_size=512,
                            creation_policy=dconn.ConnectionOptions.CreationPolicy.PRE_CREATED), policy=dparm.POLICY_INFRASTRUCTURE)
        the_desc.stderr_conn = conn
        dgchan.release_refcnt(ch.cuid)

    return the_desc


def get_create_message(exe, run_dir, args, env, user_name='', options=None,
           stdin=None, stdout=None, stderr=None, group=None,
           user=None, umask=- 1, pipesize=- 1, pmi_required=False):
    """Return a GSProcessCreate object.

    :param exe: executable to run
    :param run_dir: directory to run it in
    :param args: argv for the process
    :param env: environment for the new process, may be added to for infrastructure
    :param user_name: Requested user specified reference name
    :param options: ProcessOptions object, what options to apply to creation
    :param stdin: If `stdin=PIPE`, a Dragon Connection will be returned to enable sending data to the stdin of the child process.
    :param stdout: Capture stdout from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stdout. Using DEVNULL will cause the stdout data to be thrown away.
    :param stder: Capture stderr from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stderr. Using DEVNULL will cause the stdout data to be thrown away. Using STDOUT will cause the stderr data to be combined with the stdout data.
    :param group: Not used
    :param user: Not used
    :param umask: Not used
    :param pipesize: Set the channel capacity. Default = -1.
    :param pmi_required: This process is part of a Dragon managed MPI/PMI application group.
    :return: GSProcessCreate message object
    """

    # Like the OS, we should inherit the environment when creating a process.
    # Using dict below makes a copy.
    the_env = dict(os.environ)
    if env is not None:
        the_env.update(env) # layer on the user supplied environment.

    log.debug('creating GSProcessCreate')
    return dmsg.GSProcessCreate(tag=das.next_tag(),
                                p_uid=this_process.my_puid,
                                r_c_uid=das.get_gs_ret_cuid(),
                                exe=exe, args=args, env=the_env,
                                rundir=run_dir, user_name=user_name,
                                options=options, stdin=stdin, stdout=stdout,
                                stderr=stderr, group=group, user=user, umask=umask,
                                pipesize=pipesize, pmi_required=pmi_required)


def get_create_message_with_argdata(exe, run_dir, args, env, argdata=None, user_name='', options=None,
           stdin=None, stdout=None, stderr=None, group=None,
           user=None, umask=- 1, pipesize=- 1, pmi_required=False):
    """Return a GSProcessCreate object with starting args.

        This is an extension of the 'get_create_message' method that encapsulates our scheme for getting
        starting arguments to the new process.  This depends on how big the arguments
        turn out to be, and is tightly bound with logic in global services and in the api_setup module.

    :param exe: executable to run
    :param run_dir: directory to run it in
    :param args: argv for the process
    :param env: environment for the new process, may be added to for infrastructure
    :param argdata: bytes to deliver
    :param user_name: Requested user specified reference name
    :param options: ProcessOptions object, what options to apply to creation
    :param stdin: If `stdin=PIPE`, a Dragon Connection will be returned to enable sending data to the stdin of the child process.
    :param stdout: Capture stdout from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stdout. Using DEVNULL will cause the stdout data to be thrown away.
    :param stder: Capture stderr from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stderr. Using DEVNULL will cause the stdout data to be thrown away. Using STDOUT will cause the stderr data to be combined with the stdout data.
    :param group: Not used
    :param user: Not used
    :param umask: Not used
    :param pipesize: Set the channel capacity. Default = -1.
    :param pmi_required: This process is part of a Dragon managed MPI/PMI application group.
    :return: GSProcessCreate message object
    """

    if options is None:
        options = pdesc.ProcessOptions(make_inf_channels=True)
    else:
        options.make_inf_channels = True

    # Like the OS, we should inherit the environment when creating a process.
    # Using dict below makes a copy.
    the_env = dict(os.environ)
    if env is not None:
        the_env.update(env) # layer on the user supplied environment.

    log.debug('creating GSProcessCreate')

    if argdata is None:
        # no arguments, so set the mode and start the normal way
        options.mode = pdesc.ArgMode.NONE
        options.argdata = None
        return dmsg.GSProcessCreate(tag=das.next_tag(),
                                p_uid=this_process.my_puid,
                                r_c_uid=das.get_gs_ret_cuid(),
                                exe=exe, args=args, env=the_env,
                                rundir=run_dir, user_name=user_name,
                                options=options, stdin=stdin, stdout=stdout,
                                stderr=stderr, group=group, user=user, umask=umask,
                                pipesize=pipesize, pmi_required=pmi_required)

    if len(argdata) <= dfacts.ARG_IMMEDIATE_LIMIT:
        # deliver arguments in the argdata directly.
        options.mode = pdesc.ArgMode.PYTHON_IMMEDIATE

        # todo: make a setter for options.argdata?
        options.argdata = prepare_argdata_for_immediate(argdata)
        return dmsg.GSProcessCreate(tag=das.next_tag(),
                                p_uid=this_process.my_puid,
                                r_c_uid=das.get_gs_ret_cuid(),
                                exe=exe, args=args, env=the_env,
                                rundir=run_dir, user_name=user_name,
                                options=options, stdin=stdin, stdout=stdout,
                                stderr=stderr, group=group, user=user, umask=umask,
                                pipesize=pipesize, pmi_required=pmi_required)
    else:
        raise NotImplementedError(f"Argument data larger than {dfacts.ARG_IMMEDIATE_LIMIT} bytes is not supported at the moment.")


def create(exe, run_dir, args, env, user_name='', options=None, soft=False,
           stdin=None, stdout=None, stderr=None, group=None,
           user=None, umask=- 1, pipesize=- 1, pmi_required=False):
    """Asks Global Services to create a new process.

    :param exe: executable to run
    :param run_dir: directory to run it in
    :param args: argv for the process
    :param env: environment for the new process, may be added to for infrastructure
    :param user_name: Requested user specified reference name
    :param options: ProcessOptions object, what options to apply to creation
    :param soft: If process already exists with given name, do not create and return descriptor instead.
    :param stdin: If `stdin=PIPE`, a Dragon Connection will be returned to enable sending data to the stdin of the child process.
    :param stdout: Capture stdout from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stdout. Using DEVNULL will cause the stdout data to be thrown away.
    :param stder: Capture stderr from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stderr. Using DEVNULL will cause the stdout data to be thrown away. Using STDOUT will cause the stderr data to be combined with the stdout data.
    :param group: Not used
    :param user: Not used
    :param umask: Not used
    :param pipesize: Set the channel capacity. Default = -1.
    :param pmi_required: This process is part of a Dragon managed MPI/PMI application group.
    :return: ProcessDescriptor object
    """

    # When a new process is requested we want to check to see if we are supposed to
    # redirect output from new processes to the parent process (this process). If so, we
    # set that up here and then remove that indicator from the env vars. This
    # check and call is done here to delay the setup as long as possible. This is
    # especially needed for something like Jupyter notebooks which do an execve and
    # need this setup done after that occurs.
    if bool(strtobool(os.environ.get(DRAGON_CAPTURE_MP_CHILD_OUTPUT, 'False'))):
        _capture_child_output_setup()
        del os.environ[DRAGON_CAPTURE_MP_CHILD_OUTPUT]

    if options is None:
        options = {}

    if soft and not user_name:
        raise ProcessError('soft create requires a user supplied process name')

    log.debug('creating GSProcessCreate')
    req_msg = get_create_message(
        exe=exe,
        run_dir=run_dir,
        args=args,
        env=env,
        user_name=user_name,
        options=options,
        stdin=stdin,
        stdout=stdout,
        stderr=stderr,
        group=group,
        user=user,
        umask=umask,
        pipesize=pipesize,
        pmi_required=pmi_required,
    )

    reply_msg = das.gs_request(req_msg)
    log.debug('got GSProcessCreateResponse')
    assert isinstance(reply_msg, dmsg.GSProcessCreateResponse)

    ec = dmsg.GSProcessCreateResponse.Errors

    if ec.SUCCESS == reply_msg.err:
        the_desc = reply_msg.desc
    else:
        if soft and ec.ALREADY == reply_msg.err:
            the_desc = reply_msg.desc
        else:
            raise ProcessError(f'process create {req_msg} failed: {reply_msg.err_info}')

    return _create_stdio_connections(the_desc)


def get_argdata_from_immediate_handshake(argdata_field):
    return B64.str_to_bytes(argdata_field)


def prepare_argdata_for_immediate(argdata):
    return B64.bytes_to_str(argdata)


def create_with_argdata(exe, run_dir, args, env, argdata=None, user_name='',
                        options=None, soft=False, pmi_required=False,
                        stdin=None, stdout=None, stderr=None):
    """Asks Global services to create a new process and deliver starting args to it thru messaging.

        This is an extension of the 'create' method that encapsulates our scheme for getting
        starting arguments to the new process.  This depends on how big the arguments
        turn out to be, and is tightly bound with logic in global services and in the api_setup module.

    :param exe: executable to run
    :param run_dir: directory to run it in
    :param args: argv for the process
    :param env: environment for the new process, may be added to for infrastructure
    :param argdata: bytes to deliver
    :param user_name: Requested user specified reference name
    :param options: ProcessOptions object, what options to apply to creation
    :param soft: If process already exists with given name, do not create and return descriptor instead.
    :param pmi_required: This process is part of a Dragon managed MPI/PMI application group.
    :param stdin: If `stdin=PIPE`, a Dragon Connection will be returned to enable sending data to the stdin of the child process.
    :param stdout: Capture stdout from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stdout. Using DEVNULL will cause the stdout data to be thrown away.
    :param stder: Capture stderr from the child process. If using PIPE, a Dragon Connection object will be returned which can be used to read from stderr. Using DEVNULL will cause the stdout data to be thrown away. Using STDOUT will cause the stderr data to be combined with the stdout data.
    :return: ProcessDescriptor object
    """

    if options is None:
        options = pdesc.ProcessOptions(make_inf_channels=True)
    else:
        options.make_inf_channels = True

    the_desc = None
    if argdata is None:
        # no arguments, so set the mode and start the normal way
        options.mode = pdesc.ArgMode.NONE
        options.argdata = None
        the_desc = create(
            exe=exe,
            run_dir=run_dir,
            args=args,
            env=env,
            user_name=user_name,
            options=options,
            soft=soft,
            pmi_required=pmi_required,
            stdin=stdin,
            stdout=stdout,
            stderr=stderr)

    elif len(argdata) <= dfacts.ARG_IMMEDIATE_LIMIT:
        # deliver arguments in the argdata directly.
        options.mode = pdesc.ArgMode.PYTHON_IMMEDIATE

        # todo: make a setter for options.argdata?
        options.argdata = prepare_argdata_for_immediate(argdata)
        the_desc = create(
            exe=exe,
            run_dir=run_dir,
            args=args,
            env=env,
            user_name=user_name,
            options=options,
            soft=soft,
            pmi_required=pmi_required,
            stdin=stdin,
            stdout=stdout,
            stderr=stderr)
    else:
        # TODO: capture these comments in documentation
        # Here we don't set argdata to anything at all.
        # GS will already be creating a channel local to the new process
        # for use in responding to that process's GS requests.
        # The ID of that channel is in the process descriptor that
        # is returned, so we will borrow that channel to send
        # the argument data.
        # There is no race, because the process will need to receive
        # its arguments before it can make any GS calls.
        # See globalservices.api_setup.py.

        options.mode = pdesc.ArgMode.PYTHON_CHANNEL
        options.argdata = None

        the_desc = create(
            exe=exe,
            run_dir=run_dir,
            args=args,
            env=env,
            user_name=user_name,
            options=options,
            soft=soft,
            pmi_required=pmi_required,
            stdin=stdin,
            stdout=stdout,
            stderr=stderr)

        # Another transaction to gs but we are only here if
        # we are sending a lot of data.  Could be returned with create call
        # optionally.  Don't want to put the whole gs ret channel descriptor
        # in the process descriptor because this is the only legit
        # use for that.
        gs_ret_desc = dgchan.query(the_desc.gs_ret_cuid)
        gsret = dch.Channel.attach(gs_ret_desc.sdesc)
        # TODO PE-38745
        arg_conn = dconn.Connection(outbound_initializer=gsret,
                                    options=dconn.ConnectionOptions(min_block_size=2**16),
                                    policy=dparm.POLICY_INFRASTRUCTURE)
        arg_conn.send_bytes(argdata)
        arg_conn.ghost_close()  # avoids sending an extra EOT message
        gsret.detach()

    return _create_stdio_connections(the_desc)


def get_list():
    """Asks Global Services for a list of the p_uids of all managed processes.

    This is the moral equivalent of a very simple unix 'ps' for managed processes.

    TODO: add some options to this, and to the message itself, to have a finer
    control over which processes you get back.

    :return: list of the p_uids of all processes, alive and dead
    """
    req_msg = dmsg.GSProcessList(tag=das.next_tag(), p_uid=this_process.my_puid,
                                 r_c_uid=das.get_gs_ret_cuid())

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSProcessListResponse)

    return reply_msg.plist


def query(identifier):
    """Asks Global Services for the ProcessDescriptor of a specified managed process

    Note you can query processes that have already exited; this is how one
    can find out the process's exit code.

    :param identifier: string indicating process name or integer indicating a p_uid
    :return: ProcessDescriptor object corresponding to specified managed process
    :raises: ProcessError if there is no such process
    """
    if isinstance(identifier, str):
        req_msg = dmsg.GSProcessQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
                                      r_c_uid=das.get_gs_ret_cuid(),
                                      user_name=identifier)
    else:
        req_msg = dmsg.GSProcessQuery(tag=das.next_tag(), p_uid=this_process.my_puid,
                                      r_c_uid=das.get_gs_ret_cuid(),
                                      t_p_uid=int(identifier))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSProcessQueryResponse)

    if dmsg.GSProcessQueryResponse.Errors.SUCCESS == reply_msg.err:
        the_desc = reply_msg.desc
    else:
        raise ProcessError(f'process query {req_msg} failed: {reply_msg.err_info}')

    return _create_stdio_connections(the_desc)


def kill(identifier, sig=signal.SIGKILL):
    """Asks Global Services to kill a specified managed process with a specified signal.

    Note that this is like the unix 'kill' command - the signal given to the process
    might not necessarily be intended to cause it to terminate.

    :param identifier: string indicating process name or integer indicating a p_uid
    :param sig: signal to use to kill the process, default=signal.SIGKILL
    :return: Nothing if successful
    :raises: ProcessError if there is no such process, or if the process has not yet started.
    """
    if isinstance(identifier, str):
        req_msg = dmsg.GSProcessKill(tag=das.next_tag(), p_uid=this_process.my_puid,
                                     r_c_uid=das.get_gs_ret_cuid(),
                                     user_name=identifier, sig=int(sig))
    else:
        req_msg = dmsg.GSProcessKill(tag=das.next_tag(), p_uid=this_process.my_puid,
                                     r_c_uid=das.get_gs_ret_cuid(),
                                     t_p_uid=int(identifier), sig=int(sig))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSProcessKillResponse)

    ec = dmsg.GSProcessKillResponse.Errors
    if ec.SUCCESS == reply_msg.err or ec.DEAD == reply_msg.err:
        return
    elif ec.UNKNOWN == reply_msg.err or ec.FAIL_KILL == reply_msg.err:
        raise ProcessError(f'process kill {req_msg} failed: {reply_msg.err_info}')
    elif ec.PENDING == reply_msg.err:
        raise ProcessError(f'process kill {req_msg} failed pending: {reply_msg.err_info}')
    else:
        raise NotImplementedError('close case')


def join(identifier, timeout=None):
    """Asks Global Services to join a specified managed process.

    Returns when the process has exited.


    :param identifier: string indicating process name or integer indicating a p_uid
    :param timeout: Timeout in seconds for max time to wait.  None = default, infinite wait
    :return: The unix exit code from the process, or None if there is a timeout.
    :raises: ProcessError if there is no such process or some other error has occurred.
    """

    if timeout is None:
        msg_timeout = -1
    elif timeout < 0:
        msg_timeout = 0
    else:
        msg_timeout = int(1000000 * timeout)

    if isinstance(identifier, str):
        req_msg = dmsg.GSProcessJoin(tag=das.next_tag(), p_uid=this_process.my_puid,
                                     r_c_uid=das.get_gs_ret_cuid(), timeout=msg_timeout,
                                     user_name=identifier)
    else:
        req_msg = dmsg.GSProcessJoin(tag=das.next_tag(), p_uid=this_process.my_puid,
                                     r_c_uid=das.get_gs_ret_cuid(), timeout=msg_timeout,
                                     t_p_uid=int(identifier))

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSProcessJoinResponse)

    ec = dmsg.GSProcessJoinResponse.Errors

    if ec.SUCCESS == reply_msg.err:
        return reply_msg.exit_code
    elif ec.TIMEOUT == reply_msg.err:
        return None
    elif ec.SELF == reply_msg.err:
        raise ProcessError(f'join to {identifier} is a deadlocking self-join: {reply_msg.err_info}')
    else:
        raise ProcessError(f'process join {req_msg} failed: {reply_msg.err_info}')


def multi_join(identifiers, timeout=None, join_all=False):
    """Asks Global Services to join a list of specified managed processes.

    If join_all is False, it returns when 'any' process has exited or there is a timeout.
    If join_all is True, it returns when 'all' processes have exited or there is a timeout.

    :param identifiers: list of process identifiers indicating p_uid (int) or process name (string)
    :param timeout: Timeout in seconds for max time to wait. None = default, infinite wait
    :param join_all: indicates whether we need to wait on all processes in the list or not
    :return: If join_all is False, return a list of tuples (p_uid, unix_exit_code) for any processes exited,
             or None if none exited and there is a timeout, along with a dictionary with the status of each process.
             If join_all is True, return a list of tuples (p_uid, unix_exit_code) when all processes exited,
             or None if none exited and there is a timeout or some exited and some errored/timed out,
             along with a dictionary with the status of each process.
    :raises: ProcessError if there is no such process or some other error has occurred.
    """

    if len(identifiers) == 0:
        raise ProcessError('multi_join attempted on zero processes. Empty multi_join is disallowed.')

    if timeout is None:
        msg_timeout = -1
    elif timeout < 0:
        msg_timeout = 0
    else:
        msg_timeout = int(1000000 * timeout)

    puid_identifiers = []
    name_identifiers = []
    for item in identifiers:
        if isinstance(item, str):
            name_identifiers.append(item)
        else:
            puid_identifiers.append(item)

    req_msg = dmsg.GSProcessJoinList(tag=das.next_tag(), p_uid=this_process.my_puid,
                                     r_c_uid=das.get_gs_ret_cuid(), timeout=msg_timeout,
                                     t_p_uid_list=puid_identifiers, user_name_list=name_identifiers,
                                     join_all=join_all)

    reply_msg = das.gs_request(req_msg)
    assert isinstance(reply_msg, dmsg.GSProcessJoinListResponse)

    ec = dmsg.GSProcessJoinListResponse.Errors # get the error codes

    success_list = []
    timeout_flag = False
    for t_p_uid, status_info in reply_msg.puid_status.items():
        if status_info[0] == ec.SUCCESS.value:
            success_list.append((int(t_p_uid), status_info[1]))
        elif status_info[0] == ec.TIMEOUT.value:
            if not timeout_flag:
                timeout_flag = True

    if success_list:
        if join_all: # 'all' option
            # we want all the processes finished, otherwise return None
            if len(success_list) == len(identifiers):
                return success_list, reply_msg.puid_status # also return the dict with the status of all processes
                                                           # in the list for future use outside the
                                                           # Connection.wait() context
            else: # there is at least one process exited and at least one errored/timed out
                return None, reply_msg.puid_status
        else: # 'any' option and at least one process exited
            return success_list, reply_msg.puid_status

    elif timeout_flag: # none has exited and all timed out
        return None, reply_msg.puid_status
    else:
        log.debug(f'process join {req_msg} failed: {reply_msg.puid_status.items()}')
        raise ProcessError(f'process join {req_msg} failed')

src/dragon/globalservices/process.py
DragonHPC/dragon
"""Process related objects internal to Global Services.
"""

import copy
import logging
import signal

from .. import channels as dch
from ..globalservices import channel_int as channel_int
from ..infrastructure import process_desc as process_desc
from ..infrastructure import channel_desc as channel_desc
from ..infrastructure.group_desc import GroupDescriptor
from ..localservices import options as lsopt
from ..infrastructure import facts as dfacts
from ..infrastructure import messages as dmsg
from ..infrastructure import connection as dconn
from ..infrastructure import parameters as dp
from ..infrastructure import util as dutil
from ..utils import B64

LOG = logging.getLogger('process:')


class ProcessContext:
    """Everything to do with a single process in global services

        The methods in this object are used to communicate with
        the shepherd to which this processes is assigned to
        manage its lifecycle.
    """

    def __init__(self, *, server, request, reply_channel, p_uid, node):
        self.server = server
        self.request = request
        self.reply_channel = reply_channel
        self.exit_msg = None
        self.destroy_request = None
        self.gs_ret_channel_context = None
        self._descriptor = process_desc.ProcessDescriptor(p_uid=p_uid,
                                                          name=request.user_name,
                                                          node=node,
                                                          p_p_uid=request.p_uid)

        self.process_parms = copy.copy(dp.this_process)
        self.process_parms.my_puid = p_uid
        self.process_parms.index = node
        self.stdin_context = None
        self.stdout_context = None
        self.stderr_context = None

        self.belongs_to_group = False
        self.group_addition = False

    def __str__(self):
        return f"[[{self.__class__.__name__}]] desc:{self.descriptor!r} req:{self.request!r} exit:{self.exit_msg!r}"

    @property
    def descriptor(self):
        return self._descriptor

    def _mk_sh_proc_create(self, the_tag, which_node):

        the_env = self.request.env
        for k, v in self.process_parms.env().items():
            the_env[k] = v

        # if the parent process wants output or input redirected,
        # we encapsulate necessary messages here to be sent with the
        # SHProcessCreate. All requests are then handled at once by
        # Local Services.


        fake_reply_channel = dutil.AbsorbingChannel()
        if self.request.pipesize is not None and self.request.pipesize <= 0:
            capacity = None
        else:
            capacity = self.request.pipesize

        if self.request.stdin == dmsg.PIPE:
            puid = self.process_parms.my_puid

            channel_options = lsopt.ChannelOptions(capacity=capacity)
            options = channel_desc.ChannelOptions(ref_count=True, local_opts=channel_options)
            req_msg = dmsg.GSChannelCreate(tag=self.server.tag_inc(), p_uid=puid,
                                r_c_uid=dfacts.GS_INPUT_CUID, m_uid=dfacts.infrastructure_pool_muid_from_index(which_node),
                                user_name="stdin_"+str(puid), options=options)
            _, _, self.stdin_context = channel_int.ChannelContext.construct(self.server, req_msg, fake_reply_channel,
                                node_override=self.process_parms.index, send_msg=False)
            self.stdin_context.incref()
            stdin_msg = self.stdin_context.shchannelcreate_msg
        else:
            stdin_msg = None


        if self.request.stdout == dmsg.PIPE:
            puid = self.process_parms.my_puid

            channel_options = lsopt.ChannelOptions(capacity=capacity)
            options = channel_desc.ChannelOptions(ref_count=True, local_opts=channel_options)
            req_msg = dmsg.GSChannelCreate(tag=self.server.tag_inc(), p_uid=puid,
                                r_c_uid=dfacts.GS_INPUT_CUID, m_uid=dfacts.infrastructure_pool_muid_from_index(which_node),
                                user_name="stdout_"+str(puid), options=options)
            _, _, self.stdout_context = channel_int.ChannelContext.construct(self.server, req_msg, fake_reply_channel,
                                node_override=self.process_parms.index, send_msg=False)
            self.stdout_context.incref()
            stdout_msg = self.stdout_context.shchannelcreate_msg
        else:
            stdout_msg = None

        if self.request.stderr == dmsg.PIPE:
            puid = self.process_parms.my_puid

            channel_options = lsopt.ChannelOptions(capacity=capacity)
            options = channel_desc.ChannelOptions(ref_count=True, local_opts=channel_options)
            req_msg = dmsg.GSChannelCreate(tag=self.server.tag_inc(), p_uid=puid,
                                r_c_uid=dfacts.GS_INPUT_CUID, m_uid=dfacts.infrastructure_pool_muid_from_index(which_node),
                                user_name="stderr_"+str(puid), options=options)
            _, _, self.stderr_context = channel_int.ChannelContext.construct(self.server, req_msg, fake_reply_channel,
                                node_override=self.process_parms.index, send_msg=False)
            self.stderr_context.incref()
            stderr_msg = self.stderr_context.shchannelcreate_msg
        else:
            stderr_msg = None

        return dmsg.SHProcessCreate(tag=the_tag,
                                    p_uid=dfacts.GS_PUID,
                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                    t_p_uid=self.descriptor.p_uid,
                                    exe=self.request.exe,
                                    args=self.request.args,
                                    rundir=self.request.rundir,
                                    env=the_env,
                                    stdin = self.request.stdin,
                                    stdout = self.request.stdout,
                                    stderr=self.request.stderr,
                                    stdin_msg=stdin_msg,
                                    stdout_msg=stdout_msg,
                                    stderr_msg=stderr_msg,
                                    pmi_info=self.request._pmi_info)  #pylint: disable=protected-access

    def mk_sh_proc_kill(self, the_tag, the_sig=signal.SIGKILL):
        return dmsg.SHProcessKill(tag=the_tag,
                                  p_uid=dfacts.GS_PUID,
                                  r_c_uid=dfacts.GS_INPUT_CUID,
                                  t_p_uid=self.descriptor.p_uid,
                                  sig=the_sig)

    @classmethod
    def construct(cls, server, msg, reply_channel, head, send_msg=True, belongs_to_group=False, addition=False):
        """Makes a new context, registers it with the server, and sends a request start message

        :param server: global server context object
        :type server: dragon.globalservices.server.GlobalContext
        :param msg: GSProcessCreate message
        :type msg: dragon.infrastructure.messages.GSProcessCreate
        :param reply_channel: the handle needed to reply to this message
        :type reply_channel: dragon.infrastructure.connection.Connection
        :param head: whether this is the head process or not
        :type head: bool
        :param send_msg: whether to issue and send the response message or not, defaults to True
        :type send_msg: bool, optional
        :param belongs_to_group: whether this process is a member of a group or not, defaults to False
        :type belongs_to_group: bool, optional
        :param addition: whether this was called by gs.group_int.create_add() or gs.group_int.construct(), defaults to False
        :type addition: bool, optional
        :return: a tuple:
                First element: True if there is a pending continuation, False otherwise (see the reply channel)
                Second element: The tag if message issued, or the error message if there was a failure, None otherwise
                Third element: The newly constructed context object if issued, None otherwise
        :rtype: tuple
        """
        if msg.user_name in server.process_names:
            LOG.info(f'process name {msg.user_name} in use')
            existing_ctx = server.process_table[server.process_names[msg.user_name]]
            if send_msg:
                rm = dmsg.GSProcessCreateResponse(tag=server.tag_inc(),
                                                  ref=msg.tag,
                                                  err=dmsg.GSProcessCreateResponse.Errors.ALREADY,
                                                  desc=existing_ctx.descriptor)
                reply_channel.send(rm.serialize())
            return False, 'already', existing_ctx

        this_puid, auto_name = server.new_puid_and_default_name()

        if head:
            LOG.info(f'head puid is {this_puid}')
            server.head_puid = this_puid
            server._state = server.RunState.HAS_HEAD

        if not msg.user_name:
            msg.user_name = auto_name

        which_node = server.choose_shepherd(msg)

        context = cls(server=server, request=msg, reply_channel=reply_channel,
                      p_uid=this_puid, node=which_node)
        server.process_names[msg.user_name] = this_puid
        server.process_table[this_puid] = context

        if belongs_to_group:
            context.belongs_to_group = belongs_to_group
            if addition:
                context.group_addition = addition

        outbound_tag = None

        # start making the infrastructure channels if they are needed,
        # otherwise just start the process directly.
        if msg.options.make_inf_channels:
            LOG.debug(f'making inf channels for {which_node}')
            fake_reply_channel = dutil.AbsorbingChannel()
            gsr_msg = dmsg.GSChannelCreate(tag=server.tag_inc(), p_uid=0, r_c_uid=0,
                                           m_uid=dfacts.infrastructure_pool_muid_from_index(which_node))

            issued, tag, chan_context = channel_int.ChannelContext.construct(server, gsr_msg,
                                                                             fake_reply_channel)

            context.gs_ret_channel_context = chan_context

            if not issued:
                LOG.info(f'failed creating gs ret channel for puid {this_puid}')
                context.descriptor.state = process_desc.ProcessDescriptor.State.DEAD
                err_msg = 'gs channel create fail'
                if send_msg:
                    rm = dmsg.GSProcessCreateResponse(tag=server.tag_inc(),
                                                      ref=msg.tag,
                                                      err=dmsg.GSProcessCreateResponse.Errors.FAIL,
                                                      err_info=err_msg)
                    context.reply_channel.send(rm.serialize())
                return False, err_msg, context
            else:
                server.pending[tag] = context.check_channel_const
                return False, None, context
        else:
            context.gs_ret_channel_context = None
            outbound_tag = context.send_start()

            # if it does not belong to a group issue a pending completion now
            if not belongs_to_group:
                server.pending[outbound_tag] = context.complete_construction

            return True, outbound_tag, context

    def send_start(self):
        # send request to shep, remember pending process.
        shep_hdl = self.server.shep_inputs[self.descriptor.node]

        outbound_tag = self.server.tag_inc()
        shep_req = self._mk_sh_proc_create(outbound_tag, self._descriptor.node)

        # In cases where we are sending a large amount of
        # messages, such as with the GSGroupCreate handler,
        # we can fill the GS Input Queue with responses and
        # basically cause the GS / TA / LS to be unable to
        # send/receive any messages. To prevent this, we'll
        # enqueue pending sends and interleave sending and
        # receiving messages to allow us to process responses
        # on the input queue.

        self.server.pending_sends.put((shep_hdl, shep_req.serialize()))
        LOG.debug(f'request {shep_req} to shep')
        return outbound_tag

    def check_channel_const(self, msg):
        """This is the channel construction message - see if that worked and if so
        continue on maybe to finishing up the construction
        """
        LOG.debug('checking inf channel construction')
        channel_constructed = self.gs_ret_channel_context.complete_construction(msg)

        if self.belongs_to_group:
            guid, (_, _) = self.server.resource_to_group_map[self.descriptor.p_uid]
            group_context = self.server.group_table[guid]

        if not channel_constructed:
            LOG.info(f'failed post issue creating gs ret channel for puid {self.descriptor.p_uid}')
            self.descriptor.state = process_desc.ProcessDescriptor.State.DEAD
            rm = dmsg.GSProcessCreateResponse(tag=self.server.tag_inc(),
                                                ref=self.request.tag,
                                                err=dmsg.GSProcessCreateResponse.Errors.FAIL,
                                                err_info='gs channel create fail, post issue')
            if self.belongs_to_group:
                self.server.group_to_pending_resource_map[(self.request.tag, guid)] = self
                group_context._construction_helper(rm)
            else:
                self.reply_channel.send(rm.serialize())
            return

        sdesc = self.gs_ret_channel_context.descriptor.sdesc

        self.process_parms.gs_ret_cd = B64.bytes_to_str(sdesc)

        outbound_tag = self.send_start()
        if self.belongs_to_group:
            self.server.group_to_pending_resource_map[(outbound_tag, guid)] = self
            if self.group_addition:
                self.server.pending[outbound_tag] = group_context.complete_addition
            else:
                self.server.pending[outbound_tag] = group_context.complete_construction
        else:
            self.server.pending[outbound_tag] = self.complete_construction

    def _send_gspingproc(self):
        LOG.debug(f'sending ping to new process {self.descriptor.p_uid}')

        self.descriptor.gs_ret_cuid = self.gs_ret_channel_context.descriptor.c_uid

        pingmsg = dmsg.GSPingProc(tag=self.server.tag_inc(),
                                  mode=self.request.options.mode,
                                  argdata=self.request.options.argdata)

        # Perf issue: how long it takes to
        # attach to a remote channel and send this stuff.
        # The head of queue blocking could get troublesome if
        # there is a long wait here.  We could
        # attach to the channel and load the ping in it
        # ahead of actually starting the process if we wanted to.

        dsd = self.gs_ret_channel_context.descriptor.sdesc
        ping_chan = dch.Channel.attach(dsd)
        ping_con = dconn.Connection(outbound_initializer=ping_chan, policy=dp.POLICY_INFRASTRUCTURE)
        self.gs_ret_channel_context.attached_connection = ping_con
        ping_con.send(pingmsg.serialize())
        LOG.debug('ping sent, channel and connection kept for later')
        # note that if the argdata mode is ArgMode.PYTHON_CHANNEL the
        # starting process will be using this channel to deliver the argdata
        # directly.  There is no race because this must be complete before
        # any GS transactions get done from the new process

    def complete_construction(self, msg, send_msg=True):
        """Completes construction of a new managed process

        :param msg: SHProcessCreateResponse message
        :type msg: SHProcessCreateResponse
        :param send_msg: whether to send the response message or not, defaults to True
        :type send_msg: bool, optional
        :raises RuntimeError: when the message error is of unknown type
        :return: True or False according to success
        :rtype: bool
        """
        if dmsg.SHProcessCreateResponse.Errors.SUCCESS == msg.err:
            self.descriptor.state = process_desc.ProcessDescriptor.State.ACTIVE

            if self.stdin_context is not None:
                channel_constructed = self.stdin_context.complete_construction(msg.stdin_resp)

                if not channel_constructed:
                    err_msg = f'Failed to create stdin channel for {self.descriptor.p_uid}.'
                    LOG.info(err_msg)
                    self.descriptor.state = process_desc.ProcessDescriptor.State.DEAD
                    if send_msg:
                        rm = dmsg.GSProcessCreateResponse(tag=self.server.tag_inc(),
                                                          ref=self.request.tag,
                                                          err=dmsg.GSProcessCreateResponse.Errors.FAIL,
                                                          err_info=err_msg)
                        self.reply_channel.send(rm.serialize())
                    return False

                # After line below, refcnt is 2. This is necessary to keep the channel from
                # being deleted should the process exit before the parent has had a chance to
                # read its output. The refcnt is decremented again once the connection is handed
                # off to the parent process.
                self.stdin_context.refcnt += 1
                stdin_sdesc = B64.bytes_to_str(self.stdin_context.descriptor.sdesc)
            else:
                stdin_sdesc = None

            if self.stdout_context is not None:
                channel_constructed = self.stdout_context.complete_construction(msg.stdout_resp)

                if not channel_constructed:
                    err_msg = f'Failed to create stdout channel for {self.descriptor.p_uid}.'
                    LOG.info(err_msg)
                    self.descriptor.state = process_desc.ProcessDescriptor.State.DEAD
                    if send_msg:
                        rm = dmsg.GSProcessCreateResponse(tag=self.server.tag_inc(),
                                                          ref=self.request.tag,
                                                          err=dmsg.GSProcessCreateResponse.Errors.FAIL,
                                                          err_info=err_msg)
                        self.reply_channel.send(rm.serialize())
                    return False

                # After line below, refcnt is 2. This is necessary to keep the channel from
                # being deleted should the process exit before the parent has had a chance to
                # read its output. The refcnt is decremented again once the connection is handed
                # off to the parent process.
                self.stdout_context.refcnt += 1
                stdout_sdesc = B64.bytes_to_str(self.stdout_context.descriptor.sdesc)
            else:
                stdout_sdesc = None

            if self.stderr_context is not None:
                channel_constructed = self.stderr_context.complete_construction(msg.stderr_resp)

                if not channel_constructed:
                    err_msg = f'Failed to create stderr channel for {self.descriptor.p_uid}.'
                    LOG.info(err_msg)
                    self.descriptor.state = process_desc.ProcessDescriptor.State.DEAD
                    if send_msg:
                        rm = dmsg.GSProcessCreateResponse(tag=self.server.tag_inc(),
                                                          ref=self.request.tag,
                                                          err=dmsg.GSProcessCreateResponse.Errors.FAIL,
                                                          err_info=err_msg)
                        self.reply_channel.send(rm.serialize())
                    return False

                # After line below, refcnt is 2. This is necessary to keep the channel from
                # being deleted should the process exit before the parent has had a chance to
                # read its output. The refcnt is decremented again once the connection is handed
                # off to the parent process.
                self.stderr_context.refcnt += 1
                stderr_sdesc = B64.bytes_to_str(self.stderr_context.descriptor.sdesc)
            else:
                stderr_sdesc = None

            self.descriptor.stdin_sdesc = stdin_sdesc
            self.descriptor.stdout_sdesc = stdout_sdesc
            self.descriptor.stderr_sdesc = stderr_sdesc

            if self.descriptor.p_p_uid in self.server.process_table:
                parent_ctx = self.server.process_table[self.descriptor.p_p_uid]
                parent_ctx.descriptor.live_children.add(self.descriptor.p_uid)

            # if we made a gs return channel then we must send a GSPingProc to it.
            if self.request.options.make_inf_channels:
                self._send_gspingproc()

            if send_msg:
                response = dmsg.GSProcessCreateResponse(tag=self.server.tag_inc(),
                                                        ref=self.request.tag,
                                                        err=dmsg.GSProcessCreateResponse.Errors.SUCCESS,
                                                        desc=self.descriptor)
                self.reply_channel.send(response.serialize())
            succeeded = True
        elif dmsg.SHProcessCreateResponse.Errors.FAIL == msg.err:
            self.descriptor.state = process_desc.ProcessDescriptor.State.DEAD

            # clean up tables - don't keep ProcessDescriptor stuff around
            # for things that never were alive.
            if self.descriptor.p_uid == self.server.head_puid:
                LOG.info('head process creation failed')
                # TODO: whatever cleanup actions desired when the head process leaves.
                self.server._state = self.server.RunState.WAITING_FOR_HEAD
                self.server.head_puid = None

            del self.server.process_table[self.descriptor.p_uid]
            del self.server.process_names[self.request.user_name]

            if self.request.options.make_inf_channels:
                # TODO: can we bind this in a cleaner way?
                clean_msg = dmsg.GSChannelDestroy(tag=self.server.tag_inc(), r_c_uid=0,
                                                  p_uid=dfacts.GS_PUID,
                                                  c_uid=self.gs_ret_channel_context.descriptor.c_uid)
                channel_int.ChannelContext.destroy(self.server, clean_msg, dutil.AbsorbingChannel())

            if send_msg:
                response = dmsg.GSProcessCreateResponse(tag=self.server.tag_inc(),
                                                        ref=self.request.tag,
                                                        err=dmsg.GSProcessCreateResponse.Errors.FAIL,
                                                        err_info=msg.err_info)
                self.reply_channel.send(response.serialize())
            succeeded = False
        else:
            raise RuntimeError(f'got {msg!s} err {msg.err} unknown')

        if send_msg:
            LOG.debug(f'create response sent, tag {response.tag} ref {response.ref} pending cleared')

        return succeeded

    @staticmethod
    def kill(server, msg, reply_channel, belongs_to_group=False):
        target_uid, found, errmsg = server.resolve_puid(msg.user_name, msg.t_p_uid)
        gspkr = dmsg.GSProcessKillResponse

        if not found:
            rm = gspkr(tag=server.tag_inc(),
                       ref=msg.tag,
                       err=gspkr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            LOG.debug(f'process absent; response to {msg}: {rm}')
            return False, None
        else:
            pctx = server.process_table[target_uid]
            pdesc = pctx.descriptor
            pds = process_desc.ProcessDescriptor.State
            if pds.DEAD == pdesc.state:
                rm = gspkr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gspkr.Errors.DEAD,
                           exit_code=pdesc.ecode)
                reply_channel.send(rm.serialize())
                LOG.debug(f'process dead, response to {msg}: {rm}')
                return False, None

            # this is a highly unlikely case to happen
            # because the user must have received a create response back
            # before requesting killing this process, which means that the
            # state of the process is not pending anymore
            # so, what we really want here is to log the case and move on
            elif pds.PENDING == pdesc.state:
                rm = gspkr(tag=server.tag_inc(),
                           ref=msg.tag,
                           err=gspkr.Errors.PENDING,
                           err_info=f'process {target_uid} is pending')
                reply_channel.send(rm.serialize())
                LOG.debug(f'process pending while kill request -- this should not be happening, response to {msg}: {rm}')
                return False, None
            elif pds.ACTIVE == pdesc.state:
                pctx.descriptor.state = pds.PENDING
                pctx.destroy_request = msg
                pctx.reply_channel = reply_channel
                the_tag = server.tag_inc()
                shep_kill_msg = pctx.mk_sh_proc_kill(the_tag, msg.sig)
                target_node = pdesc.node
                if not belongs_to_group:
                    server.pending[the_tag] = pctx.complete_kill

                server.pending_sends.put((server.shep_inputs[target_node], shep_kill_msg.serialize()))
                LOG.debug(f'kill {msg} sent to shep as {shep_kill_msg} on node {target_node}')
                return True, the_tag
            else:
                raise NotImplementedError('close case')

    def complete_kill(self, msg):
        gspkr = dmsg.GSProcessKillResponse
        shpkr = dmsg.SHProcessKillResponse

        if shpkr.Errors.FAIL == msg.err:
            rm = gspkr(tag=self.server.tag_inc(), ref=self.destroy_request.tag,
                       err=gspkr.Errors.FAIL_KILL, err_info=msg.err_info)
            kill_succeeded = False
        elif shpkr.Errors.SUCCESS == msg.err:

            rm = gspkr(tag=self.server.tag_inc(),
                       ref=self.destroy_request.tag,
                       err=gspkr.Errors.SUCCESS)

            kill_succeeded = True
        else:
            raise NotImplementedError('close case')

        if self.descriptor.state == self.descriptor.State.PENDING:
            self.descriptor.state = self.descriptor.State.ACTIVE

        LOG.debug(f'sending kill response to request {self.destroy_request}: {rm}')
        self.reply_channel.send(rm.serialize())
        return kill_succeeded

src/dragon/globalservices/process_int.py
DragonHPC/dragon
"""Global server context object and entry point
"""
import enum
import logging
import os
import pickle
import sys
import time
import threading

from queue import Queue
from .. import channels as dch

from ..infrastructure import channel_desc
from ..infrastructure import pool_desc
from ..infrastructure import process_desc
from ..infrastructure import group_desc

from .channel_int import ChannelContext
from .pool_int import PoolContext
from .process_int import ProcessContext
from .node_int import NodeContext
from .group_int import GroupContext

from ..dtypes import DragonError

from . import startup
from .policy_eval import ResourceLayout, PolicyEvaluator
from ..infrastructure.policy import GS_DEFAULT_POLICY

from ..infrastructure import facts as dfacts
from ..infrastructure import messages as dmsg
from ..infrastructure import util as dutil
from ..infrastructure import connection as dconn
from ..infrastructure import parameters as dparm
from ..dlogging import util as dlog
from ..dlogging.util import DragonLoggingServices as dls
from ..utils import B64, set_procname


class ServerError(Exception):
    pass


class GlobalContext(object):
    """Global Context for the dragon runtime

        This object is meant to hold whatever info is maintained
        for all channels, pools, and processes and contain the main
        message processing and dispatch code.

        Ultimately parts of this will need to be distributed.

        Important member variables:

        process_table - key is p_uid, values are ProcessContext
        channel_table - key is c_uid, values are ChannelContext
        pool_table - key is m_uid, values are PoolContext
        node_table - key is h_uid, values are NodeContext
        group_table - key is g_uid, values are GroupContext

        process_names - key is string name of process, value is p_uid
        channel_names - key is string name of channel, value is c_uid
        pool_names - key is string name of pool, value is m_uid
        node_names - key is string name of node, value is h_uid
        group_names - key is string name of group, value is g_uid
    """

    class LaunchModes(enum.Enum):
        SINGLE = enum.auto()
        MULTI = enum.auto()
        TEST_STANDALONE_SINGLE = enum.auto()
        TEST_STANDALONE_MULTI = enum.auto()

    class RunState(enum.Enum):
        STARTING = enum.auto()
        WAITING_FOR_HEAD = enum.auto()
        HAS_HEAD = enum.auto()
        SHUTTING_DOWN = enum.auto()

    DTBL = {}  # dispatch table for handlers, no metadata

    def __init__(self, test_gs_stdout=None):
        self.process_table = dict()  # key = p_uid, value = ProcessContext
        self.channel_table = dict()  # key = c_uid, value = ChannelContext
        self.pool_table = dict()  # key = m_uid, value = PoolContext
        self.node_table = dict()  # key = h_uid, value = NodeContext
        self.group_table = dict() # key = g_uid, value = GroupContext

        self.process_names = dict()  # key = name, value = p_uid
        self.channel_names = dict()  # key = name, value = c_uid
        self.pool_names = dict()  # key = name, value = m_uid
        self.node_names = dict()  # key = name, value = h_uid
        self.group_names = dict() # key = name, value = g_uid

        # used by the group of resources objects, in group_int.py
        self.resource_to_group_map = dict()  # key = resource uid e.g. p_uid, value = tuple including g_uid, and position of the resource in the GSGroupCreate.items list
        self.group_to_pending_resource_map = dict() # used for pending resources of a group, key = (msg tag, g_uid), value = resource_context
        self.group_resource_count = dict()  # key = g_uid, value = list of items corresponding to the multiplicity of each list item in GroupDescriptor.sets
        self.group_destroy_resource_count = dict()  # is applied only when destroy has been called, key = g_uid, value = list of items corresponding to the multiplicity of each list item in GroupDescriptor.sets

        self.pending = dict()  # key = tag, value = continuation on response
        self.pending_group_destroy = dict()  # is applied only to groups of resources when destroy is called, key = uid, value = continuation on response

        self.pending_sends = Queue()

        self.pending_join = dutil.PriorityMultiMap()  # pending joins, key = p_uid val = set of join req for this p_uid
        self.pending_join_list = dutil.PriorityMultiMap()  # pending joins, key = frozenset of p_uid, val = set of join req for this set of p_uid
        self.pending_channel_joins = dutil.PriorityMultiMap()  # pending channel joins  key = channel names

        self.puid_join_list_status = dict()  # this is a dict of dicts that holds the p_uid status for every muti-join wait request
                                             # key is a tuple with of (msg.tag, msg.p_uid) and value is a dict with: key = p_uid and val = (status, info)

        self.sh_pings_recvd = 0 # added to as we receive the sh pings
        self.shep_inputs = [None]
        self.gs_input = None
        self.bela_input = None
        self.test_connections = None
        self.head_puid = None
        self._tag = 0
        self._next_puid = dfacts.FIRST_PUID
        self._next_cuid = dfacts.FIRST_CUID
        self._next_muid = dfacts.FIRST_MUID
        self._next_guid = dfacts.FIRST_GUID

        if (test_gs_stdout is None):
            self.msg_stdout = dutil.NewlineStreamWrapper(sys.stdout)
        else:
            self.msg_stdout = test_gs_stdout

        self._launch_mode = None
        self._state = GlobalContext.RunState.STARTING
        self._startup_logger = logging.getLogger(dls.GS).getChild('startup:')
        self._shutdown_logger = logging.getLogger(dls.GS).getChild('shutdown:')
        self._dispatch_logger = logging.getLogger(dls.GS).getChild('dispatch:')
        self._process_logger = logging.getLogger(dls.GS).getChild('proc:')
        self._pool_logger = logging.getLogger(dls.GS).getChild('pool:')
        self._channel_logger = logging.getLogger(dls.GS).getChild('channel:')
        self._node_logger = logging.getLogger(dls.GS).getChild('node:')
        self._group_logger = logging.getLogger(dls.GS).getChild('group:')

        self.circ_index = 0 # helper index for implementing the circular order of choosing shepherds

        self.policy_eval = None # Policy Evaluator, init'd in `handle_sh_ping`

    ################ Support Functions for GS #######################################


    def dump_to(self, fh):
        """Dumps a human readable representation of current state to a file-like object.

            May want eventually to have a version that essentially pickles
            the server state apart from handles, but that isn't this.
        """

        fh.write(f'state:\t\t{self._state.name}\n')
        fh.write(f'mode:\t\t{self._launch_mode.name}\n')
        fh.write(f'current tag:\t{self.tag}\n')
        fh.write(f'head_puid:\t{self.head_puid}\n')

        fh.write('\nProcesses:\n')
        for ctx in self.process_table.values():
            fh.write(f'{ctx!s}\n')

        # TODO: PE-35937, dump channels and pools
        fh.write('\nPending join:\n')

        for p_uid, waiters in self.pending_join.items():
            fh.write(f'\t{p_uid}:\t')
            fh.write(', '.join([f'{wr.p_uid}:{wr.r_c_uid}' for wr in waiters]))
            fh.write('\n')

        for p_uid_list, waiters in self.pending_join_list.items():
            fh.write(f'\t{p_uid_list}:\t')
            fh.write(', '.join([f'{wr.p_uid}:{wr.r_c_uid}' for wr in waiters]))
            fh.write('\n')

        fh.write('\nChannels:\n')
        fh.write('--none for now---\n')


    @property
    def tag(self):
        return self._tag


    def tag_inc(self):
        """Returns current tag counter and increments it"""
        tmp = self._tag
        self._tag += 1
        return tmp


    def new_puid_and_default_name(self):
        tmp = self._next_puid
        self._next_puid += 1
        return tmp, dfacts.DEFAULT_PROCESS_NAME_BASE + str(tmp)


    def new_cuid_and_default_name(self):
        tmp = self._next_cuid
        self._next_cuid += 1
        return tmp, dfacts.DEFAULT_CHANNEL_NAME_BASE + str(tmp)


    def new_muid_and_default_name(self):
        tmp = self._next_muid
        self._next_muid += 1
        return tmp, dfacts.DEFAULT_POOL_NAME_BASE + str(tmp)


    def new_guid_and_default_name(self):
        tmp = self._next_guid
        self._next_guid += 1
        return tmp, dfacts.DEFAULT_GROUP_NAME_BASE + str(tmp)


    # TODO: check that the message came from somewhere that makes sense.
    # for instance clients shouldn't get to send teardown messages
    def next_msg(self, timeout=0):
        # for basic unit testing
        if self.gs_input is None:
            return None

        # TODO: consider reworking call interface considering
        # single message assumption
        if self.gs_input.poll(timeout):
            rv = dmsg.parse(self.gs_input.recv())  # blocking recv
        else:
            rv = None

        return rv


    # comes up with the handle needed to reply to this message.
    def get_reply_handle(self, msg):
        log = self._dispatch_logger

        if self._launch_mode in {self.LaunchModes.MULTI,
                                 self.LaunchModes.SINGLE,
                                 self.LaunchModes.TEST_STANDALONE_SINGLE}:
            if 1 == msg.r_c_uid:
                log.debug("this is stdout")
                return self.msg_stdout
            elif 2 == msg.r_c_uid:
                raise ServerError('channel 2 requested as r_c_uid but this is gs input')
            elif dfacts.BASE_BE_CUID <= msg.r_c_uid < (
                    dfacts.BASE_BE_CUID + dfacts.RANGE_BE_CUID):
                log.debug("this is bela")
                return self.bela_input
            elif dfacts.BASE_SHEP_CUID <= msg.r_c_uid < (
                    dfacts.BASE_SHEP_CUID + dfacts.RANGE_SHEP_CUID):
                log.debug("this is shep")
                return self.shep_inputs[dfacts.index_from_shepherd_cuid(msg.r_c_uid)]
            elif self.test_connections and msg.r_c_uid in self.test_connections:
                log.debug("this is test connections")
                return self.test_connections[msg.r_c_uid]
            elif msg.r_c_uid in self.channel_table:
                # TODO PE-37246 GS Channel attachment ageing policy
                context = self.channel_table[msg.r_c_uid]
                if context.attached_connection is None:
                    underlying_channel = dch.Channel.attach(context.descriptor.sdesc)

                    if context.is_reading:
                        context.attached_connection = dconn.Connection(inbound_initializer=underlying_channel, policy=dparm.POLICY_INFRASTRUCTURE)
                    else:
                        context.attached_connection = dconn.Connection(outbound_initializer=underlying_channel, policy=dparm.POLICY_INFRASTRUCTURE)

                log.debug("this is context and channel_table")
                return context.attached_connection
            else:
                raise ServerError(f'cuid {msg.r_c_uid} requested but unknown')
        else:
            raise NotImplementedError('unknown launch mode')


    # TODO: design/collaboration approach for placement
    # interfaces.  tickets to follow...

    def choose_shepherd(self, msg:dmsg.GSProcessCreate):
        """
            Pick the shepherd to launch a process on. If the GSProcessCreate
            request includes a ResourceLayout request, use that. Otherwise,
            roundrobin the placement across the available nodes.
        """
        log = self._dispatch_logger
        assert isinstance(msg, dmsg.GSProcessCreate)

        if not msg.layout:
            index = self.circ_index
            self.circ_index = (self.circ_index + 1) % len(self.shep_inputs)
        else:
            assert isinstance(msg.layout, ResourceLayout)
            assert msg.layout.h_uid in self.node_table.keys()
            index = self.node_table[msg.layout.h_uid].ls_index

        log.info(f"choose shepherd {index} for {msg}")
        return index


    # picks the node to put a pool on
    # should raise from here if the options or state lead to nonsense
    # note this can't be a pure method on anything
    # inside the message because sometimes the choice
    # will depend on system state.
    def choose_pool_node(self, msg):
        assert isinstance(msg, dmsg.GSPoolCreate)
        if msg.p_uid in self.process_table:
            return self.process_table[msg.p_uid].descriptor.node # the node of the requesting process
        else:
            # we have to retain this way for cases where the launcher (p_uid is 0) is the requesting process
            return 0


    # this ought to respect whether the m_uid asked for
    # is on the node that is asked for.   Might need a convention
    # for this, like 0 == u pick.
    # Might also want to take into account a logging pool m_uid. Not
    # sure since GS may never be involved, but worth looking at.
    def choose_channel_node(self, msg):
        log = self._dispatch_logger

        assert isinstance(msg, dmsg.GSChannelCreate)
        if msg.m_uid < dfacts.FIRST_LOGGING_MUID:
            # Then this is an infrastructure Pool m_uid.
            if msg.p_uid in self.process_table:
                log.debug('Choosing channel node on branch 1')
                return self.process_table[msg.p_uid].descriptor.node  # the node of the requesting process
            else:
                log.debug('Choosing channel node on branch 2')
                # we have to retain this way for cases where the launcher (p_uid is 0) is the requesting process
                return dfacts.index_from_infrastructure_pool_muid(msg.m_uid)
        else:
            if msg.m_uid in self.pool_table:
                log.debug('Choosing channel node on branch 3')
                return self.pool_table[msg.m_uid].descriptor.node
            else:
                # this is a predefined memory pool
                if dfacts.is_pre_defined_pool(msg.m_uid):
                    log.debug('Choosing channel node on branch 4')
                    return dfacts.index_from_default_pool_muid(msg.m_uid)
                else:
                    log.debug('Choosing channel node on branch 5')
                    # the requested pool to create the channel on does not exist
                    log.debug(f"The requested pool with muid {msg.m_uid} to create the channel on, does not exist.")
                    raise ValueError(f'The requested m_uid is not valid {msg.m_uid}')

    # This function uses the routing table to call the correct method for an
    # infrastructure message. See the @route decorated methods.
    def dispatch(self, msg):
        log = self._dispatch_logger
        if type(msg) in GlobalContext.DTBL:
                GlobalContext.DTBL[type(msg)][0](self, msg)
        else:
            log.warning(f"{msg} isn't implemented")
            raise NotImplementedError(f"Message type {type(msg)} is not handled, yet")


    @staticmethod
    def _resolve_uid(name, uid, names, table, what, idname):
        errmsg = ''

        if name:
            if name in names:
                uid = names[name]
                assert uid in table
                found = True
            else:
                errmsg = f'no {what} named {name}'
                found = False
        else:
            found = uid in table
            if not found:
                errmsg = f'no {idname} {uid}'

        return uid, found, errmsg


    def resolve_puid(self, name=None, p_uid=None):
        """Checks if a process exists given name and/or p_uid.

                Given name or p_uid specifying a process, resolves
                the p_uid and tells whether it is present.

                If the name is specified the p_uid parm is ignored.

                Returns: (p_uid, found, errmsg) where
                    p_uid is the p_uid
                    found whether that p_uid exists
                    errmsg what is wrong if it does not exist.
                """
        return self._resolve_uid(name, p_uid,
                                 self.process_names, self.process_table,
                                 'process', 'p_uid')


    def resolve_muid(self, name, m_uid=-1):
        """Checks if a memory pool exists given name and/or m_uid.

            see docstring on _resolve_puid for more details.
        """
        return self._resolve_uid(name, m_uid,
                                 self.pool_names, self.pool_table,
                                 'pool', 'm_uid')


    def resolve_cuid(self, name, c_uid=-1):
        """Checks if a channel exists given name and/or c_uid.

            see docstring on _resolve_puid for more details.
        """
        return self._resolve_uid(name, c_uid,
                                 self.channel_names, self.channel_table,
                                 'channel', 'c_uid')


    def resolve_huid(self, name, h_uid=-1):
        """Checks if a node exists given name and/or h_uid.

            see docstring on _resolve_puid for more details.
        """
        return self._resolve_uid(name, h_uid,
                                 self.node_names, self.node_table,
                                 'node', 'h_uid')


    def resolve_guid(self, name, g_uid=-1):
        """Checks if a group exists given name and/or g_uid.

            see docstring on _resolve_puid for more details.
        """
        return self._resolve_uid(name, g_uid,
                                 self.group_names, self.group_table,
                                 'group', 'g_uid')

    def do_timeouts(self):
        log = self._process_logger

        gspjr = dmsg.GSProcessJoinResponse
        timed_out = self.pending_join.get_timed_out()
        for join_request in timed_out:
            log.debug(f'join timed out: {join_request}')
            p_uid, req_msg = join_request
            reply_channel = self.get_reply_handle(req_msg)
            rm = gspjr(tag=self.tag_inc(), ref=req_msg.tag, err=gspjr.Errors.TIMEOUT)
            reply_channel.send(rm.serialize())
            self.pending_join.remove_one(p_uid, req_msg)
            log.debug(f'timed out join response to {req_msg!s}: {rm!s}')

        gspjlr = dmsg.GSProcessJoinListResponse
        timed_out = self.pending_join_list.get_timed_out()
        for join_request in timed_out:
            log.debug(f'multi-join timed out: {join_request}')
            p_uid_list, req_msg = join_request # p_uid_list is a frozenset
            reply_channel = self.get_reply_handle(req_msg)
            for p_uid in p_uid_list:
                self.puid_join_list_status[(req_msg.tag, req_msg.p_uid)][p_uid] = (gspjlr.Errors.TIMEOUT.value, None)
            rm = gspjlr(tag=self.tag_inc(), ref=req_msg.tag, puid_status=self.puid_join_list_status[(req_msg.tag, req_msg.p_uid)])
            reply_channel.send(rm.serialize())
            self.pending_join_list.remove_one(p_uid_list, req_msg)
            log.debug(f'timed out multi-join response to {req_msg!s}: {rm!s}')
            del self.puid_join_list_status[(req_msg.tag, req_msg.p_uid)] # we are done with this msg req

        gscjr = dmsg.GSChannelJoinResponse
        timed_out = self.pending_channel_joins.get_timed_out()
        for channel_join_request in timed_out:
            log.debug(f'channel join timed out: {channel_join_request}')
            name, req_msg = channel_join_request
            reply_channel = self.get_reply_handle(req_msg)
            rm = gscjr(tag=self.tag_inc(), ref=req_msg.tag, err=gscjr.Errors.TIMEOUT)
            reply_channel.send(rm.serialize())
            self.pending_channel_joins.remove_one(name, req_msg)
            log.debug(f'timed out join response to {req_msg!s}: {rm!s}')

        join_deadline = self.pending_join.next_deadline()
        join_list_deadline = self.pending_join_list.next_deadline()
        channel_join_deadline = self.pending_channel_joins.next_deadline()

        if join_deadline is None and join_list_deadline is None and channel_join_deadline is None:
            return None
        elif join_deadline is None and join_list_deadline is None:
            return max(0, channel_join_deadline - time.time())
        elif join_deadline is None and channel_join_deadline is None:
            return max(0, join_list_deadline - time.time())
        elif join_list_deadline is None and channel_join_deadline is None:
            return max(0, join_deadline - time.time())
        elif join_deadline is None:
            nearest_deadline = min(join_list_deadline, channel_join_deadline)
            return max(0, nearest_deadline - time.time())
        elif join_list_deadline is None:
            nearest_deadline = min(join_deadline, channel_join_deadline)
            return max(0, nearest_deadline - time.time())
        elif channel_join_deadline is None:
            nearest_deadline = min(join_deadline, join_list_deadline)
            return max(0, nearest_deadline - time.time())
        else:
            nearest_deadline = min(join_deadline, join_list_deadline, channel_join_deadline)
            return max(0, nearest_deadline - time.time())


    def run_global_server(self, mode=LaunchModes.SINGLE, test_gs_input=None,
                          test_shep_inputs=None,
                          test_bela_input=None, test_gs_stdout=None, test_conns=None):
        """Main global server event loop.

            The global server loop

            :param test_conns: test connection objects - key is c_uid, value is
                               handle that you can send to or recv from.

            :param mode: How are we running this
            :param test_gs_input: Global Services input handle
            :param test_shep_inputs: List of shepherd handles for the whole allocation.
            :param test_gs_stdout: Global Services stdout message handle
            :param test_bela_input: Backend/Launcher input handle.
            :param test_conns: a c_uid keyed dictionary of test channel handles
        """
        assert mode == self._launch_mode

        self.test_connections = test_conns

        log = logging.getLogger(dls.GS).getChild('main_loop')
        log.info('entered')
        # for setting inputs in test
        if test_gs_input is not None:
            self.gs_input = test_gs_input

        if test_shep_inputs is not None:
            self.shep_inputs = test_shep_inputs

        if test_bela_input is not None:
            self.bela_input = test_bela_input

        if test_gs_stdout is not None:
            self.msg_stdout = test_gs_stdout

        while self._state != self.RunState.SHUTTING_DOWN:
            if self._state == self.RunState.HAS_HEAD:

                # In cases where we are sending a large amount of
                # messages, such as with the GSGroupCreate handler,
                # we can fill the GS Input Queue with responses and
                # basically cause the GS / TA / LS to be unable to
                # send/receive any messages. To prevent this, we'll
                # enqueue pending sends and interleave sending and
                # receiving messages to allow us to process responses
                # on the input queue.

                while self.pending_sends.qsize() > 0 and not self.gs_input.poll(timeout=0):
                    chan, msg = self.pending_sends.get()
                    chan.send(msg)

                next_timeout = self.do_timeouts()
            else:
                next_timeout = None

            the_msg = self.next_msg(next_timeout)
            if the_msg is not None:
                self.dispatch(the_msg)

        log = self._shutdown_logger
        log.info('main loop exit')


    def run_startup(self, mode=LaunchModes.SINGLE, test_gs_input=None, test_shep_inputs=None,
                    test_bela_input=None, test_gs_stdout=None):
        """Complete startup activities with other actors.

            Normal exit indicates bringup went fine.
            Exception indicates bringup failed for some reason.

            The idea is to vary this for different scenarios,
            for instance tests.

            :param mode: How are we bringing this up
            :param test_shep_inputs: List of shepherd handles for the whole allocation.
            :param test_gs_input: Global Services input handle
            :param test_bela_input: Launcher/back end input handle
            :param test_gs_stdout: Global Services stdout message handle

        """

        self._launch_mode = mode

        if self.LaunchModes.SINGLE == self._launch_mode:
            self.shep_inputs, self.gs_input, self.bela_input, self.num_nodes = startup.startup_single(self)
        elif self.LaunchModes.MULTI == self._launch_mode:
            self.shep_inputs, self.gs_input, self.bela_input, self.num_nodes = startup.startup_multi(self)
        elif self.LaunchModes.TEST_STANDALONE_SINGLE == self._launch_mode:
            if test_shep_inputs is None:
                test_shep_inputs = [None]

            my_inputs = startup.startup_single(self, test_gs_input,
                                               test_shep_inputs[0],
                                               test_bela_input)

            self.shep_inputs, self.gs_input, self.bela_input, self.num_nodes = my_inputs

        elif self.LaunchModes.TEST_STANDALONE_MULTI == self._launch_mode:
            my_inputs = startup.startup_multi(self,
                                              test_gs_input,
                                              test_shep_inputs,
                                              test_bela_input)
            self.shep_inputs, self.gs_input, self.bela_input, self.num_nodes = my_inputs
        else:
            raise NotImplementedError('unknown mode')


    def add_node_from_SHPingGS(self, msg) -> None:
        """ Create node contexts and h_uids from the node descriptors we have received
            from LS during startup via SHPingGS.

            :param messages: SHPingGS messages received during startup sequence
        """
        ctx = NodeContext.construct(msg)

        self.node_table[ctx.h_uid] = ctx
        self.node_names[ctx.name] = ctx.h_uid

        self._node_logger.info(f"Added node {ctx.name}, h_uid={ctx.h_uid}")


    def detach_from_channels(self):
        # this does the infrastructure channels.  Detaching from
        # the return channels or anything else that is open should happen in the
        # cleaning up user state part, which should have happened before this
        # gets called
        log = logging.getLogger('teardown_detach')

        real_channels = []
        if isinstance(self.gs_input, dconn.Connection):
            real_channels.append(self.gs_input)

        for shep in self.shep_inputs:
            if isinstance(shep, dconn.Connection):
                real_channels.append(shep)

        if isinstance(self.bela_input, dconn.Connection):
            real_channels.append(self.bela_input)

        if not real_channels:
            return

        log.debug('detaching from channels')

        # no need to detach from pool(s) due to refcounting
        # inside the channels library
        # also, GS does not explicitly attach to any pool
        for ch in real_channels:
            if ch.outbound_chan is not None:
                ch.outbound_chan.detach()

            if ch.inbound_chan is not None:
                ch.inbound_chan.detach()

        log.debug('channels detached')


    def run_teardown(self, test_gs_input=None, test_shep_inputs=None,
                     test_bela_input=None, test_gs_stdout=None):
        """Complete teardown activities with other actors.

            Normal exit = happy path
            Exception indicates a failure

            :param test_gs_input: Global Services input handle
            :param test_shep_inputs: list of shepherd handles for the whole allocation.
            :param test_gs_stdout: Global Services stdout message handle
            :param test_bela_input: Back end/launcher message input handle
        """

        assert self.RunState.SHUTTING_DOWN == self._state

        if test_gs_stdout is not None:
            self.msg_stdout = test_gs_stdout

        if test_bela_input is not None:
            self.bela_input = test_bela_input

        log = self._shutdown_logger

        gait = self.tag_inc  # 'get and inc tag'

        log.info('teardown procedure entered')

        self.detach_from_channels()
        log.debug('detached from channels')

        self.msg_stdout.send(dmsg.GSHalted(tag=gait()).serialize())
        log.info(f"gs sent GSHalted to primary shep via stdout")

        log.info('teardown complete')


    ################ Message Handlers for GS #######################################
    @dutil.route(dmsg.SHPingGS, DTBL)
    def handle_sh_ping(self, msg):
        log = self._startup_logger
        gait = self.tag_inc  # 'get and inc tag'

        log.info(f'Received SHPingGS from Local Services on node {msg.idx}')

        # Add the node into Global Services
        self.add_node_from_SHPingGS(msg)

        self.sh_pings_recvd+=1
        log.info(f'GS has received {self.sh_pings_recvd} of an expected {self.num_nodes} SHPingGS messages.')

        if self.sh_pings_recvd == self.num_nodes:
            # And done receiving SHPingGS's
            log.info('received SHPingGS from all ls - m11')
            self.bela_input.send(dmsg.GSIsUp(tag=gait()).serialize())
            log.info('gs signaled to la_be it is up (GSIsUP) - m12.1')
            self._state = self.RunState.WAITING_FOR_HEAD
            # Set up global PolicyEvaluator here
            # When we get to elasticity we can do more by creating the evaluator here
            # Get a list of NodeDescriptors from NodeContext
            node_descrs = [x.descriptor for x in self.node_table.values()]
            self.policy_eval = PolicyEvaluator(node_descrs, GS_DEFAULT_POLICY)



    @dutil.route(dmsg.GSTeardown, DTBL)
    def handle_gs_teardown(self, msg):
        log = self._shutdown_logger
        log.info('Received the GSTeardown message. Now shutting down Global Services.')
        self._state = self.RunState.SHUTTING_DOWN


    @dutil.route(dmsg.GSProcessCreate, DTBL)
    def handle_process_create(self, msg):
        log = self._process_logger
        is_head_proc = self._state == self.RunState.WAITING_FOR_HEAD

        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)
        # here the constructor might have multiple steps so
        # it has to control its own pending completions.
        # The RunState is set to HAS_HEAD in this function as well.
        ProcessContext.construct(self, msg, reply_channel, is_head_proc)


    @dutil.route(dmsg.SHProcessCreateResponse, DTBL)
    def handle_process_create_response(self, msg):
        log = self._process_logger
        log.debug(f'process create response tag {msg.tag} ref {msg.ref}')
        assert msg.ref in self.pending
        handler = self.pending.pop(msg.ref)
        handler(msg)


    @dutil.route(dmsg.GSProcessList, DTBL)
    def handle_process_list(self, msg):
        log = self._process_logger
        log.debug(f'handling {msg!s}')

        reply_channel = self.get_reply_handle(msg)

        rm = dmsg.GSProcessListResponse(tag=self.tag_inc(),
                                        ref=msg.tag,
                                        err=dmsg.GSProcessListResponse.Errors.SUCCESS,
                                        plist=list(self.process_table.keys()))

        reply_channel.send(rm.serialize())
        log.debug(f'response to {msg!s}: {rm!s}')


    @dutil.route(dmsg.GSProcessQuery, DTBL)
    def handle_process_query(self, msg):
        log = self._process_logger
        log.debug(f'handling {msg}')

        reply_channel = self.get_reply_handle(msg)
        target_uid, found, errmsg = self.resolve_puid(msg.user_name, msg.t_p_uid)

        if not found:
            rm = dmsg.GSProcessQueryResponse(tag=self.tag_inc(),
                                             ref=msg.tag,
                                             err=dmsg.GSProcessQueryResponse.Errors.UNKNOWN,
                                             err_info=errmsg)
        else:
            pdesc = self.process_table[target_uid].descriptor
            rm = dmsg.GSProcessQueryResponse(tag=self.tag_inc(),
                                             ref=msg.tag,
                                             err=dmsg.GSProcessQueryResponse.Errors.SUCCESS,
                                             desc=pdesc)

        reply_channel.send(rm.serialize())
        log.debug(f'response to {msg!s}: {rm!s}')


    @dutil.route(dmsg.GSProcessKill, DTBL)
    def handle_process_kill(self, msg):
        log = self._process_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)
        ProcessContext.kill(self, msg, reply_channel)


    @dutil.route(dmsg.SHProcessKillResponse, DTBL)
    def handle_process_kill_response(self, msg):
        log = self._process_logger
        log.debug(f'handling {msg!s}')
        assert msg.ref in self.pending

        handler = self.pending.pop(msg.ref)
        handler(msg)


    @dutil.route(dmsg.GSProcessJoin, DTBL)
    def handle_process_join(self, msg):
        log = self._process_logger
        log.debug(f'handling {msg!s}')
        gspjr = dmsg.GSProcessJoinResponse
        reply_channel = self.get_reply_handle(msg)
        target_uid, found, errmsg = self.resolve_puid(msg.user_name, msg.t_p_uid)

        if not found:
            rm = gspjr(tag=self.tag_inc(),
                       ref=msg.tag,
                       err=gspjr.Errors.UNKNOWN,
                       err_info=errmsg)
            reply_channel.send(rm.serialize())
            log.debug(f'unknown: response to {msg}: {rm}')
        elif target_uid == msg.p_uid:
            rm = gspjr(tag=self.tag_inc(), ref=msg.tag,
                       err=gspjr.Errors.SELF, err_info='self-join forbidden')
            reply_channel.send(rm.serialize())
            log.debug(f'self join not allowed: {msg}')
        else:
            pctx = self.process_table[target_uid]
            pdesc = pctx.descriptor
            if process_desc.ProcessDescriptor.State.DEAD == pdesc.state:
                rm = gspjr(tag=self.tag_inc(),
                           ref=msg.tag,
                           err=gspjr.Errors.SUCCESS,
                           exit_code=pdesc.ecode)
                reply_channel.send(rm.serialize())
                log.debug(f'process dead, response to {msg}: {rm}')
            else:
                if msg.timeout < 0:
                    timeout = None
                else:
                    timeout = msg.timeout / 1000000.0  # microseconds

                if timeout is not None and timeout <= 0.000100:  # 100 microseconds from now, is now.
                    rm = gspjr(tag=self.tag_inc(), ref=msg.tag, err=gspjr.Errors.TIMEOUT)
                    reply_channel.send(rm.serialize())
                else:
                    self.pending_join.put(target_uid, msg, timeout=timeout)
                    log.debug(f'join stored for {target_uid}: {msg} timeout {timeout}')

    @dutil.route(dmsg.GSProcessJoinList, DTBL)
    def handle_process_join_list(self, msg):
        log = self._process_logger
        log.debug(f'handling {msg!s}')
        gspjlr = dmsg.GSProcessJoinListResponse
        reply_channel = self.get_reply_handle(msg)
        success_num = 0
        pending_puid = []

        # holds the status of the p_uids from the join list for this particular msg
        join_key = (msg.tag, msg.p_uid)
        self.puid_join_list_status[join_key] = dict()

        if msg.timeout < 0:
            timeout = None
        else:
            timeout = msg.timeout / 1000000.0  # microseconds

        # helper function which updates the puid_join_list_status dict with the status of each puid
        def proc_join_list_helper(target_uid, found, errmsg, success_num, item):
            if not found:
                self.puid_join_list_status[join_key][item] = (gspjlr.Errors.UNKNOWN.value, errmsg)
                log.debug(f'unknown target_uid: {item}')
            elif target_uid == msg.p_uid:
                self.puid_join_list_status[join_key][target_uid] = (gspjlr.Errors.SELF.value, 'self-join forbidden')
                log.debug(f'self join not allowed for target_uid: {target_uid}')
            else:
                pctx = self.process_table[target_uid]
                pdesc = pctx.descriptor
                if process_desc.ProcessDescriptor.State.DEAD == pdesc.state:
                    self.puid_join_list_status[join_key][target_uid] = (gspjlr.Errors.SUCCESS.value, pdesc.ecode)
                    success_num += 1
                    log.debug(f'process dead, target_uid: {target_uid}')
                else:
                    if timeout is not None and timeout <= 0.000100:  # 100 microseconds from now, is now.
                        self.puid_join_list_status[join_key][target_uid] = (gspjlr.Errors.TIMEOUT.value, None)
                    else:
                        pending_puid.append(target_uid)
                        self.puid_join_list_status[join_key][target_uid] = (gspjlr.Errors.PENDING.value, None)
            return success_num

        if msg.t_p_uid_list:
            for item in msg.t_p_uid_list:
                item = int(item)
                target_uid, found, errmsg = self.resolve_puid(None, item)
                success_num = proc_join_list_helper(target_uid, found, errmsg, success_num, item)
        if msg.user_name_list:
            for item in msg.user_name_list:
                target_uid, found, errmsg = self.resolve_puid(item, None)
                success_num = proc_join_list_helper(target_uid, found, errmsg, success_num, item)

        # Do not send a message if all the processes are still pending
        # if no exits no timeouts, or, 'all' option and some are pending (even though some may have exited)
        if (not success_num and pending_puid) or (msg.join_all and pending_puid):
            self.pending_join_list.put(frozenset(pending_puid), msg, timeout=timeout)
            log.debug(f'join stored for target_uid: {pending_puid} timeout {timeout}')
        else:
            rm = gspjlr(tag=self.tag_inc(), ref=msg.tag, puid_status=self.puid_join_list_status[join_key])
            reply_channel.send(rm.serialize())
            del self.puid_join_list_status[join_key]

    @dutil.route(dmsg.SHProcessExit, DTBL)
    def handle_process_exit(self, msg):
        log = self._process_logger
        log.info(f'rec {msg}')
        gspjr = dmsg.GSProcessJoinResponse
        gspjlr = dmsg.GSProcessJoinListResponse
        log.debug(f'p_uid {msg.p_uid} exited')

        ctx = self.process_table[msg.p_uid]
        ctx.descriptor.state = process_desc.ProcessDescriptor.State.DEAD
        ctx.descriptor.ecode = msg.exit_code

        if ctx.request.options.make_inf_channels:
            clean_chan = dmsg.GSChannelDestroy(tag=self.tag_inc(),
                                               p_uid=dfacts.GS_PUID, r_c_uid=0,
                                               c_uid=ctx.descriptor.gs_ret_cuid)
            ChannelContext.destroy(self, clean_chan, dutil.AbsorbingChannel())

        if ctx.stdin_context is not None:
            clean_chan = dmsg.GSChannelDestroy(tag=self.tag_inc(),
                                               p_uid=dfacts.GS_PUID, r_c_uid=0,
                                               c_uid=ctx.stdin_context.shchannelcreate_msg.c_uid, dec_ref=True)
            ChannelContext.destroy(self, clean_chan, dutil.AbsorbingChannel())

        if ctx.stdout_context is not None:
            clean_chan = dmsg.GSChannelDestroy(tag=self.tag_inc(),
                                               p_uid=dfacts.GS_PUID, r_c_uid=0,
                                               c_uid=ctx.stdout_context.shchannelcreate_msg.c_uid, dec_ref=True)
            ChannelContext.destroy(self, clean_chan, dutil.AbsorbingChannel())

        if ctx.stderr_context is not None:
            clean_chan = dmsg.GSChannelDestroy(tag=self.tag_inc(),
                                               p_uid=dfacts.GS_PUID, r_c_uid=0,
                                               c_uid=ctx.stderr_context.shchannelcreate_msg.c_uid, dec_ref=True)
            ChannelContext.destroy(self, clean_chan, dutil.AbsorbingChannel())

        if ctx.descriptor.p_p_uid in self.process_table:
            parent_ctx = self.process_table[ctx.descriptor.p_p_uid]
            parent_ctx.descriptor.live_children.discard(ctx.descriptor.p_uid)

        if msg.p_uid in self.pending_join:
            for join_req in self.pending_join.get(msg.p_uid):
                reply_channel = self.get_reply_handle(join_req)
                rm = gspjr(tag=self.tag_inc(),
                           ref=join_req.tag,
                           err=gspjr.Errors.SUCCESS,
                           exit_code=ctx.descriptor.ecode)
                reply_channel.send(rm.serialize())
                log.debug(f'join response to {join_req}: {rm}')
            self.pending_join.remove(msg.p_uid)

        to_be_erased = []
        join_reqs_to_be_erased = []
        to_be_added = []
        for puid_set in self.pending_join_list.keys():
            if msg.p_uid in puid_set:
                for join_req in self.pending_join_list.get(puid_set):
                    join_key = (join_req.tag, join_req.p_uid)
                    self.puid_join_list_status[join_key][msg.p_uid] = (gspjlr.Errors.SUCCESS.value, ctx.descriptor.ecode)

                    # if this is multi_join on 'any' processes request, or, set with a single member
                    if not join_req.join_all or len(puid_set) == 1:
                        reply_channel = self.get_reply_handle(join_req)
                        rm = gspjlr(tag=self.tag_inc(), ref=join_req.tag,
                                    puid_status=self.puid_join_list_status[join_key])
                        reply_channel.send(rm.serialize())
                        log.debug(f'join response to {join_req}: {rm}')

                        # we are done with this join request, so let's clean things up
                        del self.puid_join_list_status[join_key]
                        join_reqs_to_be_erased.append((puid_set, join_req))

                    else:  # this is multi_join on 'all' processes request
                        # we create a new entry for waiting on the new set puid_set - msg.p_uid
                        new_puid_set = set(puid_set)  # puid_set is a frozenset
                        new_puid_set.remove(msg.p_uid)
                        to_be_added.append((puid_set, new_puid_set, join_req))
                to_be_erased.append(puid_set)

        for oldset, puidset, joinreq in to_be_added:
            if joinreq.timeout < 0:
                timeout = None
            elif joinreq.timeout == 0:
                timeout = 0
            else:
                timeout = self.pending_join_list.time_left(oldset, joinreq) - time.time()  # this is the remaining time on waiting until timeout
            self.pending_join_list.put(frozenset(puidset), joinreq, timeout=timeout)

        for puid_set, join_req in join_reqs_to_be_erased:
            self.pending_join_list.remove_one(puid_set, join_req)

        for puid_set in to_be_erased:
            self.pending_join_list.remove(puid_set)

        if msg.p_uid == self.head_puid:
            log.info(f'head process id {self.head_puid} exited')
            self.bela_input.send(dmsg.GSHeadExit(tag=self.tag_inc(),
                                                 exit_code=msg.exit_code).serialize())
            # TODO: head process exit cleanup
            self._state = self.RunState.WAITING_FOR_HEAD
            self.head_puid = None

        # GROUP related work
        # if this process is a member of a group, then update the group context
        if msg.p_uid in self.resource_to_group_map:
            guid, (lst_idx, item_idx) = self.resource_to_group_map[msg.p_uid]
            log.debug(f'process {msg.p_uid} just exited and belongs to group {guid}')
            groupctx = self.group_table[guid]
            groupctx.descriptor.sets[lst_idx][item_idx].state = process_desc.ProcessDescriptor.State.DEAD
            groupctx.descriptor.sets[lst_idx][item_idx].desc.state = process_desc.ProcessDescriptor.State.DEAD

            # if group destroy was called
            if msg.p_uid in self.pending_group_destroy:
                # this is needed to complete the pending continuation on the group_int
                self.group_to_pending_resource_map[(msg.p_uid, guid)] = self.process_table[msg.p_uid]
                handler = self.pending_group_destroy.pop(msg.p_uid)
                handler(msg)


    @dutil.route(dmsg.GSPoolCreate, DTBL)
    def handle_pool_create(self, msg):
        log = self._pool_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)
        issued, tag, context = PoolContext.construct(self, msg, reply_channel)
        if issued:
            self.pending[tag] = context.complete_construction


    @dutil.route(dmsg.SHPoolCreateResponse, DTBL)
    def handle_pool_create_response(self, msg):
        log = self._pool_logger
        log.debug(f'create response tag {msg.tag} ref {msg.ref} received')
        assert msg.ref in self.pending
        handler = self.pending.pop(msg.ref)
        handler(msg)


    @dutil.route(dmsg.GSPoolList, DTBL)
    def handle_pool_list(self, msg):
        log = self._pool_logger
        log.debug(f'handling {msg!s}')

        reply_channel = self.get_reply_handle(msg)

        rm = dmsg.GSPoolListResponse(tag=self.tag_inc(),
                                     ref=msg.tag,
                                     err=dmsg.GSPoolListResponse.Errors.SUCCESS,
                                     mlist=list(self.pool_table.keys()))

        reply_channel.send(rm.serialize())
        log.debug(f'response to {msg!s}: {rm!s}')


    @dutil.route(dmsg.GSPoolQuery, DTBL)
    def handle_pool_query(self, msg):
        log = self._pool_logger
        log.debug(f'handling {msg}')

        reply_channel = self.get_reply_handle(msg)
        target_uid, found, errmsg = self.resolve_muid(msg.user_name, msg.m_uid)

        if not found:
            rm = dmsg.GSPoolQueryResponse(tag=self.tag_inc(),
                                          ref=msg.tag,
                                          err=dmsg.GSPoolQueryResponse.Errors.UNKNOWN,
                                          err_info=errmsg)
        else:
            pdesc = self.pool_table[target_uid].descriptor
            rm = dmsg.GSPoolQueryResponse(tag=self.tag_inc(),
                                          ref=msg.tag,
                                          err=dmsg.GSPoolQueryResponse.Errors.SUCCESS,
                                          desc=pdesc)

        reply_channel.send(rm.serialize())

        log.debug(f'response to {str(msg)}: {str(rm)}')


    @dutil.route(dmsg.GSPoolDestroy, DTBL)
    def handle_pool_destroy(self, msg):
        log = self._pool_logger
        log.debug(f'handling {str(msg)}')
        reply_channel = self.get_reply_handle(msg)
        PoolContext.destroy(self, msg, reply_channel)


    @dutil.route(dmsg.SHPoolDestroyResponse, DTBL)
    def handle_pool_destroy_response(self, msg):
        log = self._pool_logger
        log.debug(f'handling {msg}')
        assert msg.ref in self.pending
        handler = self.pending.pop(msg.ref)
        handler(msg)


    @dutil.route(dmsg.GSChannelCreate, DTBL)
    def handle_channel_create(self, msg):
        log = self._channel_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)
        issued, tag, context = ChannelContext.construct(self, msg, reply_channel)
        if issued:
            self.pending[tag] = context.complete_construction


    @dutil.route(dmsg.SHChannelCreateResponse, DTBL)
    def handle_channel_create_response(self, msg):
        log = self._channel_logger
        log.debug(f'create response tag {msg.tag} ref {msg.ref} received')

        assert msg.ref in self.pending
        handler = self.pending.pop(msg.ref)
        handler(msg)


    @dutil.route(dmsg.GSChannelDestroy, DTBL)
    def handle_channel_destroy(self, msg):
        log = self._channel_logger
        log.debug(f'handling {msg!s}')
        if msg.reply_req:
            reply_channel = self.get_reply_handle(msg)
        else:
            reply_channel = dutil.AbsorbingChannel()

        ChannelContext.destroy(self, msg, reply_channel)


    @dutil.route(dmsg.SHChannelDestroyResponse, DTBL)
    def handle_channel_destroy_response(self, msg):
        log = self._channel_logger
        log.debug(f'handling {msg}')
        assert msg.ref in self.pending
        handler = self.pending.pop(msg.ref)
        handler(msg)


    @dutil.route(dmsg.GSChannelList, DTBL)
    def handle_channel_list(self, msg):
        log = self._channel_logger
        log.debug(f'handling {msg!s}')

        reply_channel = self.get_reply_handle(msg)

        rm = dmsg.GSChannelListResponse(tag=self.tag_inc(),
                                        ref=msg.tag,
                                        err=dmsg.GSChannelListResponse.Errors.SUCCESS,
                                        clist=list(self.channel_table.keys()))

        reply_channel.send(rm.serialize())
        log.debug(f'response to {str(msg)}: {str(rm)}')


    @dutil.route(dmsg.GSChannelQuery, DTBL)
    def handle_channel_query(self, msg):
        log = self._channel_logger
        # log.debug(f'handling {msg}')

        reply_channel = self.get_reply_handle(msg)
        # log.debug(f'channel query to {msg.user_name} - {msg.c_uid}')
        target_uid, found, errmsg = self.resolve_cuid(msg.user_name, msg.c_uid)

        if not found:
            rm = dmsg.GSChannelQueryResponse(tag=self.tag_inc(),
                                             ref=msg.tag,
                                             err=dmsg.GSChannelQueryResponse.Errors.UNKNOWN,
                                             err_info=errmsg)
        else:
            cctx = self.channel_table[target_uid]
            cdesc = cctx.descriptor

            if cctx.is_refcounted and msg.inc_refcnt and cdesc.State.ACTIVE == cdesc.state:
                cctx.refcnt += 1
                log.debug(f"channel {cctx.descriptor.c_uid}: refcnt={cctx.refcnt}")

            rm = dmsg.GSChannelQueryResponse(tag=self.tag_inc(), ref=msg.tag,
                                             err=dmsg.GSChannelQueryResponse.Errors.SUCCESS,
                                             desc=cdesc)
            # log.debug(f'found descriptor: {cdesc}')

        reply_channel.send(rm.serialize())
        # log.debug(f'response to {msg!s}: {rm!s}')


    @dutil.route(dmsg.GSChannelJoin, DTBL)
    def handle_channel_join(self, msg):
        log = self._channel_logger
        log.debug(f'handling {msg}')

        target_uid, found, errmsg = self.resolve_cuid(msg.name)

        handled_immediately = False
        if found:
            cdesc = self.channel_table[target_uid].descriptor
            if cdesc.State.ACTIVE == cdesc.state:
                reply_channel = self.get_reply_handle(msg)
                rm = dmsg.GSChannelJoinResponse(tag=self.tag_inc(), ref=msg.tag,
                                                err=dmsg.GSChannelJoinResponse.Errors.SUCCESS,
                                                desc=cdesc)
                reply_channel.send(rm.serialize())
                log.debug(f'response to {msg!s}: {rm!s}')
                handled_immediately = True
            elif cdesc.State.DEAD == cdesc.state:
                reply_channel = self.get_reply_handle(msg)
                rm = dmsg.GSChannelJoinResponse(tag=self.tag_inc(),
                                                ref=msg.tag,
                                                err=dmsg.GSChannelJoinResponse.Errors.DEAD,
                                                desc=cdesc)
                reply_channel.send(rm.serialize())
                log.debug(f'response to {msg!s}: {rm!s}')
                handled_immediately = True
            else:  # close case
                assert cdesc.State.PENDING == cdesc.state

        # In two cases we remember the join request:
        # - nobody has asked this channel to be created yet
        # - the channel is in the process of creation
        # In either case this gets handled when the channel create response
        # is returned or the timeout expires.
        if not handled_immediately:
            if msg.timeout < 0:
                the_timeout = None
            else:
                the_timeout = msg.timeout / 1000000.0  # microseconds

            self.pending_channel_joins.put(msg.name, msg, timeout=the_timeout)
            log.debug(f'channel join stored for {target_uid}: {msg} timeout {the_timeout}')


    @dutil.route(dmsg.GSDump, DTBL)
    def handle_dump(self, msg):
        log = self._process_logger
        log.debug(f'dumping to {msg.filename}')

        with open(msg.filename, 'w') as fh:
            self.dump_to(fh)

        log.debug('dump complete')


    @dutil.route(dmsg.GSNodeList, DTBL)
    def handle_node_list(self, msg):
        log = self._pool_logger
        log.debug(f'handling {msg!s}')

        reply_channel = self.get_reply_handle(msg)

        rm = dmsg.GSNodeListResponse(tag=self.tag_inc(),
                                     ref=msg.tag,
                                     err=dmsg.GSNodeListResponse.Errors.SUCCESS,
                                     hlist=list(self.node_table.keys()))

        reply_channel.send(rm.serialize())
        log.debug(f'response to {msg!s}: {rm!s}')


    @dutil.route(dmsg.GSNodeQueryTotalCPUCount, DTBL)
    def handle_node_query_total_cpu_count(self, msg):
        log = self._node_logger

        log.debug(f'handling {msg}')

        reply_channel = self.get_reply_handle(msg)

        total_cpus = 0
        for desc in self.node_table.values():
            total_cpus += desc.descriptor.num_cpus

        rm = dmsg.GSNodeQueryTotalCPUCountResponse(tag=self.tag_inc(), ref=msg.tag,
                                                   err=dmsg.GSNodeQueryTotalCPUCountResponse.Errors.SUCCESS,
                                                   total_cpus=total_cpus)

        reply_channel.send(rm.serialize())
        log.debug(f'response to {msg!s}: {rm!s}')


    @dutil.route(dmsg.GSNodeQuery, DTBL)
    def handle_node_query(self, msg):
        log = self._node_logger
        log.debug(f'handling {msg}')

        reply_channel = self.get_reply_handle(msg)

        ## Get descriptor
        target_uid, found, errmsg = self.resolve_huid(msg.name, msg.h_uid)

        sdesc = self.node_table[target_uid].descriptor.sdesc

        if not found:
            log.error(errmsg)
            rm = dmsg.GSNodeQueryResponse(tag=self.tag_inc(),
                                          ref=msg.tag,
                                          err=dmsg.GSNodeQueryResponse.Errors.UNKNOWN,
                                          err_info=errmsg)
        else:
            rm = dmsg.GSNodeQueryResponse(tag=self.tag_inc(), ref=msg.tag,
                                          err=dmsg.GSNodeQueryResponse.Errors.SUCCESS,
                                          desc=sdesc)
            log.debug(f'found descriptor={sdesc}')

        reply_channel.send(rm.serialize())
        log.debug(f'response to {msg!s}: {rm!s}')


    @dutil.route(dmsg.GSGroupCreate, DTBL)
    def handle_group_create(self, msg):
        log = self._group_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)

        # here the constructor has to control its
        # own pending completions
        GroupContext.construct(self, msg, reply_channel)


    @dutil.route(dmsg.GSGroupKill, DTBL)
    def handle_group_kill(self, msg):
        log = self._group_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)

        # here the constructor has to control its
        # own pending completions
        GroupContext.kill(self, msg, reply_channel)


    @dutil.route(dmsg.GSGroupDestroy, DTBL)
    def handle_group_destroy(self, msg):
        log = self._group_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)

        # here the constructor has to control its
        # own pending completions
        GroupContext.destroy(self, msg, reply_channel)


    @dutil.route(dmsg.GSGroupAddTo, DTBL)
    def handle_group_add_to(self, msg):
        log = self._group_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)

        GroupContext.add(self, msg, reply_channel)


    @dutil.route(dmsg.GSGroupCreateAddTo, DTBL)
    def handle_group_create_add_to(self, msg):
        log = self._group_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)

        # here the constructor has to control its
        # own pending completions
        GroupContext.create_add(self, msg, reply_channel)


    @dutil.route(dmsg.GSGroupRemoveFrom, DTBL)
    def handle_group_remove_from(self, msg):
        log = self._group_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)

        GroupContext.remove(self, msg, reply_channel)


    @dutil.route(dmsg.GSGroupDestroyRemoveFrom, DTBL)
    def handle_group_destroy_remove_from(self, msg):
        log = self._group_logger
        log.debug(f'handling {msg!s}')
        reply_channel = self.get_reply_handle(msg)

        # here the constructor has to control its
        # own pending completions
        GroupContext.destroy_remove(self, msg, reply_channel)


# separate process server entry:
# python3 -c "import dragon.globalservices.server as dgs; dgs.single()"
def single():
    dlog.setup_logging(basename=os.environ.get(dfacts.GS_LOG_BASE, 'gs'),
                       level=logging.DEBUG)
    log = logging.getLogger('single entry')
    log.info('Starting GS in single mode')
    log.info(f'start in pid {os.getpid()}')
    the_ctx = GlobalContext()
    log.info('GS context constructed, running startup')

    gs_stdout = dutil.NewlineStreamWrapper(sys.stdout)

    the_test_conns = {}

    if dfacts.GS_TEST_CH_EV in os.environ:
        log.debug('getting test channels')
        channel_init_bytes = B64.str_to_bytes(os.environ[dfacts.GS_TEST_CH_EV])
        channel_init = pickle.loads(channel_init_bytes)

        for initvals in channel_init:
            cuid, spoold, schand, rdng = initvals
            log.debug(f'attaching to test channel: cuid {cuid} reading {rdng}')
            chan = dch.Channel.attach(schand)
            if rdng:
                the_test_conns[cuid] = dconn.Connection(inbound_initializer=chan, policy=dparm.POLICY_INFRASTRUCTURE)
            else:
                the_test_conns[cuid] = dconn.Connection(outbound_initializer=chan, policy=dparm.POLICY_INFRASTRUCTURE)

    else:
        log.debug('no test channels here!')

    try:
        the_ctx.run_startup()
    except Exception as e:
        msg = f'fatal startup error: {type(e)}:{e}'
        log.exception(msg)
        gs_stdout.send(dmsg.AbnormalTermination(tag=0, err_info=msg).serialize())
        exit(dfacts.GS_ERROR_EXIT)

    try:
        the_ctx.run_global_server(test_conns=the_test_conns)
    except Exception as e:
        msg = f'fatal error: {type(e)}:{e}'
        log.exception(msg)
        gs_stdout.send(dmsg.AbnormalTermination(tag=0, err_info=msg).serialize())
        exit(dfacts.GS_ERROR_EXIT)

    try:
        the_ctx.run_teardown()
    except Exception as e:
        msg = f'fatal shutdown error: {type(e)}:{e}'
        log.exception(msg)
        gs_stdout.send(dmsg.AbnormalTermination(tag=0, err_info=msg).serialize())
        exit(dfacts.GS_ERROR_EXIT)


# thread internal entry
def single_thread(gs_stdout):
    log = logging.getLogger('gs')
    log.info('Starting GS in single thread mode')
    log.info(f'start in {os.getpid()}')
    the_ctx = GlobalContext(gs_stdout)
    log.info('GS context constructed, running startup')

    try:
        the_ctx.run_startup(mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE)
        the_ctx.run_global_server(mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE)
        the_ctx.run_teardown()
        log.info('normal exit')
    except Exception as e:
        log.exception('fatal error')
        gs_stdout.send(dmsg.AbnormalTermination(tag=0, err_info='fatal exception').serialize())
        time.sleep(1)
        raise e


# server entry:
# python3 -c "import dragon.globalservices.server as dgs; dgs.multi()"
def multi():

    set_procname(dfacts.PROCNAME_GS)
    dch.register_gateways_from_env()

    # Set up logging
    _, fname = dlog.setup_BE_logging(service=dls.GS)
    log = logging.getLogger(dls.GS).getChild('server.multi')
    log.info('Starting GS in multi mode')
    log.debug('getting logger sdesc')
    logger_sdesc = os.environ.get(dfacts.DRAGON_LOGGER_SDESC)
    if logger_sdesc is not None:
        logger_sdesc = B64.from_str(logger_sdesc)
    dlog.setup_BE_logging(service=dls.GS, logger_sdesc=logger_sdesc, fname=fname)
    log.info('gs logging stood up')

    the_ctx = GlobalContext()

    the_test_conns = {}
    if dfacts.GS_TEST_CH_EV in os.environ:
        log.debug('getting test channels')
        channel_init_bytes = B64.str_to_bytes(os.environ[dfacts.GS_TEST_CH_EV])
        channel_init = pickle.loads(channel_init_bytes)

        for initvals in channel_init:
            cuid, spoold, schand, rdng = initvals
            log.debug(f'attaching to test channel: cuid {cuid} reading {rdng}')
            chan = dch.Channel.attach(schand)
            if rdng:
                the_test_conns[cuid] = dconn.Connection(inbound_initializer=chan, policy=dparm.POLICY_INFRASTRUCTURE)
            else:
                the_test_conns[cuid] = dconn.Connection(outbound_initializer=chan, policy=dparm.POLICY_INFRASTRUCTURE)
    else:
        log.debug('no test channels here!')

    try:
        the_ctx.run_startup(mode=the_ctx.LaunchModes.MULTI)
    except Exception as e:
        msg = f'fatal startup error: {type(e)}:{e}'
        log.exception(msg)
        the_ctx.msg_stdout.send(dmsg.AbnormalTermination(tag=0, err_info=msg).serialize())
        exit(dfacts.GS_ERROR_EXIT)

    try:
        the_ctx.run_global_server(mode=the_ctx.LaunchModes.MULTI, test_conns=the_test_conns)
    except Exception as e:
        msg = f'fatal error: {type(e)}:{e}'
        log.exception(msg)
        the_ctx.msg_stdout.send(dmsg.AbnormalTermination(tag=0, err_info=msg).serialize())
        exit(dfacts.GS_ERROR_EXIT)

    try:
        the_ctx.run_teardown()
    except Exception as e:
        msg = f'fatal shutdown error: {type(e)}:{e}'
        log.exception(msg)
        the_ctx.msg_stdout.send(dmsg.AbnormalTermination(tag=0, err_info=msg).serialize())

src/dragon/globalservices/server.py
DragonHPC/dragon
"""Startup processing for Global Services server.

    These functions implement the infrastructure startup sequence.
"""

import logging
import os
import sys

import dragon.channels as dch

import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.util as dutil
from dragon.dlogging.util import DragonLoggingServices as dls
from dragon.utils import B64


class StartupError(Exception):
    pass


def single_connect_to_default_channels(gs_input, shep_input, bela_input):
    """Connects to default channels.

        Any argument that is not None is returned un
        Purpose of doing it like this is to make it easy to substitute
        different handles in test.

        :return: Connection objects or overrides from parameters.
        :rtype: tuple, of gs_input, shep_input, bela_input
    """

    # if shep_input or gs_input or bela_input isn't a message send/recv obj, do
    # the work, and if it is an object just echo it because
    # we are in some test bench scenario.

    if gs_input is None:
        gs_input_cd = B64.str_to_bytes(dparm.this_process.gs_cd)
        gs_input_channel = dch.Channel.attach(gs_input_cd)
        gs_input = dconn.Connection(inbound_initializer=gs_input_channel, policy=dparm.POLICY_INFRASTRUCTURE)

    if shep_input is None:
        shep_input_cd = B64.str_to_bytes(dparm.this_process.local_shep_cd)
        shep_input_channel = dch.Channel.attach(shep_input_cd)
        shep_input = dconn.Connection(outbound_initializer=shep_input_channel, policy=dparm.POLICY_INFRASTRUCTURE)

    if bela_input is None:
        bela_input_cd = B64.str_to_bytes(dparm.this_process.local_be_cd)
        bela_input_channel = dch.Channel.attach(bela_input_cd)
        bela_input = dconn.Connection(outbound_initializer=bela_input_channel, policy=dparm.POLICY_INFRASTRUCTURE)

    return gs_input, shep_input, bela_input


def startup_single(the_ctx, gs_input=None, shep_input=None, bela_input=None):
    """Single node Global Services startup.

        This goes through the protocol steps needed for global services
        to come up on a single node.

        The gs_input, shep_input, bela_input parameters are here for test.  Normally
        they are gotten from interacting with the Channels library.

        Todo: add a timeout to this

        :param the_ctx: global services runtime context object
        :param shep_input: testbench override handle to Shepherd input
        :param gs_input: testbench override handle to Global Services Server (this process) input
        :param bela_input: testbench override handle to back end/launcher input channel.
        :raises StartupError: Reports errors in starting up.

        :return: shepherd input channel and gs input channel
        :rtype: tuple: ([shepherd input channel], gs input channel)
    """

    gait = the_ctx.tag_inc  # 'get and inc tag'

    log = logging.getLogger('startup_single')

    log.info('single node startup procedure entered')
    inputs = single_connect_to_default_channels(gs_input, shep_input, bela_input)
    gs_input, shep_input, bela_input = inputs
    log.info('connected to channels')

    shep_input.send(dmsg.GSPingSH(tag=gait()).serialize())
    log.debug('handshake message sent to Shepherd')

    return [shep_input], gs_input, bela_input, 1 # == number of nodes in single mode


def startup_multi(the_ctx, gs_input=None, shep_inputs=None, bela_input=None):
    """Multi-node Global services startup

        :param the_ctx: global services runtime context object
        :param shep_inputs: List of handles to Shepherd input for debug setup
        :param gs_input: handle to Global Services Server (this process) input
        :param bela_input: handle to back end/launcher input channel.

    """
    gait = the_ctx.tag_inc  # 'get and inc tag'

    log = logging.getLogger(dls.GS).getChild('startup.startup_multi')
    log.info('beginning gs startup...')

    # Get the LAChannelsInfo from ls via stdin
    ls_stdin = os.fdopen(sys.stdin.fileno(), 'rb')
    ls_stdin_recv = dutil.NewlineStreamWrapper(ls_stdin, write_intent=False)
    la_chs_info = dmsg.parse(ls_stdin_recv.recv())
    log.info(f'received all channels info, LAChannelsInfo - m9')
    log.debug(f'la_channels.nodes_desc: {la_chs_info.nodes_desc}')
    log.debug(f'la_channels.gs_cd: {la_chs_info.gs_cd}')

    # Attach to the gs channel for recv'ing and the node_idx==0 ls for this test bringup
    gs_ch = la_chs_info.gs_cd
    gs_in_ch = dch.Channel.attach(B64.str_to_bytes(gs_ch))
    gs_in_rh = dconn.Connection(inbound_initializer=gs_in_ch, policy=dparm.POLICY_INFRASTRUCTURE)
    log.info('attached to its recv channel - a17')

    ls_in_whs = []
    n_ls = len(la_chs_info.nodes_desc)
    for val in range(n_ls):
        ls_ch = dch.Channel.attach(B64.str_to_bytes(la_chs_info.nodes_desc[str(val)].shep_cd))
        ls_in_whs.append(dconn.Connection(outbound_initializer=ls_ch, policy=dparm.POLICY_INFRASTRUCTURE))
    log.info('GS now attached to all ls channels a17')

    # Ping all the shepherds and wait for responses.
    for ls_in_wh in ls_in_whs:
        ls_in_wh.send(dmsg.GSPingSH(tag=gait()).serialize())
    log.info('pinged all LS (GSPingSH) - m10')

    # Attach to the launcher backend on this node.
    be_ch = dparm.this_process.local_be_cd
    be_in_ch = dch.Channel.attach(B64.from_str(be_ch).decode())
    be_in_wh = dconn.Connection(outbound_initializer=be_in_ch, policy=dparm.POLICY_INFRASTRUCTURE)
    log.info('GS attached to its launcher BE channel - a17')

    return ls_in_whs, gs_in_rh, be_in_wh, n_ls

src/dragon/globalservices/startup.py
DragonHPC/dragon
"""A class that uniquely defines  properties of a managed channel for infrastructure communication.
"""

import enum

import dragon.infrastructure.parameters as dparm
import dragon.localservices.options as dso
import dragon.utils as du


class ChannelDescriptor:
    """Globally available description of a managed channel.

    Attributes:

        node : nonnegative int
            which node this channel

        name : str
            user assigned unique name of the process

        c_uid : nonnegative int
            unique channel id

        m_uid : nonnegative int
            memory pool containing this channel

        sdesc : bytes-like object
            serialized descriptor from the Channels library

        state : cls.State enum member
            current run state of this process
    """

    @enum.unique
    class State(enum.Enum):
        """Current state of the channel."""
        PENDING = enum.auto()
        ACTIVE = enum.auto()
        DEAD = enum.auto()

    def __init__(self, m_uid, c_uid, name, node, state=None, sdesc=''):
        """
            sdesc - serialized library descriptor
        """
        self.m_uid = m_uid
        self.c_uid = c_uid
        self.name = name
        self.sdesc = sdesc
        self.node = node

        if state is None:
            self.state = self.State.PENDING
        else:
            if isinstance(state, self.State):
                self.state = state
            elif isinstance(state, int):
                self.state = self.State(state)
            else:
                raise NotImplementedError('unknown state init')

    @property
    def sdesc(self):
        """Returns the serialized library descriptor of the channel."""
        return self._sdesc

    @sdesc.setter
    def sdesc(self, value):
        """Populate the serialized library descriptor of the channel.

        :param value: value to be populated as the serialzied descriptor
        :type value: either bytes (or) string passed through the environment
        """
        if isinstance(value, str):
            self._sdesc = du.B64.str_to_bytes(value)
        else:
            self._sdesc = value

    def __str__(self):
        return f'{self.name}:{self.c_uid} {self.state.name} on node {self.node}'

    def get_sdict(self):
        """Collect the entire information of the channel.

        :return: A dictionary with all key-value pairs of the available description of the channel.
        :rtype: Dictionary
        """
        rv = {'sdesc': du.B64.bytes_to_str(self.sdesc),
              'name': self.name,
              'm_uid': self.m_uid,
              'c_uid': self.c_uid,
              'node': self.node,
              'state': self.state.value}

        return rv

    @classmethod
    def from_sdict(cls, sdict):
        """Returns the descriptor of the managed channel."""
        return ChannelDescriptor(**sdict)


class ChannelOptions:
    """Open type, placeholder, for options on how to make a channel

    Attributes:

    local_opts
        local channel creation options, such as block size and capacity.
        see dragon.localservices.options.ChannelOptions

    ref_count
        bool, default False, whether GS should count references on this
        channel and destroy if count reaches zero

        NOTE: these options will be growing and the scheme needs
        thought for organization.
    """

    def __init__(self, *, local_opts=None, ref_count=False):
        self.ref_count = ref_count
        if local_opts is None:
            self.local_opts = dso.ChannelOptions()
        elif isinstance(local_opts, dso.ChannelOptions):
            self.local_opts = local_opts
        else:
            self.local_opts = dso.ChannelOptions(**local_opts)

    def get_sdict(self):
        """Collect the different attributes/options of the channel.

        :return: A dictionary with key-value pairs of the channel options.
        :rtype: Dictionary
        """
        return {'local_opts': self.local_opts.get_sdict(),
                'ref_count': self.ref_count}

    @staticmethod
    def from_sdict(sdict):
        """Returns the options of the managed channel."""
        return ChannelOptions(**sdict)

src/dragon/infrastructure/channel_desc.py
DragonHPC/dragon
"""Internal objects required for infrastructure process communication.
"""


import multiprocessing  # for BufferTooShort
import enum
import numbers
import pickle
import threading
import logging
import time
from typing import ClassVar

import dragon.dtypes as dtypes
import dragon.channels as dch
from .channel_desc import ChannelOptions
import dragon.globalservices.channel as dgchan
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts
from dragon.managed_memory import MemoryPool
from dragon.utils import B64

LOG = logging.getLogger('infrastructure.connection')

class ConnectionOptions:
    """Options for dragon.Connections

        Separate object because this will expand over time.

        There are a number of different facts about how one might
        want a Connection object to behave vs. its underlying
        Channels and its interaction with Global Services, together
        with anticipated need to let someone customize how the object
        behaves for performance.  This object is meant to organize
        all these into one object.
    """

    class CreationPolicy(enum.Enum):
        EXTERNALLY_MANAGED = enum.auto()
        PRE_CREATED = enum.auto()
        RECEIVER_CREATES = enum.auto()
        SENDER_CREATES = enum.auto()

    def __init__(self, *, creation_policy=CreationPolicy.EXTERNALLY_MANAGED,
                 min_block_size=None, large_block_size=None, huge_block_size=None, default_pool=None):
        self.min_block_size = min_block_size
        self.large_block_size = large_block_size
        self.huge_block_size = huge_block_size
        self.creation_policy = creation_policy
        self.default_pool = default_pool

class PipeOptions:
    """Options for the Pipe factory.

        Separate object because this, too, will expand over time.

        Same rationale as for ConnectionOptions.
    """

    # TODO: PE-38342, CreationPolicy.FIRST for first user to make channel
    # where-ever they happen to be, joining with a soft create
    class CreationPolicy(enum.Enum):
        EXTERNAL = enum.auto()
        EARLY = enum.auto()
        RECEIVER_CREATES = enum.auto()
        SENDER_CREATES = enum.auto()

    def __init__(self, *,
                 creation_policy=CreationPolicy.EARLY,
                 conn_options=None):

        self.creation_policy = creation_policy

        if conn_options is None:
            self.conn_options = ConnectionOptions()
        else:
            self.conn_options = conn_options

        if self.creation_policy == self.CreationPolicy.EARLY:
            self.conn_options.creation_policy = self.conn_options.CreationPolicy.PRE_CREATED
        elif self.creation_policy == self.CreationPolicy.SENDER_CREATES:
            self.conn_options.creation_policy = self.conn_options.CreationPolicy.SENDER_CREATES
        elif self.creation_policy == self.CreationPolicy.RECEIVER_CREATES:
            self.conn_options.creation_policy = self.conn_options.CreationPolicy.RECEIVER_CREATES
        elif self.creation_policy == self.CreationPolicy.EXTERNAL:
            self.creation_policy = self.conn_options.CreationPolicy.EXTERNALLY_MANAGED
        else:
            raise NotImplementedError('close case')


class CUID(int):
    """A wrapper around int to allow waiting on a cuid.
    """

    def thread_wait(self, timeout, done_ev, ready):
        if timeout is None:
            timeout = 1000000
        start = time.monotonic()
        delta = min(0.01, timeout) # TODO: figure out the value for delta
        timed_out = False

        if not done_ev.is_set():
            while not timed_out and not done_ev.is_set():
                done_res = self.conn.poll(delta)
                if not done_ev.is_set():
                    if done_res:
                        ready.append(self)
                        done_ev.set()
                timed_out = (time.monotonic() - start) > timeout
            done_ev.set()  # if not set here, timed out.


class Connection:
    """Dragon infrastructure's internal connection class. """
    class State(enum.Enum):
        CLOSED = enum.auto()  # we have closed and may have deleted backing channels we think we own
        ATTACHED = enum.auto()  # we have channels and we're attached, but haven't sent anything
        READY = enum.auto()  # we have got channel objects and a send was called
        UNBACKED = enum.auto()  # initialized with names but we have not created yet
        UNATTACHED = enum.auto()  # initialized with descriptors but we have not attached yet

    def __init__(self, *, inbound_initializer=None, outbound_initializer=None,
                 options=None, policy=dparm.POLICY_USER):
        """Initializes the Connection object.

            The object is initialized in an inbound direction, an outbound
            direction, or both.  These initializers can be either
            - a Channel object, whose lifecycle is assumed externally managed
            - a bytes-like descriptor, attached to lazily, destroyed automatically
            - a string name, for a rendezvous in Global services.

            TODO: document current behavior and design intent MUCH more thoroughly

        :param inbound_initializer: an initializer for the inbound channel
        :param outbound_initializer: an initializer for the outbound channel
        :param options: ConnectionOptions object
        :param policy: A Policy object.
        """

        self.at_eof = False
        self.ghost = False  # whether to omit sending EOT on close.

        if inbound_initializer is None and outbound_initializer is None:
            raise ConnectionError('at least one initializer is required')

        if options is None:
            self.options = ConnectionOptions()
        else:
            self.options = options

        self.write_adapter_options = {}
        if self.options.min_block_size is not None:
            self.write_adapter_options["small_blk_size"] = self.options.min_block_size

        if self.options.large_block_size is not None:
            self.write_adapter_options["large_blk_size"] = self.options.large_block_size

        if self.options.huge_block_size is not None:
            self.write_adapter_options["huge_blk_size"] = self.options.huge_block_size

        if not isinstance(policy, dparm.Policy):
            raise ValueError('The class of service must be a dragon.infrastructure.parameters.Policy value.')

        self.policy = policy
        self.inbound_chan = None
        self.outbound_chan = None
        self.write_adapter = None
        self.read_adapter = None
        self._detach_pool = None
        my_cp = self.options.creation_policy

        if isinstance(inbound_initializer, dch.Channel) or isinstance(outbound_initializer, dch.Channel):
            if my_cp != my_cp.EXTERNALLY_MANAGED and my_cp != my_cp.PRE_CREATED:
                msg = f'init with Channels and {my_cp.name} lifecycle not supported'
                raise ConnectionError(msg)

            self.inbound_chan = inbound_initializer
            self.outbound_chan = outbound_initializer
            self.inbound_initializer = None
            self.outbound_initializer = None
            self.state = self.State.ATTACHED
            if my_cp == my_cp.PRE_CREATED:
                self._incref_inbound_chan()
        else:
            self.inbound_initializer = inbound_initializer
            self.outbound_initializer = outbound_initializer
            if isinstance(inbound_initializer, str) or isinstance(outbound_initializer, str):
                self.state = self.State.UNBACKED
            else:  # assume channel descriptor
                if my_cp != my_cp.PRE_CREATED:
                    msg = f'init with Channels and {my_cp.name} lifecycle not supported'
                    raise ConnectionError(msg)
                self.state = self.State.UNATTACHED

    def _incref_inbound_chan(self):
        if isinstance(self.inbound_chan, dch.Channel):
            dgchan.get_refcnt(self.inbound_chan.cuid)

    def _decref_inbound_chan(self):
        if isinstance(self.inbound_chan, dch.Channel):
            dgchan.release_refcnt(self.inbound_chan.cuid)

    def __getstate__(self):
        state = dict(self.__dict__)
        state["write_adapter"] = None
        state["read_adapter"] = None
        state["_detach_pool"] = None
        # assume we're passing this to another process. Not a good final solution.
        if self.options.creation_policy == ConnectionOptions.CreationPolicy.PRE_CREATED:
            self._incref_inbound_chan()
        return state

    def open(self):
        # should construct or attach to the actual channels if not ready.
        if self.state == self.State.CLOSED:
            raise ConnectionError('reopening closed Connection not allowed')
        elif self.state == self.State.ATTACHED:
            self.state = self.State.READY
            return
        elif self.state == self.State.READY:
            return
        elif self.state == self.State.UNBACKED:
            if self.options.creation_policy == self.options.CreationPolicy.SENDER_CREATES:
                if self.outbound_initializer is not None:
                    target_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
                    descriptor = dgchan.create(target_muid, self.outbound_initializer)
                    self.outbound_chan = dch.Channel.attach(descriptor.sdesc)
                    self.outbound_initializer = None

                if self.inbound_initializer is not None:
                    chan_desc = dgchan.join(self.inbound_initializer)
                    self.inbound_chan = dch.Channel.attach(chan_desc.sdesc)
                    self.inbound_initializer = None
            elif self.options.creation_policy == self.options.CreationPolicy.RECEIVER_CREATES:
                if self.inbound_initializer is not None:
                    target_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
                    descriptor = dgchan.create(target_muid, self.inbound_initializer)
                    self.inbound_chan = dch.Channel.attach(descriptor.sdesc)
                    self.inbound_initializer = None

                if self.outbound_initializer is not None:
                    chan_desc = dgchan.join(self.outbound_initializer)
                    self.outbound_chan = dch.Channel.attach(chan_desc.sdesc)
                    self.outbound_initializer = None
            else:
                raise NotImplementedError('close case')

            self.state = self.State.READY

        elif self.state == self.State.UNATTACHED:
            if self.outbound_initializer is not None:
                self.outbound_chan = dch.Channel.attach(self.outbound_initializer)
                self.outbound_initializer = None

            if self.inbound_initializer is not None:
                self.inbound_chan = dch.Channel.attach(self.inbound_initializer)
                self.inbound_initializer = None

            self.state = self.State.READY

        else:
            raise NotImplementedError('open case')

    def _setup_write_adapter(self):
        if not self.outbound_chan.is_local:
            write_options = self.write_adapter_options
            if self.options.default_pool is None:
                try:
                    self._detach_pool = MemoryPool.attach(B64.str_to_bytes(dparm.this_process.default_pd))
                    write_options["buffer_pool"] = self._detach_pool
                except:
                    raise OSError("unable to attach to default Managed Memory pool")
            else:
                write_options["buffer_pool"] = self.options.default_pool
            self.write_adapter = dch.Peer2PeerWritingChannelFile(self.outbound_chan, options=write_options)
        else:
            self.write_adapter = dch.Peer2PeerWritingChannelFile(self.outbound_chan,
                                                                 options=self.write_adapter_options, wait_mode=self.policy.wait_mode)

    def _setup_read_adapter(self):
        self.read_adapter = dch.Peer2PeerReadingChannelFile(self.inbound_chan, wait_mode=self.policy.wait_mode)

    def _check_inbound(self):
        if self.inbound_chan is None and self.inbound_initializer is None:
            raise OSError(f'No receiving on a {self.__class__.__name__} that is not enabled for reading')

        if self.at_eof:
            raise EOFError()

        if self.state != self.State.READY:
            self.open()

    def _check_outbound(self):
        if self.outbound_chan is None and self.outbound_initializer is None:
            raise OSError(f'No sending on a {self.__class__.__name__} that is not enabled for writing')

        if self.state != self.State.READY:
            self.open()

    def send(self, obj):
        self._check_outbound()
        if self.write_adapter is None:
            self._setup_write_adapter()

        try:
            self.write_adapter.open()
            pickle.dump(obj, file=self.write_adapter, protocol=5)

        except Exception as e:
            raise ConnectionError(f"Could not complete send operation: {e}")
        finally:
            self.write_adapter.close()

    def recv(self):
        self._check_inbound()
        if self.read_adapter is None:
            self._setup_read_adapter()

        try:
            self.read_adapter.open()
            send_type, msg_len = self.read_adapter.check_header()

            if send_type == dch.ChannelAdapterMsgTypes.EOT:
                self.at_eof = True
                raise EOFError()
            elif send_type == dch.ChannelAdapterMsgTypes.RAW_BYTES:
                # multiprocessing.Connection unit tests assume you can go both ways.
                self.read_adapter.advance_raw_header(msg_len)
                buf = bytearray(msg_len)
                self.read_adapter.readinto(memoryview(buf))
                obj = pickle.loads(buf)
            else:
                obj = pickle.load(self.read_adapter)
        except EOFError:
            raise EOFError
        except Exception as e:
            raise ConnectionError(f"Could not complete receive operation: {e}")
        finally:
            self.read_adapter.close()

        return obj

    def ghost_close(self):
        """Force the connection into a closed state without sending EOT

        This 'ghosts' the receiver on the other end, which won't
        be getting an EOT message from this object, but when using
        externally managed Channels on this object that is what one
        might want to do, in order to transfer use of that Channel.
        """
        assert self.options.creation_policy == self.options.CreationPolicy.EXTERNALLY_MANAGED
        self.state = self.State.CLOSED
        if self.options.creation_policy == ConnectionOptions.CreationPolicy.PRE_CREATED:
            self._decref_inbound_chan()

    def thread_wait(self, timeout, done_ev, ready):
        """Thread waiter signaling with an ev.
        """
        if timeout is None:
            timeout = 1000000
        start = time.monotonic()
        delta = min(0.01, timeout) # TODO: figure out the value for delta
        timed_out = False

        if not done_ev.is_set():
            while not timed_out and not done_ev.is_set():
                done_res = self.poll(delta)
                if not done_ev.is_set():
                    if done_res:
                        ready.append(self)
                        done_ev.set()
                timed_out = (time.monotonic() - start) > timeout
            done_ev.set()  # if not set here, timed out.

    def close(self):
        if self.state == self.State.CLOSED:
            pass
        elif self.state == self.State.UNBACKED:
            pass  # nothing got created, so nothing to clean up.
        elif self.state == self.State.UNATTACHED:
            # todo: Should get the c_uid from the inbound_initializer (here a serialized descriptor) and destroy it.
            # todo: ask for a way to get this from the descriptor directly without attaching to anything
            pass
        elif self.state == self.State.READY or self.state == self.State.ATTACHED:
            # todo: look at cleanup policy here.  If we are holding around adapters
            # we may need to flush them.
            if self.state == self.State.READY and self.outbound_chan is not None and not self.ghost:
                try:
                    if self.write_adapter is None:
                        self._setup_write_adapter()
                    self.write_adapter.open()
                    self.write_adapter.write_eot()
                except dch.ChannelError:  # receiver may have destroyed channel
                    pass
                finally:
                    self.write_adapter.close()

                if self._detach_pool is not None:
                    self._detach_pool.detach()

            if self.options.creation_policy == ConnectionOptions.CreationPolicy.EXTERNALLY_MANAGED:
                # don't detach from or destroy anything
                pass
            elif self.options.creation_policy == ConnectionOptions.CreationPolicy.PRE_CREATED:
                # TODO: technically we should detach here, but we need refcounting in Channels to be merged first
                self._decref_inbound_chan()
            else:
                if self.outbound_chan is not None:
                    try:
                        self.outbound_chan.detach()
                    except dch.ChannelError:
                        self.outbound_chan = None

                if self.inbound_chan is not None:
                    inbound_cuid = self.inbound_chan.cuid

                    try:
                        self.inbound_chan.detach()
                    except dch.ChannelError:
                        pass

                    dgchan.destroy(inbound_cuid)
        else:
            raise NotImplementedError('close case')

        self.state = self.State.CLOSED
        return

    @property
    def inbound_channel(self):
        return self.inbound_chan

    @property
    def outbound_channel(self):
        return self.outbound_chan

    @property
    def closed(self):
        """True if the connection is closed"""
        return self.state == self.State.CLOSED

    @property
    def readable(self):
        """True if the connection is readable"""
        return self.inbound_initializer is not None or self.inbound_chan is not None

    @property
    def writable(self):
        """True if the connection is writable"""
        return self.outbound_initializer is not None or self.outbound_chan is not None

    def __enter__(self):
        self.open()
        return self

    def __exit__(self, exc_type, exc_inst, exc_tb):
        self.close()

    # Just a way to hang an attribute on an integer,
    # TODO: should be the real c_uid.
    def fileno(self):
        if self.state == self.State.READY:
            the_cuid = CUID(-2)
        else:
            the_cuid = CUID(-1)

        the_cuid.conn = self
        return the_cuid

    def poll(self, timeout=0.0, event_mask=dch.POLLIN):
        if self.at_eof:
            return True

        self._check_inbound()

        # multiprocessing tests interpret negative timeouts as 0.
        if isinstance(timeout, numbers.Number):
            if timeout < 0:
                timeout = 0

        return self.inbound_chan.poll(wait_mode=self.policy.wait_mode, event_mask=event_mask, timeout=timeout)

    def send_bytes(self, buffer, offset=0, size=None):
        self._check_outbound()

        # lifted from multiprocessing.connection.Connection
        m = memoryview(buffer)
        if m.itemsize > 1:
            m = memoryview(bytes(m))
        n = len(m)
        if offset < 0:
            raise ValueError("offset is negative")
        if n < offset:
            raise ValueError("buffer length < offset")
        if size is None:
            size = n - offset
        elif size < 0:
            raise ValueError("size is negative")
        elif offset + size > n:
            raise ValueError("buffer length < offset + size")
        # end lift

        try:
            if self.write_adapter is None:
                self._setup_write_adapter()
            self.write_adapter.open()
            self.write_adapter.write_raw_header(size)
            self.write_adapter.write(m[offset:offset + size])
        except Exception as e:
            raise ConnectionError(f"Could not complete send bytes operation: {e}")
        finally:
            self.write_adapter.close()

    # Todo: refactor recv_bytes and recv_bytes_into? duplicative
    def recv_bytes(self, maxlength=None):
        self._check_inbound()

        try:
            if self.read_adapter is None:
                self._setup_read_adapter()

            self.read_adapter.open()
            send_type, msg_len = self.read_adapter.check_header()

            if send_type == dch.ChannelAdapterMsgTypes.EOT:
                self.at_eof = True
                raise EOFError()
            elif send_type == dch.ChannelAdapterMsgTypes.PICKLE_PROT_5:
                # Would like to raise, but base case multiprocessing unit tests
                # abuse the interface, so we need to implement a way to get
                # the pickled rep out.  This is slow but people shouldn't
                # be doing it.
                # raise RuntimeError('recv_bytes called on a Connection whose sender called send')
                buf = bytearray()
                self.read_adapter.set_side_buf(buf)
                _ = pickle.load(self.read_adapter)
                if maxlength is not None and len(buf) > maxlength:
                    self.close()
                    raise OSError(f'recv_bytes (from send) maxlength={maxlength} but msg_len={msg_len}')

            else:
                if maxlength is not None and msg_len > maxlength:
                    self.close()  # base multiprocessing says connection is unusable after maxlength error; we could
                    # potentially stay around
                    raise OSError(f'recv_bytes maxlength={maxlength} but msg_len={msg_len}')

                self.read_adapter.advance_raw_header(msg_len)
                buf = bytearray(msg_len)
                self.read_adapter.readinto(memoryview(buf))
        except Exception as e:
            raise ConnectionError(f"Could not complete receive bytes operation: {e}")
        finally:
            self.read_adapter.close()

        return buf

    def recv_bytes_into(self, buffer, offset=0):
        self._check_inbound()

        with memoryview(buffer) as m:
            # Get byte size of arbitrary buffer
            byte_size = m.itemsize * len(m)
            if offset < 0:
                raise ValueError("negative offset")
            elif offset > byte_size:
                raise ValueError("offset too large")

            try:
                if self.read_adapter is None:
                    self._setup_read_adapter()
                self.read_adapter.open()
                send_type, msg_len = self.read_adapter.check_header()

                if send_type == dch.ChannelAdapterMsgTypes.EOT:
                    self.at_eof = True
                    raise EOFError()
                elif send_type == dch.ChannelAdapterMsgTypes.PICKLE_PROT_5:
                    raise RuntimeError('recv_bytes_into called on a Connection whose sender called send')

                if msg_len + offset > byte_size:
                    self.read_adapter.advance_raw_header(msg_len)
                    msg = self.read_adapter.read(msg_len)
                    raise multiprocessing.BufferTooShort(msg)

                self.read_adapter.advance_raw_header(msg_len)
                target = m.cast('B')
                self.read_adapter.readinto(target[offset:offset + msg_len])
            except Exception as e:
                raise ConnectionError(f"Could not complete recv_bytes_into operation: {e}")
            finally:
                self.read_adapter.close()

        return msg_len


_UNIQ_CHAN_LOCK = threading.Lock()
_CHAN_CTR = 0


def get_next_ctr():
    global _CHAN_CTR
    with _UNIQ_CHAN_LOCK:
        tmp = _CHAN_CTR
        _CHAN_CTR += 1

    return tmp


def make_uniq_chan_names():
    base_name = f'anon_{dparm.this_process.my_puid}_{get_next_ctr()}'
    return base_name, base_name + '_rev'


def Pipe(duplex=True, *, channels=None, options=None):
    """Pipe constructor, producing two Connections.

    The channels parameter can be None, indicating that
    this is an anonymous Pipe requiring a new channel to be constructed
    and eventually destroyed by the receiver.

    It can also be a Channel object or descriptor in which case
    the two Connections returned will communicate through the channel, with
    one being in read mode and the other in write mode.

    Finally, if it is a tuple, it is expected to have two elements which are
    both channels or channel descriptors and each of which can send and receive.

    If the channels parameter is None, then the duplex parameter governs
    whether the connections are full duplex and whether two new channels
    get created or only one.

    If not duplex, the first one is the reader and the second the writer

    The options object carries lifecycle and creation options for the Pipe call
    and also policy settings for the Connection objects that result, as well as
    policy on how the Connections objects should interact with the pool when
    sending data.

    :param channels: None or a Channel object or a 2-element tuple of distinct Channel objects
    :param duplex: default True, whether to produce
    :param options: a PipeOptions object
    :return: a pair of Connection objects.
    """

    if options is None:
        options = PipeOptions()

    if channels is not None:
        options.creation_policy = options.CreationPolicy.EXTERNAL
        options.conn_options.creation_policy = options.conn_options.CreationPolicy.EXTERNALLY_MANAGED

        if isinstance(channels, dch.Channel):
            first = Connection(inbound_initializer=channels, options=options.conn_options)
            second = Connection(outbound_initializer=channels, options=options.conn_options)
        elif 2 == len(channels):
            assert channels[0] is not channels[1]
            first = Connection(inbound_initializer=channels[0], outbound_initializer=channels[1])
            second = Connection(inbound_initializer=channels[1], outbound_initializer=channels[0])
        else:
            raise ValueError('unexpected channels/initializer: {channels}')
    else:
        if options.creation_policy == options.CreationPolicy.EARLY:
            options.conn_options.creation_policy = options.conn_options.CreationPolicy.PRE_CREATED
            target_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
            chan_options = ChannelOptions(ref_count=True)
            first_desc = dgchan.create(target_muid, options=chan_options).sdesc
            first_channel = dch.Channel.attach(first_desc)

            if duplex:
                second_desc = dgchan.create(target_muid, options=chan_options).sdesc
                second_channel = dch.Channel.attach(second_desc)
                first = Connection(inbound_initializer=first_channel,
                                   outbound_initializer=second_channel,
                                   options=options.conn_options)
                second = Connection(inbound_initializer=second_channel,
                                    outbound_initializer=first_channel,
                                    options=options.conn_options)
                # decref because all counting of the Channel happens in Connection now. Our create did one inc.
                dgchan.release_refcnt(first_channel.cuid)
                dgchan.release_refcnt(second_channel.cuid)
            else:
                first = Connection(inbound_initializer=first_channel, options=options.conn_options)
                second = Connection(outbound_initializer=first_channel, options=options.conn_options)
                # decref because all counting of the Channel happens in Connection now. Our create did one inc.
                dgchan.release_refcnt(first_channel.cuid)
        else:
            assert options.creation_policy != options.CreationPolicy.EXTERNAL

            firstname, secondname = make_uniq_chan_names()

            if duplex:
                first = Connection(inbound_initializer=firstname, outbound_initializer=secondname,
                                   options=options.conn_options)
                second = Connection(inbound_initializer=secondname, outbound_initializer=firstname,
                                    options=options.conn_options)
            else:
                first = Connection(inbound_initializer=firstname, options=options.conn_options)
                second = Connection(outbound_initializer=firstname, options=options.conn_options)

    return first, second


class Address:
    """Address class for dragon servers and listeners and all that
    """

    def __init__(self, user_name=''):
        self.user_name = user_name
        self.salt = dparm.this_process.my_puid

    def __str__(self):
        return self.user_name

    @property
    def cname(self):
        return f'{self.user_name!s}.{self.salt!s}'


class Listener:
    def __init__(self, address='', family='', backlog=1, authkey=None):
        pass


def Client(address='', family='', authkey=None):
    pass

src/dragon/infrastructure/connection.py
DragonHPC/dragon
import atexit
import os
import pdb
import sys
import traceback

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.parameters as dp
import dragon.utils as du
import dragon.infrastructure.connection as dconn

_DEBUG_OUT = None  # once initialized, attached dch.Channel
_DEBUG_IN = None
_LOCAL_BE = None
_LOCAL_BE_CONN = None

_BREAKPOINT_FH = None


def handle_breakpoint(msg):
    # print some stuff out and update the breakpoint file
    global _BREAKPOINT_FH
    if _BREAKPOINT_FH is None:
        _BREAKPOINT_FH = open(dfacts.BREAKPOINT_FILENAME, 'w')

    _BREAKPOINT_FH.write(msg.serialize() + '\n')
    _BREAKPOINT_FH.flush()
    sys.stdout.write(f'### breakpoint:  p_uid {msg.p_uid} on node {msg.index}\n')
    sys.stdout.flush()


def _cleanup_debug_channels():
    # TODO: figure out other cleanup to do here.  Should there be a switch to
    # try to destroy all the debug channels?  Some of this might better be under
    # the control of the reader, but the input channel can probably go
    out_adapter = DbgAdapter(_DEBUG_OUT)
    out_adapter.write(f'### debug exit for: {dp.this_process.my_puid} on {dp.this_process.index}\n\n')
    _DEBUG_OUT.detach()
    _LOCAL_BE.detach()
    _DEBUG_IN.destroy()


def _connect_debug_channels():
    global _DEBUG_OUT
    global _DEBUG_IN
    global _LOCAL_BE
    global _LOCAL_BE_CONN

    inf_pool_descr = du.B64.str_to_bytes(dp.this_process.inf_pd)
    inf_pool = dmm.MemoryPool.attach(inf_pool_descr)

    my_pid = os.getpid()

    dbg_out_cuid = dfacts.BASE_DEBUG_CUID * (2 * my_pid)
    dbg_in_cuid = dfacts.BASE_DEBUG_CUID * (2 * my_pid + 1)

    _DEBUG_OUT = dch.Channel(mem_pool=inf_pool, c_uid=dbg_out_cuid)
    _DEBUG_IN = dch.Channel(mem_pool=inf_pool, c_uid=dbg_in_cuid)

    be_ch_descr = du.B64.str_to_bytes(dp.this_process.local_be_cd)
    _LOCAL_BE = dch.Channel.attach(be_ch_descr)
    _LOCAL_BE_CONN = dconn.Connection(outbound_initializer=_LOCAL_BE)

    atexit.register(_cleanup_debug_channels)


class DbgAdapter:
    def __init__(self, chan):
        self.chan = chan
        self.sendh = self.chan.sendh()
        self.sendh.open()
        self.recvh = self.chan.recvh()
        self.recvh.open()

    def read(self):
        return self.recvh.recv_bytes().decode()

    def readline(self):
        return self.read()

    def write(self, data):
        self.sendh.send_bytes(data.encode())
        return len(data)

    def close(self):
        pass

    def flush(self):
        pass


# if true, write to the breakpoints file locally.
_TESTING_DEBUG_HOOK = False


def _get_bk_msg():
    if _DEBUG_OUT is None:
        _connect_debug_channels()

    bk_msg = dmsg.Breakpoint(tag=0, p_uid=dp.this_process.my_puid,
                             index=dp.this_process.index,
                             out_desc=du.B64.bytes_to_str(_DEBUG_OUT.serialize()),
                             in_desc=du.B64.bytes_to_str(_DEBUG_IN.serialize()))

    if _TESTING_DEBUG_HOOK:
        handle_breakpoint(bk_msg)

    return bk_msg


def dragon_debug_hook():
    bk_msg = _get_bk_msg()

    _LOCAL_BE_CONN.send(bk_msg.serialize())
    out_adapter = DbgAdapter(_DEBUG_OUT)
    out_adapter.write(f'### breakpoint: p_uid {dp.this_process.my_puid} on node {dp.this_process.index}')
    pdb.Pdb(stdin=DbgAdapter(_DEBUG_IN), stdout=out_adapter,
            skip=[f'{__name__}*']).set_trace()


def dragon_exception_hook(ex_type, ex_obj, ex_tb):
    bk_msg = _get_bk_msg()

    _LOCAL_BE_CONN.send(bk_msg.serialize())
    out_adapter = DbgAdapter(_DEBUG_OUT)
    out_adapter.write(f'### breakpoint: uncaught exception {ex_type!s} '
                      f'p_uid {dp.this_process.my_puid} on node '
                      f'{dp.this_process.index}')

    ex_lines = traceback.format_exception(ex_type, ex_obj, ex_tb)
    for line in ex_lines:
        out_adapter.write('### ' + line)

    pdb.Pdb(stdin=DbgAdapter(_DEBUG_IN), stdout=out_adapter).post_mortem(ex_tb)

src/dragon/infrastructure/debug_support.py
DragonHPC/dragon
"""Specified constants and names used in the Dragon runtime.
"""
import shlex
import enum
import sys
from .. import dtypes

PREFIX = 'DRAGON_'

# Number of gateway channels per node
DRAGON_DEFAULT_NUM_GW_CHANNELS_PER_NODE = 0
DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE = 1
NUM_GW_CHANNELS_PER_NODE_VAR = 'NUM_GW_CHANNELS_PER_NODE'
# needed for naming convention for the gateway channels
# should be unique on a given node, not globally
# Note: there is also a dependency on dragon_channel_register_gateways_from_env()
GW_ENV_PREFIX = PREFIX + 'GW'

# For environment variable passing, this set is the list of dragon parameters
# in capital letters.
env_vars = frozenset({'MODE', 'INDEX', 'DEFAULT_PD', 'INF_PD', 'LOCAL_SHEP_CD', 'LOCAL_BE_CD', 'GS_RET_CD',
                 'SHEP_RET_CD', 'GS_CD', 'DEFAULT_SEG_SZ', 'INF_SEG_SZ', 'TEST', 'DEBUG', 'MY_PUID',
                 'BE_CUID', 'INF_WAIT_MODE', 'USER_WAIT_MODE', 'USER_RETURN_WHEN_MODE', 'INF_RETURN_WHEN_MODE',
                 'GW_CAPACITY', NUM_GW_CHANNELS_PER_NODE_VAR, 'TRANSPORT_AGENT', 'HSTA_MAX_EJECTION_MB', 'HSTA_MAX_GETMSG_MB',
                 'PMOD_COMMUNICATION_TIMEOUT', 'BASEPOOL'})


# TODO:PE-37770  This list of names should NOT need to appear
# both here and in the parameters module.

# Get environment variable name from parameter name. Conversion routine used by the parameters module.
def env_name(parameter_name: str) -> str:
    if parameter_name not in env_vars:
        raise ValueError(f'{parameter_name} is not a known parameter')

    return PREFIX + parameter_name


# This sets the environment variable names as attributes of this module. So for instance, ALLOC_SZ is a
# variable name of this module with value 'DRAGON_ALLOC_SZ'. These are the environment variable names only
# not their values, which can be accessed through the parameters module.
for var_name in env_vars:
    setattr(sys.modules[__name__], var_name, var_name)
    setattr(sys.modules[__name__], "ENV_" + var_name, env_name(var_name))

# Mode Values
CRAYHPE_MULTI_MODE = 'crayhpe-multi'
SINGLE_MODE = 'single'
TEST_MODE = 'test'

# Pool Values
PATCHED = 'PATCHED'
NATIVE = 'NATIVE'
DEFAULT_POOL = NATIVE

#: Input channel unique ID for Global Services head
GS_INPUT_CUID = 2

#: Unique process ID for head user process started by the launcher
LAUNCHER_PUID = 0
#: Unique process ID given for the global services process
GS_PUID = 1

#: The first PUID for user applications launched by local services
FIRST_PUID = 2 ** 32

#: Base index used by the transport services for the unique process ID (puid)
FIRST_TRANSPORT_PUID = 2 ** 8
#: Range of the unique process ID's for the transport services
RANGE_TRANSPORT_PUID = FIRST_PUID - FIRST_TRANSPORT_PUID

#: Starting value for PMOD launch Channel ID (cuid) that
#: communicates with the client pmod at process launch
BASE_PMOD_LAUNCH_CUID = 2 ** 54

#: Range for the PMOD launch channel ID's (cuid) that
#: communicates with the client pmod at process launch
RANGE_PMOD_LAUNCH_CUID = 2 ** 54

# needed by the frontend for its infrastructure channels
FE_CUID = 2 ** 55 - 1
FE_GW_CUID = 2 ** 55 - 2
FE_LOCAL_IN_CUID = 2 ** 55 - 3
FE_LOCAL_OUT_CUID = 2 ** 55 - 4

#: Starting value of a reserved range of cuids that can be managed by a user
BASE_USER_MANAGED_CUID = 2 ** 54
#: Range for the user managed cuids
RANGE_USER_MANAGED_CUID = 2 ** 54

#: Starting value of the local transport agent channel unique ID (cuid)
BASE_TA_CUID = 2 ** 55
#: Range for the local transport agent channel ID's (cuid)
RANGE_TA_CUID = 2 ** 55

#: Starting value of the local services shepherd channel unique ID (cuid)
BASE_SHEP_CUID = 2 ** 56
#: Range for the local services channel ID's (cuid)
RANGE_SHEP_CUID = 2 ** 56

#: Starting value for the local launcher backend channel unique ID (cuid)
BASE_BE_CUID = 2 ** 57
#: Range for the local launcher backend channel ID's (cuid)
RANGE_BE_CUID = 2 ** 57

#: Starting value for the gateway channel unique ID (cuid)
BASE_GW_CUID = 2 ** 58
#: Range for the gateway channel ID's (cuid)
RANGE_GW_CUID = 2 ** 58

#: Base value for the logging channel ID (cuid)
BASE_LOG_CUID = 2 ** 59
#: Range for the logging channel ID's (cuid)
RANGE_LOG_CUID = 2 ** 59

#: Starting value for the backend channel ID (cuid)
#: of the infrastructure transport gateway channel
BASE_BE_GW_CUID = 2 ** 60
#: Range for the backend channel ID's (cuid)
#: of the infrastructure transport gateway channel
RANGE_BE_GW_CUID = 2 ** 60

#: Starting value for the backend channel ID (cuid) that
#: communicates with its transport agent
BASE_BE_LOCAL_CUID = 2 ** 61
#: Range for the backend channel ID's (cuid) that
#: communicate with its transport agent
RANGE_BE_LOCAL_CUID = 2 ** 61

#: Starting value for the backend channel ID (cuid) that
#: communicates with the frontend
BASE_BE_FE_CUID = 2 ** 62
#: Range for the backend channel ID's (cuid) that
#: communicate with the frontend
RANGE_BE_FE_CUID = 2 ** 62

#: Starting value of the User created channel ID (cuid)
FIRST_CUID = 2 ** 63

#: First backend pool ID (MUID) needed for the overlay network transport
FIRST_BE_MUID = 2 ** 60
#: Range of the backend pool ID's (MUID) needed for the overlay network transport
RANGE_BE_MUID = 2 ** 60

#: Pool for the frontend overlay network transport agent
FE_OVERLAY_TRANSPORT_AGENT_MUID = 2 ** 60 - 1

#: First Logging MUID
FIRST_LOGGING_MUID = 2 ** 61

#: First MUID of Default Pools.
FIRST_DEFAULT_MUID = 2 ** 62

#: User Pool MUID starting value.
FIRST_MUID = 2 ** 63

#: Starting value for the Group of resources (guid)
FIRST_GUID = 2 ** 63

#: Frontend's magic host ID
FRONTEND_HOSTID = 1

# Names
DEFAULT_PROCESS_NAME_BASE = 'dragon_process_'
DEFAULT_POOL_NAME_BASE = 'dragon_pool_'
DEFAULT_DICT_POOL_NAME_BASE = 'dragon_dict_pool_'
DEFAULT_CHANNEL_NAME_BASE = 'dragon_channel_'
DEFAULT_GROUP_NAME_BASE = 'dragon_group_'

GS_ERROR_EXIT = -1

# Reserved Shared Memory Segments
INFRASTRUCTURE_POOL_SUFFIX = "_inf"
DEFAULT_POOL_SUFFIX = "_def"
LOGGING_POOL_SUFFIX = "_log"


# Pre-defined process PUIDs:
def is_transport_puid(puid) -> bool:
    """Check whether a puid is associated with a transport agent

    :param puid: puid of a process
    :type puid: int
    :return: Whether the puid is associated with pre-defined puid value
    :type: bool
    """

    if puid >= FIRST_TRANSPORT_PUID and puid <= (FIRST_TRANSPORT_PUID + RANGE_TRANSPORT_PUID):
        return True
    else:
        return False


def transport_puid_from_index(index) -> int:
    """Get the unique process ID from the node index.
    Otherwise, throw an `AssertionError`

    :param index: node index
    :type index: int
    :return: p_uid of the local transport
    :rtype: int
    """
    assert (index >= 0)
    assert (index < RANGE_TRANSPORT_PUID)
    return FIRST_TRANSPORT_PUID + index


# Pre-defined Channel CUIDs
def launcher_cuid_from_index(index:int) -> int:
    """Get the unique Channel ID of the local Launcher backend from the node
    index.

    :param index: local node index.
    :type index: int
    :return: c_uid of the Launcher backend
    :rtype: int
    """
    assert (index >= 0)
    assert (index < RANGE_BE_CUID)
    return BASE_BE_CUID + index


def transport_cuid_from_index(index:int)->int:
    """Get the unique Channel ID of the local transport agent from the nodex index.

    :param index: local node index
    :type index: int
    :return: c_uid of the local transport agent
    :rtype: int
    """

    assert (index >= 0)
    assert (index < RANGE_TA_CUID)
    return BASE_TA_CUID + index


def shepherd_cuid_from_index(index:int) -> int:
    """Get the unique Channel ID of local services from the node index.

    :param index: Local node index
    :type index: int
    :return: c_uid of the Local Services
    :rtype: int
    """

    assert (index >= 0)
    assert (index < RANGE_SHEP_CUID)
    return BASE_SHEP_CUID + index


def gw_cuid_from_index(index:int, num_gw_channels:int) -> int:
    """Return the unique Channel ID of the first gateway channel from the local
    node index

    :param index: local node index
    :type index: int
    :param num_gw_channels: number of gateway channels
    :type num_gw_channels: int
    :return: c_uid of first gateway channel
    :rtype: int
    """
    assert (index >= 0)
    assert (index * num_gw_channels < RANGE_GW_CUID)
    return BASE_GW_CUID + index * num_gw_channels


def be_gw_cuid_from_hostid(host_id):
    """Return the unique backend channel ID of the overlay network
    transport gateway channel, based on the host id

    :param host_id: host id
    :type host_id: int
    :return: c_uid of backend transport gateway channel
    :rtype: int
    """
    assert (host_id >= 0)
    cand_cuid = BASE_BE_GW_CUID + host_id
    cand_cuid = cand_cuid % (BASE_BE_GW_CUID + RANGE_BE_GW_CUID)
    assert (cand_cuid < BASE_BE_GW_CUID + RANGE_BE_GW_CUID)
    return cand_cuid


def be_fe_cuid_from_hostid(host_id):
    """Return the unique backend Channel ID that communicates
    with the fronted, based on the host id

    :param host_id: host id
    :type host_id: int
    :return: c_uid of backend channel that communicates with fronted
    :rtype: int
    """
    assert (host_id >= 0)
    cand_cuid = BASE_BE_FE_CUID + host_id
    cand_cuid = cand_cuid % (BASE_BE_FE_CUID + RANGE_BE_FE_CUID)
    assert (cand_cuid < BASE_BE_FE_CUID + RANGE_BE_FE_CUID)
    return cand_cuid


def be_local_cuid_in_from_hostid(host_id):
    """Return the unique channel ID based on host id, for a backend channel that acts
    as an inbound initializer for its communication with transport agent

    :param host_id: host id
    :type host_id: int
    :return: c_uid of backend inbound channel that communicates with transport agent
    :rtype: int
    """
    assert (host_id >= 0)
    cand_cuid = BASE_BE_LOCAL_CUID + host_id
    cand_cuid = cand_cuid % (BASE_BE_LOCAL_CUID + RANGE_BE_LOCAL_CUID)
    assert (cand_cuid < BASE_BE_LOCAL_CUID + RANGE_BE_LOCAL_CUID)
    return cand_cuid


def be_local_cuid_out_from_hostid(host_id):
    """Return the unique channel ID based on host id, for a backend channel that acts
    as an outbound initializer for its communication with transport agent

    :param host_id: host id
    :type host_id: int
    :return: c_uid of backend outbound channel that communicates with transport agent
    :rtype: int
    """
    assert (host_id >= 0)
    cand_cuid = BASE_BE_LOCAL_CUID + host_id + 1
    cand_cuid = cand_cuid % (BASE_BE_LOCAL_CUID + RANGE_BE_LOCAL_CUID)
    assert (cand_cuid < BASE_BE_LOCAL_CUID + RANGE_BE_LOCAL_CUID)
    return cand_cuid

def pmod_launch_cuid_from_jobinfo(host_id, job_id, lrank):
    """Return the unique channel ID based on host id, for a channel that initializes
    a client PMI process with its required pmod data.

    :param host_id: host id
    :type host_id: int
    :return: c_uid of backend outbound channel that communicates with transport agent
    :rtype: int
    """
    assert (host_id >= 0)
    job_spacing = 2 ** 20
    host_spacing = 2 ** 14
    pmod_launch_cuid = BASE_PMOD_LAUNCH_CUID + (job_id * job_spacing) + (host_id * host_spacing) + lrank + 1
    pmod_launch_cuid = pmod_launch_cuid % (BASE_PMOD_LAUNCH_CUID + RANGE_PMOD_LAUNCH_CUID)
    assert (pmod_launch_cuid < BASE_PMOD_LAUNCH_CUID + RANGE_PMOD_LAUNCH_CUID)
    return pmod_launch_cuid



# Pre-defined Pool MUIDs
def default_pool_muid_from_index(index:int) -> int:
    """Return the unique managed memory pool ID of the default user pool from
    the node index.

    :param index: index of this node
    :type index: int
    :return: m_uid of user pool
    :rtype: int
    """
    assert (index >= 0)
    return FIRST_DEFAULT_MUID + index


def logging_pool_muid_from_index(index:int) -> int:
    """Return the unique managed memory pool ID of the default logging pool from
    the node index.

    :param index: index of this node
    :type index: int
    :return: m_uid of logging pool
    :rtype: int
    """
    assert (index >= 0)
    return FIRST_LOGGING_MUID + index


def infrastructure_pool_muid_from_index(index:int) -> int:
    """Return the unique managed memory pool ID of the default infrastructure pool from
    the node index.

    :param index: index of this node
    :type index: int
    :return: m_uid of the infrastructure pool
    :rtype: int
    """
    assert (index >= 0)
    return index


def index_from_infrastructure_pool_muid(m_uid:int) -> int:
    """Return the node index from the unique managed memory pool ID of the
    default infrastructure pool.

    :param m_uid: m_uid of the infrastructure pool
    :type index: int
    :return: unique node index of this node
    :rtype: int
    """
    assert (m_uid < FIRST_DEFAULT_MUID)
    return m_uid


def be_pool_muid_from_hostid(host_id):
    """Return the backend memory pool unique ID based on the host index

    :param host_id: host id
    :type index: int
    :return: unique ID of the backend memory pool
    :rtype: int
    """
    assert (host_id >= 0)
    cand_muid = FIRST_BE_MUID + host_id
    cand_muid = cand_muid % (FIRST_BE_MUID + RANGE_BE_MUID)
    assert (cand_muid < FIRST_BE_MUID + RANGE_BE_MUID)
    return cand_muid


def is_pre_defined_pool(m_uid:int) -> bool:
    """Return if the pool corresponding to the given `m_uid` is pre-defined by
    the runtime.

    :param m_uid: uid to check
    :type m_uid: int
    :return: True if the m_uid is predefined, otherwise throw an `AssertionError`.
    :rtype: bool
    """

    assert m_uid >= 0
    return m_uid < FIRST_MUID


def index_from_shepherd_cuid(cuid:int)->int:
    """Retun the local node index from the `c_uid` of local services.

    :param cuid: unique channel id of the local Shepherd
    :type cuid: int
    :return: Unique local node index.
    :rtype: int
    """
    assert (cuid >= BASE_SHEP_CUID)
    assert (cuid < BASE_SHEP_CUID + RANGE_SHEP_CUID)
    return cuid - BASE_SHEP_CUID


def is_default_pool(m_uid:int) -> bool:
    """Check if the given `m_uid` is a default `m_uid`. Otherwise throw an `AssertionError`.

    :param m_uid: uid to check
    :type m_uid: int
    :return: True
    :rtype: bool
    """
    assert m_uid >= FIRST_DEFAULT_MUID
    assert m_uid < FIRST_MUID
    return True


def index_from_default_pool_muid(m_uid:int) -> int:
    """Return the unique node index from the default pool id.

    :param m_uid: A default pool id.
    :type m_uid: int
    :return: node index
    :rtype: int
    """
    assert m_uid >= FIRST_DEFAULT_MUID
    assert is_default_pool(m_uid)
    return m_uid - FIRST_DEFAULT_MUID


# Size constants
ONE_MB = 2 ** 20
ONE_GB = 2 ** 30
TWO_GB = 2 ** 31
FOUR_GB = 2 ** 32
SIXTEEN_GB = 2 ** 34

# decision constant for argument delivery
# TODO: tune this - 32K is probably too big. check size of e.g. pool worker Process call.
ARG_IMMEDIATE_LIMIT = 2 ** 15  # 32K

# environment variable name to use to pass in
# some pickled pre-made test channels to gs
# for testing purposes.
GS_TEST_CH_EV = '_XDRAGON_GS_TEST_CHANS'
GS_LOG_BASE = '_XDRAGON_GS_LOG_BASE'

GS_SINGLE_LAUNCH_CMD = 'import dragon.globalservices.server as dgs; dgs.single()'
GS_MULTI_LAUNCH_CMD = 'import dragon.globalservices.server as dgs; dgs.multi()'

# add names for more entry points for e.g. multinode
DRAGON_LOGGER_SDESC = "DRAGON_LOGGER_SDESC"
DEFAULT_GS_SINGLE_MAX_MSGS = str(1000)
DEFAULT_SH_SINGLE_MAX_MSGS = str(1000)
DEFAULT_SINGLE_DEF_SEG_SZ = str(FOUR_GB)
DEFAULT_SINGLE_INF_SEG_SZ = str(ONE_GB)
DEFAULT_BE_OVERLAY_TRANSPORT_SEG_SZ = str(2 ** 8 * ONE_MB) # 256MB

DRAGON_CFG_FILENAME = '.dragon_launch_config'

BREAKPOINT_FILENAME = '.dragon_breakpoints'
BASE_DEBUG_CUID = 2 ** 33  # range of pid is 32 bits, need 2 per, so < 2**34

# time to wait for messages in ls teardown before declaring an error
TEARDOWN_PATIENCE = 1.0  # seconds, 1 second

# This may become configurable and part of the managed process creation in the future.
MANAGED_PROCESS_MAX_OUTPUT_SIZE_PER_MESSAGE = 5000

# needed for naming the environment variable keys for the
# three channels when piping stdout/stderr/stdin for a process
# created by the Dragon run-time services.
STDOUT_DESC = 'dragon_stdout_desc'
STDERR_DESC = 'dragon_stderr_desc'
STDIN_DESC = 'dragon_stdin_desc'

# default capacity of gateway channels. This might need to be raised
# in circumstances where we're flooded in ops in a sequence that lead to a hang
GW_DEFAULT_CAPACITY = 2048

# This defines an enivironment variable name for an isolated transport test
# environment where global services is not started.
TRANSPORT_TEST_ENV = 'DRAGON_TRANSPORT_TEST'

# This provides env. var. to perform a simple bring-up and teardown of all
# multi-node services in order to measure bounce time of our services.
BOUNCE_TEST_ENV = 'DRAGON_BOUNCE_RUNTIME_TEST'

# This defines the preferred waiting mode of connections, either idle waiting
# or spin waiting for both user connections (those made directly by user code)
# and infrastructure connections (made by dragon infrastructure code).
INFRASTRUCTURE_DEFAULT_WAIT_MODE = dtypes.IDLE_WAIT
USER_DEFAULT_WAIT_MODE = dtypes.ADAPTIVE_WAIT

# This defines the preferred send return when mode for connections.
# WHEN_DEPOSITED means to return control to the sender when the message
# is deposited into the destination channel. WHEN_IMMEDIATE means to
# return control to the sender earlier, allowing more asynchronous behavior
# but not guaranteeing order between other operations on the channel like
# receives and polling operations.
INFRASTRUCTURE_DEFAULT_RETURN_WHEN_MODE = dtypes.WHEN_BUFFERED
USER_DEFAULT_RETURN_WHEN_MODE = dtypes.WHEN_DEPOSITED

# Service process names
from ..cli import (PROCNAME_GS,
                   PROCNAME_LA_FE,
                   PROCNAME_LA_BE,
                   PROCNAME_LS,
                   PROCNAME_TCP_TA,
                   PROCNAME_OVERLAY_TA,
                   console_script_args)

# Aliases for transport agent commands. Transport agent commands are always
# called with two positional arguments:
#   1. Node index
#   2. Input channel serialized descriptor
@enum.unique
class TransportAgentOptions(enum.Enum):
    """Enumerated list of supported transport agents"""
    HSTA = 'hsta'
    TCP = 'tcp'

    def __str__(self):
        return self.value

    @staticmethod
    def from_str(s):
        """Obtain enum value from TransportAgentOptions string

        :param s: string representation of enumerated TransportAgentOptions
        :type s: str
        :return: name of an available transport agent
        :rtype: TransportAgentOptions
        """
        try:
            return TransportAgentOptions(s)
        except KeyError:
            raise ValueError()


TRANSPORT_AGENT_ALIASES = {
    # NOTE Cannot use dragon.cli.console_script_args() as it would cause a
    # NOTE circular import.
    str(TransportAgentOptions.TCP): console_script_args(PROCNAME_TCP_TA, '--no-tls'),
    # TODO Launcher does not yet support dynamic cert/key generation
    # TODO See dragon.transport.x509 to generate keys and certs. Possible
    # TODO approaches for integrating X509 creds:
    # TODO
    # TODO   1. Add new messages to support CSRs and Certs and instrument
    # TODO      agents, local services, launcher BE and FE to support X509
    # TODO      lifecycle.
    # TODO
    # TODO   2. Launcher/Local services detects TLS-enabled agent and
    # TODO      generates keys and certs and passes them as command line
    # TODO      arguments (as shown below).
    # TODO
    # TODO Either way we should probably consider the launcher to generate
    # TODO a root CA which is used to issue a sub-CA to each local
    # TODO services. That would enable local services to more easily manage
    # TODO the X509 lifecycle with (option 1) or for (option 2) transport
    # TODO agents.
    #'tcp+tls': console_script_args(PROCNAME_TCP_TA, '--cafile', CA, '--certfile', CERT, '--keyfile', KEY)),
}

DEFAULT_TRANSPORT_AGENT = shlex.split(str(TransportAgentOptions.TCP))

# Prefer to use high-speed network interfaces. Be cautious about matching any
# interface (e.g., r'^.*$') as it may expose the transport agent on an
# unintended network (e.g., the internet).
DEFAULT_TRANSPORT_NETIF = r'^(hsn|ipogif|ib)\d+$'
DEFAULT_TRANSPORT_PORT = 7575
DEFAULT_OVERLAY_NETWORK_PORT = 6565
DEFAULT_FRONTEND_PORT = 6566
DEFAULT_PMI_CONTROL_PORT = 8575
DEFAULT_PORT_RANGE=1000

# GS_DEFAULT_POLICY -- To prevent circular imports, this lives in policy_eval.py

src/dragon/infrastructure/facts.py
DragonHPC/dragon
"""A class that uniquely defines properties of a managed Group of resources for infrastructure communication."""

import enum
from dataclasses import dataclass, field
from ..infrastructure import parameters as dparm
from ..infrastructure.process_desc import ProcessDescriptor as process_desc
from ..infrastructure.channel_desc import ChannelDescriptor as channel_desc
from ..infrastructure.pool_desc import PoolDescriptor as pool_desc
from ..infrastructure.node_desc import NodeDescriptor as node_desc
from ..globalservices.policy_eval import Policy


@dataclass
class GroupDescriptor:
    """The group descriptor holds a list of tuples of puids, muids, cuids or huids."""

    @enum.unique
    class State(enum.IntEnum):
        """Current state of a group's member."""

        PENDING = 0
        ACTIVE = 1
        DEAD = 2

    @dataclass
    class GroupMember:
        """The GroupMember class holds information about each object that is a member of the group."""

        @enum.unique
        class Errors(enum.Enum):
            SUCCESS = 0  #: Resource was created
            FAIL = 1  #: Resource was not created
            ALREADY = 2  #: Resource exists already

        state : object = None
        uid : int = None
        placement : int = None
        desc : object = None
        error_code : int = None
        error_info : str = None

        @classmethod
        def from_sdict(cls, d):
            """Returns the launched group member object.

            :param d: dictionary with all key-value pairs of the available information of the group's member
            :type d: Dictionary
            :raises ValueError: when there is an error during deserialization
            :return: the GroupMember object
            :rtype: GroupMember
            """
            try:
                return cls(**d)
            except Exception as exc:
                raise ValueError(f'Error deserializing {cls.__name__} {d=}') from exc

        def get_sdict(self):
            """Collects the entire information of a member of the group.

            :return: A dictionary with all key-value pairs of the available information of the group's member.
            :rtype: Dictionary
            """
            rv = {'state': self.state.value,
                  'uid': self.uid,
                  'placement': self.placement}

            rv['desc'] = self.desc.get_sdict()
            rv['error_code'] = self.error_code.value
            rv['error_info'] = self.error_info

            return rv

    state : State = State.PENDING
    g_uid : int = None
    name : str = None
    sets : list(list()) = field(default_factory=list)
    policy : Policy = None

    def __post_init__(self):
        if type(self.policy) is dict:
            self.policy = Policy.from_sdict(self.policy)

        if self.sets:
            old_sets, self.sets = self.sets, []
            for i, lst in enumerate(old_sets):
                self.sets.append([])
                for j, member in enumerate(lst):
                    self.sets[i] += [member if type(member) is GroupDescriptor.GroupMember else GroupDescriptor.GroupMember.from_sdict(member)]
                    if type(self.sets[i][j].desc) is dict:
                        # self.sets[i][j].desc is a dictionary and we need to call from_sdict() on it
                        # to convert it to the appropriate Descriptor object
                        # figure out what type of resource object this member is
                        if 'c_uid' in self.sets[i][j].desc.keys():
                            self.sets[i][j].desc = channel_desc.from_sdict(self.sets[i][j].desc)
                        elif 'm_uid' in self.sets[i][j].desc.keys():
                            self.sets[i][j].desc = pool_desc.from_sdict(self.sets[i][j].desc)
                        elif 'h_uid' in self.sets[i][j].desc.keys():
                            self.sets[i][j].desc = node_desc.from_sdict(self.sets[i][j].desc)
                        else:
                            self.sets[i][j].desc = process_desc.from_sdict(self.sets[i][j].desc)

    @classmethod
    def from_sdict(cls, d):
        """Returns the descriptor of the launched group of resources.

        :param d: dictionary with all key-value pairs of the available information of the group
        :type d: Dictionary
        :raises ValueError: when there is an error during deserialization
        :return: the GroupDescriptor object
        :rtype: GroupDescriptor
        """

        try:
            return cls(**d)
        except Exception as exc:
            raise ValueError(f'Error deserializing {cls.__name__} {d=}') from exc

    def get_sdict(self):
        """Collects the entire information of the group.

        :return: A dictionary with all key-value pairs of the available information of the group.
        :rtype: Dictionary
        """
        rv = {'name': self.name,
              'g_uid': self.g_uid,
              'state': self.state}

        if isinstance(self.policy, Policy):
            rv['policy'] = self.policy.get_sdict()
        else:
            rv['policy'] = self.policy

        if self.sets is not None:
            rv['sets'] = []
            for i, tup in enumerate(self.sets):
                rv['sets'].append(())
                for member in tup:
                    rv['sets'][i] += (member.get_sdict(),)

        return rv

src/dragon/infrastructure/group_desc.py
DragonHPC/dragon
"""Dragon infrastructure messages are the internal API used for service communication.
"""

import enum
import json
import zlib
import base64
import subprocess

from dataclasses import dataclass, asdict

from ..infrastructure import channel_desc
from ..infrastructure import pool_desc
from ..infrastructure import node_desc
from ..infrastructure import process_desc
from ..infrastructure import parameters as parms
from ..infrastructure import facts as dfacts
from ..localservices import options as dso
from ..infrastructure.util import to_str_iter
from ..utils import B64

from ..globalservices.policy_eval import ResourceLayout, Policy

from ..infrastructure import group_desc
# from ..infrastructure.policy import DefaultPolicy


# This enum class lists the type codes in infrastructure
# messages.  The values are significant for interoperability.


@enum.unique
class MessageTypes(enum.Enum):
    """
    These are the enumerated values of message type identifiers within
    the Dragon infrastructure messages.
    """
    INVALID = 0  #:
    GS_PROCESS_CREATE = 1  #:
    GS_PROCESS_CREATE_RESPONSE = 2  #:
    GS_PROCESS_LIST = 3  #:
    GS_PROCESS_LIST_RESPONSE = 4  #:
    GS_PROCESS_QUERY = 5  #:
    GS_PROCESS_QUERY_RESPONSE = 6  #:
    GS_PROCESS_KILL = 7  #:
    GS_PROCESS_KILL_RESPONSE = 8  #:
    GS_PROCESS_JOIN = 9  #:
    GS_PROCESS_JOIN_RESPONSE = 10  #:
    GS_CHANNEL_CREATE = 11  #:
    GS_CHANNEL_CREATE_RESPONSE = 12  #:
    GS_CHANNEL_LIST = 13  #:
    GS_CHANNEL_LIST_RESPONSE = 14  #:
    GS_CHANNEL_QUERY = 15  #:
    GS_CHANNEL_QUERY_RESPONSE = 16  #:
    GS_CHANNEL_DESTROY = 17  #:
    GS_CHANNEL_DESTROY_RESPONSE = 18  #:
    GS_CHANNEL_JOIN = 19  #:
    GS_CHANNEL_JOIN_RESPONSE = 20  #:
    GS_CHANNEL_DETACH = 21  #:
    GS_CHANNEL_DETACH_RESPONSE = 22  #:
    GS_CHANNEL_GET_SENDH = 23  #:
    GS_CHANNEL_GET_SENDH_RESPONSE = 24  #:
    GS_CHANNEL_GET_RECVH = 25  #:
    GS_CHANNEL_GET_RECVH_RESPONSE = 26  #:
    ABNORMAL_TERMINATION = 27  #:
    GS_STARTED = 28  #:
    GS_PING_SH = 29  #:
    GS_IS_UP = 30  #:
    GS_HEAD_EXIT = 31  #:
    GS_CHANNEL_RELEASE = 32  #:
    GS_HALTED = 33  #:
    SH_PROCESS_CREATE = 34  #:
    SH_PROCESS_CREATE_RESPONSE = 35  #:
    SH_PROCESS_KILL = 36  #:
    SH_PROCESS_EXIT = 37  #:
    SH_CHANNEL_CREATE = 38  #:
    SH_CHANNEL_CREATE_RESPONSE = 39  #:
    SH_CHANNEL_DESTROY = 40  #:
    SH_CHANNEL_DESTROY_RESPONSE = 41  #:
    SH_LOCK_CHANNEL = 42  #:
    SH_LOCK_CHANNEL_RESPONSE = 43  #:
    SH_ALLOC_MSG = 44  #:
    SH_ALLOC_MSG_RESPONSE = 45  #:
    SH_ALLOC_BLOCK = 46  #:
    SH_ALLOC_BLOCK_RESPONSE = 47  #:
    SH_IS_UP = 48  #:
    SH_CHANNELS_UP = 49  #:
    SH_PING_GS = 50  #:
    SH_HALTED = 51  #:
    SH_FWD_INPUT = 52  #:
    SH_FWD_INPUT_ERR = 53  #:
    SH_FWD_OUTPUT = 54  #:
    GS_TEARDOWN = 55  #:
    SH_TEARDOWN = 56  #:
    SH_PING_BE = 57  #:
    BE_PING_SH = 58  #:
    TA_PING_SH = 59  #:
    SH_HALT_TA = 60  #:
    TA_HALTED = 61  #:
    SH_HALT_BE = 62  #:
    BE_HALTED = 63  #:
    TA_UP = 64  #:
    GS_PING_PROC = 65  #:
    GS_DUMP = 66  #:
    SH_DUMP = 67  #:
    LA_BROADCAST = 68  #:
    LA_PASSTHRU_FB = 69  #:
    LA_PASSTHRU_BF = 70  #:
    GS_POOL_CREATE = 71  #:
    GS_POOL_CREATE_RESPONSE = 72  #:
    GS_POOL_DESTROY = 73  #:
    GS_POOL_DESTROY_RESPONSE = 74  #:
    GS_POOL_LIST = 75  #:
    GS_POOL_LIST_RESPONSE = 76  #:
    GS_POOL_QUERY = 77  #:
    GS_POOL_QUERY_RESPONSE = 78  #:
    SH_POOL_CREATE = 79  #:
    SH_POOL_CREATE_RESPONSE = 80  #:
    SH_POOL_DESTROY = 81  #:
    SH_POOL_DESTROY_RESPONSE = 82  #:
    SH_EXEC_MEM_REQUEST = 83  #:
    SH_EXEC_MEM_RESPONSE = 84  #:
    GS_UNEXPECTED = 85  #:
    LA_SERVER_MODE = 86  #:
    LA_SERVER_MODE_EXIT = 87  #:
    LA_PROCESS_DICT = 88  #:
    LA_PROCESS_DICT_RESPONSE = 89  #:
    LA_DUMP = 90  #:
    BE_NODEINDEX_SH = 91  #:
    LA_CHANNELS_INFO = 92  #:
    SH_PROCESS_KILL_RESPONSE = 93  #:
    BREAKPOINT = 94  #:
    GS_PROCESS_JOIN_LIST = 95  #:
    GS_PROCESS_JOIN_LIST_RESPONSE = 96  #:
    GS_NODE_QUERY = 97  #:
    GS_NODE_QUERY_RESPONSE = 98  #:
    LOGGING_MSG = 99  #:
    LOGGING_MSG_LIST = 100  #:
    LOG_FLUSHED = 101  #:
    GS_NODE_LIST = 102  #:
    GS_NODE_LIST_RESPONSE = 103  #:
    GS_NODE_QUERY_TOTAL_CPU_COUNT = 104  #:
    GS_NODE_QUERY_TOTAL_CPU_COUNT_RESPONSE = 105  #:
    BE_IS_UP = 106  #:
    FE_NODE_IDX_BE = 107  #:
    HALT_OVERLAY = 108  #:
    HALT_LOGGING_INFRA = 109
    OVERLAY_PING_BE = 110  #:
    OVERLAY_PING_LA = 111  #:
    LA_HALT_OVERLAY = 112  #:
    BE_HALT_OVERLAY = 113  #:
    OVERLAY_HALTED = 114  #:
    EXCEPTIONLESS_ABORT = 115  #: Communicate abnormal termination without raising exception
    LA_EXIT = 116  #:
    GS_GROUP_LIST = 117 #:
    GS_GROUP_LIST_RESPONSE = 118 #:
    GS_GROUP_QUERY = 119 #:
    GS_GROUP_QUERY_RESPONSE = 120 #:
    GS_GROUP_DESTROY = 121 #:
    GS_GROUP_DESTROY_RESPONSE = 122 #:
    GS_GROUP_ADD_TO = 123 #:
    GS_GROUP_ADD_TO_RESPONSE = 124 #:
    GS_GROUP_REMOVE_FROM = 125 #:
    GS_GROUP_REMOVE_FROM_RESPONSE = 126 #:
    GS_GROUP_CREATE = 127 #:
    GS_GROUP_CREATE_RESPONSE = 128 #:
    GS_GROUP_KILL = 129 #:
    GS_GROUP_KILL_RESPONSE = 130 #:
    GS_GROUP_CREATE_ADD_TO = 131 #:
    GS_GROUP_CREATE_ADD_TO_RESPONSE = 132 #:
    GS_GROUP_DESTROY_REMOVE_FROM = 133 #:
    GS_GROUP_DESTROY_REMOVE_FROM_RESPONSE = 134 #:


@enum.unique
class FileDescriptor(enum.Enum):
    stdin = 0
    stdout = 1
    stderr = 2


PIPE = subprocess.PIPE
STDOUT = subprocess.STDOUT
DEVNULL = subprocess.DEVNULL


class AbnormalTerminationError(Exception):

    def __init__(self, msg=''):
        self._msg = msg

    def __str__(self):
        return f'{self._msg}'

    def __repr__(self):
        return f"{str(__class__)}({repr(self._msg)})"

@dataclass
class PMIInfo():
    """
    Required information to enable the launching of pmi based applications.
    """

    job_id: int
    lrank: int
    ppn: int
    nid: int
    nnodes: int
    nranks: int
    nidlist: list[int]
    hostlist: list[str]
    control_port: int
    pid_base: int

    @classmethod
    def fromdict(cls, d):
        try:
            return cls(**d)
        except Exception as exc:
            raise ValueError(f'Error deserializing {cls.__name__} {d=}') from exc


class _MsgBase(object):
    """Common base for all messages.

        This common base type for all messages sets up the
        default fields and the serialization strategy for
        now.
    """

    _tc = MessageTypes.INVALID  # deliberately invalid value, overridden

    @enum.unique
    class Errors(enum.Enum):
        INVALID = -1  # deliberately invalid, overridden

    def __init__(self, tag, ref=None, err=None):
        assert isinstance(tag, int)

        self._tag = int(tag)

        if ref is None:
            self._ref = None
        else:
            self._ref = int(ref)

        if err is not None:
            if isinstance(err, self.Errors):
                self._err = err
            elif isinstance(err, int):
                self._err = self.Errors(err)
            else:
                raise NotImplementedError('invalid error parameter')
        else:
            self._err = err

    def get_sdict(self):

        rv = {'_tc': self._tc.value,
              'tag': self.tag}

        if self.err is not None:
            rv['err'] = self.err.value

        if self.ref is not None:
            assert isinstance(self.ref, int)
            rv['ref'] = self.ref

        return rv

    @property
    def tc(self):
        return self._tc

    @classmethod
    def tcv(cls):
        return cls._tc.value

    @property
    def tag(self):
        return self._tag

    @tag.setter
    def tag(self, value):
        self._tag = value

    @property
    def ref(self):
        return self._ref

    @property
    def err(self):
        return self._err

    # the keys in the serialization dictionary must match up
    # with the arguments in the __init__ constructor
    # for all the subclasses
    @classmethod
    def from_sdict(cls, sdict):
        return cls(**sdict)

    def serialize(self):
        return json.dumps(self.get_sdict())

    def __str__(self):
        cn = self.__class__.__name__
        msg = f'{cn}: {self.tag}'
        if hasattr(self, 'p_uid'):
            msg += f' {self.p_uid}'

        if hasattr(self, 'r_c_uid'):
            msg += f'->{self.r_c_uid}'
        return msg

    def __repr__(self):
        fields_to_set = self.get_sdict()
        del fields_to_set['_tc']
        fs = ', '.join([f'{k!s}={v!r}' for k, v in fields_to_set.items()])
        return f'{self.__class__.__name__}({fs})'


# class setup methodology:
# 1) the _tc class variable has the value of the typecode
# for this class.
#
# 2) To allow keyword-initialization of the object
# from a dictionary sent over the wire - which will have the
# _tc member in it - we need to repeat the _tc parameter
# simply to allow __init__ to absorb that parameter without
# an error. We could take the _tc out manually
# in parse, but this would mutate the dictionary which
# is worth avoiding.
#
# 3) This setup lets one initialize the object normally.  In
#    the object constructor the _tc parameter is ignored, and
#    in the serialization the typecode is gotten from the
#    class attribute _tc.

class GSProcessCreate(_MsgBase):
    """
        Refer to :ref:`definition<gsprocesscreate>`
        and :ref:`Common Fields<cfs>` for a description of
        the message structure.
    """

    _tc = MessageTypes.GS_PROCESS_CREATE

    def __init__(self, tag, p_uid, r_c_uid, exe, args, env=None, rundir='',
                 user_name='', options=None, stdin=None, stdout=None, stderr=None,
                 group=None, user=None, umask=- 1, pipesize=None, pmi_required=False,
                 _pmi_info=None, layout=None, _tc=None):

        # Coerce args to a list of strings
        args = list(to_str_iter(args))

        super().__init__(tag)

        if options is None:
            options = {}

        if env is None:
            env = {}

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.exe = exe
        self.args = args
        self.env = env
        self.rundir = rundir
        self.user_name = user_name
        self.options = options
        self.stdin = stdin
        self.stdout = stdout
        self.stderr = stderr
        self.group = group
        self.user = user
        self.umask = umask
        self.pipesize = pipesize

        self.pmi_required = pmi_required

        if _pmi_info is None:
            self._pmi_info = None
        elif isinstance(_pmi_info, dict):
            self._pmi_info = PMIInfo.fromdict(_pmi_info)
        elif isinstance(_pmi_info, PMIInfo):
            self._pmi_info = _pmi_info
        else:
            raise ValueError(f'GS unsupported _pmi_info value {_pmi_info=}')

        if layout is None:
            self.layout = None
        elif isinstance(layout, dict):
            self.layout = ResourceLayout(**_pmi_info)
        elif isinstance(layout, ResourceLayout):
            self.layout = _pmi_info
        else:
            raise ValueError(f'GS unsupported layout value {layout=}')

    @property
    def options(self):
        return self._options

    @options.setter
    def options(self, value):

        if isinstance(value, process_desc.ProcessOptions):
            self._options = value
        else:
            self._options = process_desc.ProcessOptions.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['exe'] = self.exe
        rv['args'] = self.args
        rv['env'] = self.env
        rv['rundir'] = self.rundir
        rv['user_name'] = self.user_name
        rv['options'] = self.options.get_sdict()
        rv['stdin'] = self.stdin
        rv['stdout'] = self.stdout
        rv['stderr'] = self.stderr
        rv['group'] = self.group
        rv['user'] = self.user
        rv['umask'] = self.umask
        rv['pipesize'] = self.pipesize
        rv['pmi_required'] = self.pmi_required
        rv['_pmi_info'] = None if self._pmi_info is None else asdict(self._pmi_info)
        rv['layout'] = None if self.layout is None else asdict(self.layout)

        return rv

    def __str__(self):
        return super().__str__() + f'{self.exe} {self.args}'


class GSProcessCreateResponse(_MsgBase):
    """
        Refer to :ref:`definition<gsprocesscreateresponse>`
        and :ref:`Common Fields<cfs>` for a
        description of the message structure.
    """

    _tc = MessageTypes.GS_PROCESS_CREATE_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Process was created
        FAIL = 1  #: Process was not created
        ALREADY = 2  #: Process exists already

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        if self.Errors.SUCCESS == self.err or self.Errors.ALREADY == self.err:
            self.desc = desc
        elif self.err == self.Errors.FAIL:
            self.err_info = err_info
        else:
            raise NotImplementedError('missing case')

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, process_desc.ProcessDescriptor):
            self._desc = value
        else:
            self._desc = process_desc.ProcessDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if self.err == self.Errors.SUCCESS or self.err == self.Errors.ALREADY:
            rv['desc'] = self.desc.get_sdict()
        elif self.err == self.Errors.FAIL:
            rv['err_info'] = self.err_info

        return rv


class GSProcessList(_MsgBase):
    """
        Refer to :ref:`definition<gsprocesslist>` and :ref:`Common Fields<cfs>` for a description of
        the message structure.
    """

    _tc = MessageTypes.GS_PROCESS_LIST

    def __init__(self, tag, p_uid, r_c_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class GSProcessListResponse(_MsgBase):
    """
        Refer to :ref:`definition<gsprocesslistresponse>` and :ref:`Common Fields<cfs>` for a
        description of the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_LIST_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Always succeeds

    def __init__(self, tag, ref, err, plist=None, _tc=None):
        super().__init__(tag, ref, err)

        if plist is None:
            self.plist = []
        else:
            self.plist = list(plist)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['plist'] = self.plist
        return rv


class GSProcessQuery(_MsgBase):
    """
        Refer to :ref:`definition<gsprocessquery>` and :ref:`Common Fields<cfs>`
        for a description of the message structure.
    """

    _tc = MessageTypes.GS_PROCESS_QUERY

    def __init__(self, tag, p_uid, r_c_uid,
                 t_p_uid=None, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)

        if t_p_uid is not None:
            self.t_p_uid = int(t_p_uid)
        else:
            self.t_p_uid = t_p_uid

        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.t_p_uid is not None:
            rv['t_p_uid'] = self.t_p_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSProcessQueryResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsprocessqueryresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_QUERY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        if desc is None:
            desc = {}
            self._desc = None

        if self.Errors.SUCCESS == self.err:
            self.desc = desc
        elif self.Errors.UNKNOWN == self.err:
            self.err_info = err_info
        else:
            raise NotImplementedError('close case')

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, process_desc.ProcessDescriptor):
            self._desc = value
        else:
            self._desc = process_desc.ProcessDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err:
            rv['desc'] = self.desc.get_sdict()
        else:
            rv['err_info'] = self.err_info

        return rv


class GSProcessKill(_MsgBase):
    """
            Refer to :ref:`definition<gsprocesskill>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_KILL

    def __init__(self, tag, p_uid, r_c_uid, sig, t_p_uid=None, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        self.t_p_uid = t_p_uid
        self.sig = int(sig)
        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['sig'] = self.sig

        if self.t_p_uid is not None:
            rv['t_p_uid'] = self.t_p_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSProcessKillResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsprocesskillresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_KILL_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        FAIL_KILL = 2  #:
        DEAD = 3  #:
        PENDING = 4  #:

    def __init__(self, tag, ref, err, exit_code=0, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info
        self.exit_code = exit_code

    def get_sdict(self):
        rv = super().get_sdict()
        if (self.Errors.UNKNOWN == self.err or
                self.Errors.FAIL_KILL == self.err):
            rv['err_info'] = self.err_info
        else:
            rv['exit_code'] = self.exit_code

        return rv


class GSProcessJoin(_MsgBase):
    """
            Refer to :ref:`definition<gsprocessjoin>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_JOIN

    def __init__(self, tag, p_uid, r_c_uid, timeout=-1, t_p_uid=None, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        self.t_p_uid = t_p_uid
        self.timeout = timeout
        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        rv['timeout'] = self.timeout

        if self.t_p_uid is not None:
            rv['t_p_uid'] = self.t_p_uid
        else:
            rv['user_name'] = self.user_name

        return rv

    def __str__(self):
        first = super().__str__()
        return first + f' {self.t_p_uid}:{self.user_name}'


class GSProcessJoinResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsprocessjoinresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_JOIN_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        TIMEOUT = 2  #:
        SELF = 3  #:

    def __init__(self, tag, ref, err, exit_code=0, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info
        self.exit_code = exit_code

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.SUCCESS == self.err:
            rv['exit_code'] = self.exit_code
        elif self.Errors.UNKNOWN == self.err:
            rv['err_info'] = self.err_info
        elif self.Errors.SELF == self.err:
            rv['err_info'] = self.err_info

        return rv

    def __str__(self):
        msg = super().__str__()

        if self.Errors.SUCCESS == self.err:
            msg += f' exit_code: {self.exit_code}'
        elif self.Errors.UNKNOWN == self.err:
            msg += f' unknown: {self.err_info}'
        else:
            msg += f' timeout'

        return msg


class GSProcessJoinList(_MsgBase):
    """
            Refer to :ref:`definition<gsprocessjoinlist>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_JOIN_LIST

    def __init__(self, tag, p_uid, r_c_uid, timeout=-1, t_p_uid_list=None, user_name_list=None, join_all=False, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        if t_p_uid_list is None:
            t_p_uid_list = []
        self.t_p_uid_list = t_p_uid_list
        if user_name_list is None:
            user_name_list = []
        self.user_name_list = user_name_list
        self.timeout = timeout
        self.join_all = join_all

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['timeout'] = self.timeout
        rv['join_all'] = self.join_all
        if self.t_p_uid_list:
            rv['t_p_uid_list'] = self.t_p_uid_list
        if self.user_name_list:
            rv['user_name_list'] = self.user_name_list

        return rv

    def __str__(self):
        first = super().__str__()
        return first + f' {self.t_p_uid_list}:{self.user_name_list}'


class GSProcessJoinListResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsprocessjoinlistresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_PROCESS_JOIN_LIST_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        TIMEOUT = 2  #:
        SELF = 3  #:
        PENDING = 4  #:

    def __init__(self, tag, ref, puid_status, _tc=None):
        super().__init__(tag, ref)
        self.puid_status = puid_status

    def get_sdict(self):
        rv = super().get_sdict()
        rv['puid_status'] = self.puid_status

        return rv

    def __str__(self):
        msg = super().__str__()
        msg += f' ref: {self.ref}, puid_status: {self.puid_status}'

        return msg


class GSPoolCreate(_MsgBase):
    """
            Refer to :ref:`definition<gspoolcreate>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_POOL_CREATE

    def __init__(self, tag, p_uid, r_c_uid, size, user_name='', options=None, _tc=None):
        super().__init__(tag)

        if options is None:
            options = {}

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.size = int(size)
        self.user_name = user_name
        self.options = options

    @property
    def options(self):
        return self._options

    @options.setter
    def options(self, value):

        if isinstance(value, pool_desc.PoolOptions):
            self._options = value
        else:
            self._options = pool_desc.PoolOptions.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['size'] = self.size
        rv['user_name'] = self.user_name
        rv['options'] = self.options.get_sdict()
        return rv


class GSPoolCreateResponse(_MsgBase):
    """
            Refer to :ref:`definition<gspoolcreateresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_POOL_CREATE_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Pool was created
        FAIL = 1  #: Pool was not created
        ALREADY = 2  #: Pool exists already

    def __init__(self, tag, ref, err, err_code=0, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        if self.Errors.SUCCESS == self.err or self.Errors.ALREADY == self.err:
            self.desc = desc
        elif self.err == self.Errors.FAIL:
            self.err_info = err_info
            self.err_code = err_code
        else:
            raise NotImplementedError('missing case')

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, pool_desc.PoolDescriptor):
            self._desc = value
        else:
            self._desc = pool_desc.PoolDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if self.err == self.Errors.SUCCESS or self.err == self.Errors.ALREADY:
            rv['desc'] = self.desc.get_sdict()
        elif self.err == self.Errors.FAIL:
            rv['err_info'] = self.err_info
            rv['err_code'] = self.err_code

        return rv


class GSPoolList(_MsgBase):
    """
            Refer to :ref:`definition<gspoollist>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_POOL_LIST

    def __init__(self, tag, p_uid, r_c_uid,
                 _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class GSPoolListResponse(_MsgBase):
    """
            Refer to :ref:`definition<gspoollistresponse>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.GS_POOL_LIST_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Always succeeds

    def __init__(self, tag, ref, err, mlist=None, _tc=None):
        super().__init__(tag, ref, err)

        if mlist is None:
            self.mlist = []
        else:
            self.mlist = list(mlist)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['mlist'] = self.mlist
        return rv


class GSPoolQuery(_MsgBase):
    """
            Refer to :ref:`definition<gspoolquery>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_POOL_QUERY

    def __init__(self, tag, p_uid, r_c_uid,
                 m_uid=None, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)

        if m_uid is not None:
            self.m_uid = int(m_uid)
        else:
            self.m_uid = m_uid

        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.m_uid is not None:
            rv['m_uid'] = self.m_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSPoolQueryResponse(_MsgBase):
    """
            Refer to :ref:`definition<gspoolqueryresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_POOL_QUERY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        if desc is None:
            desc = {}
            self._desc = None

        if self.Errors.SUCCESS == self.err:
            self.desc = desc
        elif self.Errors.UNKNOWN == self.err:
            self.err_info = err_info
        else:
            raise NotImplementedError('close case')

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, pool_desc.PoolDescriptor):
            self._desc = value
        else:
            self._desc = pool_desc.PoolDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err:
            rv['desc'] = self.desc.get_sdict()
        else:
            rv['err_info'] = self.err_info

        return rv


class GSPoolDestroy(_MsgBase):
    """
            Refer to :ref:`definition<gspooldestroy>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_POOL_DESTROY

    def __init__(self, tag, p_uid, r_c_uid, m_uid=0, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.m_uid = m_uid
        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.m_uid is not None:
            rv['m_uid'] = self.m_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSPoolDestroyResponse(_MsgBase):
    """
            Refer to :ref:`definition<gspooldestroyresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_POOL_DESTROY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        FAIL = 2  #:
        GONE = 3  #:
        PENDING = 4  #:

    def __init__(self, tag, ref, err, err_code=0, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info
        self.err_code = err_code

    def get_sdict(self):
        rv = super().get_sdict()

        if (self.Errors.UNKNOWN == self.err or
                self.Errors.FAIL == self.err):
            rv['err_info'] = self.err_info
            rv['err_code'] = self.err_code

        return rv


class GSGroupCreate(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupcreate>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_GROUP_CREATE

    def __init__(self, tag, p_uid, r_c_uid, items=None, policy=None, user_name='', _tc=None):
        super().__init__(tag)

        if items is None:
            items = []

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.items = list(items)
        self.user_name = user_name

        if policy is None:
            self.policy = None
        elif isinstance(policy, dict):
            self.policy = Policy(**policy)
        elif isinstance(policy, Policy):
            self.policy = policy
        else:
            raise ValueError(f'GS Groups unsupported policy value {policy=}')

    def get_sdict(self):
        rv = super().get_sdict()

        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['items'] = self.items
        rv['policy'] = self.policy.get_sdict()
        rv['user_name'] = self.user_name

        return rv


class GSGroupCreateResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupcreateresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_CREATE_RESPONSE

    def __init__(self, tag, ref, desc=None, _tc=None):
        super().__init__(tag, ref)

        if desc is None:
            desc = {}
            self._desc = None
        self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['desc'] = None if self.desc is None else self.desc.get_sdict()
        return rv

class GSGroupList(_MsgBase):
    """
            Refer to :ref:`definition<gsgrouplist>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_GROUP_LIST

    def __init__(self, tag, p_uid, r_c_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class GSGroupListResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgrouplistresponse>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_LIST_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Always succeeds

    def __init__(self, tag, ref, err, glist=None, _tc=None):
        super().__init__(tag, ref, err)

        if glist is None:
            self.glist = []
        else:
            self.glist = list(glist)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['glist'] = self.glist
        return rv


class GSGroupQuery(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupquery>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_GROUP_QUERY

    def __init__(self, tag, p_uid, r_c_uid,
                 g_uid=None, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)

        if g_uid is not None:
            self.g_uid = int(g_uid)
        else:
            self.g_uid = g_uid

        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.g_uid is not None:
            rv['g_uid'] = self.g_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSGroupQueryResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupqueryresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_QUERY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        if desc is None:
            desc = {}
            self._desc = None

        if self.Errors.SUCCESS == self.err:
            self.desc = desc
        elif self.Errors.UNKNOWN == self.err:
            self.err_info = err_info
        else:
            raise NotImplementedError('close case')

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err:
            rv['desc'] = self.desc.get_sdict()
        else:
            rv['err_info'] = self.err_info

        return rv


class GSGroupKill(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupkill>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_GROUP_KILL

    def __init__(self, tag, p_uid, r_c_uid, sig, g_uid=None, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.sig = int(sig)
        self.g_uid = g_uid
        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['sig'] = self.sig

        if self.g_uid is not None:
            rv['g_uid'] = self.g_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSGroupKillResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupkillresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_KILL_RESPONSE

    @enum.unique
    class Errors(enum.IntEnum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        DEAD = 2  #:
        PENDING = 3  #:
        ALREADY = 4  #:

    def __init__(self, tag, ref, err, err_info='', desc=None, _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info

        if desc is None:
            desc = {}
            self._desc = None
        self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.UNKNOWN == self.err:
            rv['err_info'] = self.err_info
        if self.Errors.SUCCESS == self.err or self.Errors.ALREADY == self.err:
            rv['desc'] = None if self.desc is None else self.desc.get_sdict()

        return rv


class GSGroupDestroy(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupdestroy>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_GROUP_DESTROY

    def __init__(self, tag, p_uid, r_c_uid, g_uid=None, user_name='', _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.g_uid = g_uid
        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.g_uid is not None:
            rv['g_uid'] = self.g_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSGroupDestroyResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupdestroyresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_DESTROY_RESPONSE

    @enum.unique
    class Errors(enum.IntEnum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        DEAD = 3  #:
        PENDING = 4  #:

    def __init__(self, tag, ref, err, err_info='', desc=None, _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info

        if desc is None:
            desc = {}
            self._desc = None
        self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.UNKNOWN == self.err :
            rv['err_info'] = self.err_info
        if self.Errors.SUCCESS == self.err:
            rv['desc'] = None if self.desc is None else self.desc.get_sdict()

        return rv


class GSGroupAddTo(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupaddto>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_GROUP_ADD_TO

    def __init__(self, tag, p_uid, r_c_uid, g_uid=None, user_name='', items=None, _tc=None):
        super().__init__(tag)

        if items is None:
            items = []

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.g_uid = g_uid
        self.user_name = user_name
        self.items = list(items)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.g_uid is not None:
            rv['g_uid'] = self.g_uid
        else:
            rv['user_name'] = self.user_name

        rv['items'] = self.items

        return rv


class GSGroupAddToResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupaddtoresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_ADD_TO_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        FAIL = 2  #:
        DEAD = 3  #:
        PENDING = 4  #:

    def __init__(self, tag, ref, err, err_info='', desc=None, _tc=None):
        super().__init__(tag, ref, err)

        self.err_info = err_info

        if desc is None:
            desc = {}
            self._desc = None
        self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if (self.Errors.UNKNOWN == self.err or
                self.Errors.FAIL == self.err):
            rv['err_info'] = self.err_info

        rv['desc'] = None if self.desc is None else self.desc.get_sdict()

        return rv

class GSGroupCreateAddTo(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupcreateaddto>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_GROUP_CREATE_ADD_TO

    def __init__(self, tag, p_uid, r_c_uid, g_uid=None, user_name='', items=None, policy=None, _tc=None):
        super().__init__(tag)

        if items is None:
            items = []

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.g_uid = g_uid
        self.user_name = user_name
        self.items = list(items)

        if policy is None:
            self.policy = None
        elif isinstance(policy, dict):
            self.policy = Policy(**policy)
        elif isinstance(policy, Policy):
            self.policy = policy
        else:
            raise ValueError(f'GS Groups unsupported policy value {policy=}')

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.g_uid is not None:
            rv['g_uid'] = self.g_uid
        else:
            rv['user_name'] = self.user_name

        rv['items'] = self.items
        rv['policy'] = self.policy.get_sdict()

        return rv


class GSGroupCreateAddToResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupcreateaddtoresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_CREATE_ADD_TO_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        DEAD = 2  #:
        PENDING = 3  #:

    def __init__(self, tag, ref, err, err_info='', desc=None, _tc=None):
        super().__init__(tag, ref, err)

        self.err_info = err_info

        if desc is None:
            desc = {}
            self._desc = None
        self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if (self.Errors.UNKNOWN == self.err or
                self.Errors.PENDING == self.err):
            rv['err_info'] = self.err_info

        rv['desc'] = None if self.desc is None else self.desc.get_sdict()

        return rv


class GSGroupRemoveFrom(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupremovefrom>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_GROUP_REMOVE_FROM

    def __init__(self, tag, p_uid, r_c_uid, g_uid=None, user_name='', items=None, _tc=None):
        super().__init__(tag)

        if items is None:
            items = []

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.g_uid = g_uid
        self.user_name = user_name
        self.items = list(items)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.g_uid is not None:
            rv['g_uid'] = self.g_uid
        else:
            rv['user_name'] = self.user_name

        rv['items'] = self.items

        return rv


class GSGroupRemoveFromResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupremovefromresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_REMOVE_FROM_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        FAIL = 2  #:
        DEAD = 3  #:
        PENDING = 4  #:

    def __init__(self, tag, ref, err, err_info='', desc=None, _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info

        if desc is None:
            desc = {}
            self._desc = None
        self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if (self.Errors.UNKNOWN == self.err or
                self.Errors.FAIL == self.err):
            rv['err_info'] = self.err_info

        rv['desc'] = None if self.desc is None else self.desc.get_sdict()

        return rv


class GSGroupDestroyRemoveFrom(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupdestroyremovefrom>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_GROUP_DESTROY_REMOVE_FROM

    def __init__(self, tag, p_uid, r_c_uid, g_uid=None, user_name='', items=None, _tc=None):
        super().__init__(tag)

        if items is None:
            items = []

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.g_uid = g_uid
        self.user_name = user_name
        self.items = list(items)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        if self.g_uid is not None:
            rv['g_uid'] = self.g_uid
        else:
            rv['user_name'] = self.user_name

        rv['items'] = self.items

        return rv


class GSGroupDestroyRemoveFromResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsgroupdestroyremovefromresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_GROUP_DESTROY_REMOVE_FROM_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        FAIL = 2  #:
        DEAD = 3  #:
        PENDING = 4  #:

    def __init__(self, tag, ref, err, err_info='', desc=None, _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info

        if desc is None:
            desc = {}
            self._desc = None
        self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, group_desc.GroupDescriptor):
            self._desc = value
        else:
            self._desc = group_desc.GroupDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        if (self.Errors.UNKNOWN == self.err or
                self.Errors.FAIL == self.err):
            rv['err_info'] = self.err_info

        rv['desc'] = None if self.desc is None else self.desc.get_sdict()

        return rv


class GSChannelCreate(_MsgBase):
    """
            Refer to :ref:`definition<gschannelcreate>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_CREATE

    def __init__(self, tag, p_uid, r_c_uid, m_uid, options=None, user_name='', _tc=None):
        super().__init__(tag)

        if options is None:
            options = {}

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.m_uid = int(m_uid)
        self.user_name = user_name
        self.options = options

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['user_name'] = self.user_name
        rv['m_uid'] = self.m_uid
        rv['options'] = self.options.get_sdict()
        return rv

    @property
    def options(self):
        return self._options

    @options.setter
    def options(self, value):
        if isinstance(value, channel_desc.ChannelOptions):
            self._options = value
        else:
            self._options = channel_desc.ChannelOptions.from_sdict(value)


class GSChannelCreateResponse(_MsgBase):
    """
        Refer to :ref:`definition<gschannelcreateresponse>` and :ref:`Common Fields<cfs>` for a
        description of the message structure.
    """

    _tc = MessageTypes.GS_CHANNEL_CREATE_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:
        ALREADY = 2  #:

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        if desc is None:
            desc = {}

        if self.Errors.SUCCESS == self.err or self.Errors.ALREADY == self.err:
            self.desc = desc
        else:
            self.err_info = err_info

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, channel_desc.ChannelDescriptor):
            self._desc = value
        else:
            self._desc = channel_desc.ChannelDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err or self.Errors.ALREADY == self.err:
            rv['desc'] = self.desc.get_sdict()
        else:
            rv['err_info'] = self.err_info

        return rv


class GSChannelList(_MsgBase):
    """
        Refer to :ref:`definition<gschannellist>` and to
        :ref:`Common Fields<cfs>` for a description of the message structure.
    """

    _tc = MessageTypes.GS_CHANNEL_LIST

    def __init__(self, tag, p_uid, r_c_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class GSChannelListResponse(_MsgBase):
    """
        Refer to :ref:`definition<gschannellistresponse>` and :ref:`Common Fields<cfs>` for a
        description of the message structure.
    """

    _tc = MessageTypes.GS_CHANNEL_LIST_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Always succeeds

    def __init__(self, tag, ref, err, clist=None, _tc=None):
        super().__init__(tag, ref, err)

        if clist is None:
            self.clist = []
        else:
            self.clist = list(clist)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['clist'] = self.clist
        return rv


class GSChannelQuery(_MsgBase):
    """
            Refer to :ref:`definition<gschannelquery>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_QUERY

    def __init__(self, tag, p_uid, r_c_uid, c_uid=None, user_name='',
                 inc_refcnt=False, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        self.inc_refcnt = inc_refcnt

        if c_uid is not None:
            self.c_uid = int(c_uid)
        else:
            self.c_uid = c_uid

        self.user_name = user_name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['inc_refcnt'] = self.inc_refcnt

        if self.c_uid is not None:
            rv['c_uid'] = self.c_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSChannelQueryResponse(_MsgBase):
    """
            Refer to :ref:`definition<gschannelqueryresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_QUERY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        self.err_info = err_info
        if desc is None:
            desc = {}

        if self.Errors.SUCCESS == self.err:
            self.desc = desc
        elif self.Errors.UNKNOWN == self.err:
            self.err_info = err_info
        else:
            raise NotImplementedError('open enum')

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, channel_desc.ChannelDescriptor):
            self._desc = value
        else:
            self._desc = channel_desc.ChannelDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err:
            rv['desc'] = self.desc.get_sdict()
        elif self.Errors.UNKNOWN == self.err:
            rv['err_info'] = self.err_info
        else:
            raise NotImplementedError('open enum')

        return rv


class GSChannelDestroy(_MsgBase):
    """
            Refer to :ref:`definition<gschanneldestroy>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_DESTROY

    def __init__(self, tag, p_uid, r_c_uid, c_uid=None, user_name='',
                 reply_req=True, dec_ref=False, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.c_uid = c_uid
        self.user_name = user_name
        self.reply_req = reply_req
        self.dec_ref = dec_ref

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['reply_req'] = self.reply_req
        rv['dec_ref'] = self.dec_ref

        if self.c_uid is not None:
            rv['c_uid'] = self.c_uid
        else:
            rv['user_name'] = self.user_name

        return rv


class GSChannelDestroyResponse(_MsgBase):
    """
            Refer to :ref:`definition<gschanneldestroyresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_DESTROY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        UNKNOWN_CHANNEL = 2  #:
        BUSY = 3  #:

    def __init__(self, tag, ref, err, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        self.err_info = err_info

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS != self.err:
            rv['err_info'] = self.err_info

        return rv


class GSChannelJoin(_MsgBase):
    """
            Refer to :ref:`definition<gschanneljoin>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_JOIN

    def __init__(self, tag, p_uid, r_c_uid, name, timeout=-1, _tc=None):
        super().__init__(tag)
        self.timeout = timeout
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.name = name

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['name'] = self.name
        rv['timeout'] = self.timeout
        return rv


class GSChannelJoinResponse(_MsgBase):
    """
            Refer to :ref:`definition<gschanneljoinresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_JOIN_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        TIMEOUT = 1  #:
        DEAD = 2  #:

    def __init__(self, tag, ref, err, desc=None, _tc=None):
        super().__init__(tag, ref, err)

        if desc is None:
            desc = {}

        if self.Errors.SUCCESS == self.err:
            self.desc = desc

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, channel_desc.ChannelDescriptor):
            self._desc = value
        else:
            self._desc = channel_desc.ChannelDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err:
            rv['desc'] = self.desc.get_sdict()

        return rv


class GSChannelDetach(_MsgBase):
    """
            Refer to :ref:`definition<gschanneldetach>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_DETACH

    def __init__(self, tag, p_uid, r_c_uid, c_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        self.c_uid = int(c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['c_uid'] = self.c_uid
        return rv


class GSChannelDetachResponse(_MsgBase):
    """
            Refer to :ref:`definition<gschanneldetachresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_DETACH_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        UNKNOWN_CHANNEL = 2  #:
        NOT_ATTACHED = 3  #:

    def __init__(self, tag, ref, err, _tc=None):
        super().__init__(tag, ref, err)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSChannelGetSendH(_MsgBase):
    """
            Refer to :ref:`definition<gschannelgetsendh>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    # TODO: REMOVE, DEPRECATED
    _tc = MessageTypes.GS_CHANNEL_GET_SENDH

    def __init__(self, tag, p_uid, r_c_uid, c_uid, _tc=None):
        raise NotImplementedError('OBE')
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.c_uid = int(c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['c_uid'] = self.c_uid
        return rv


class GSChannelGetSendHResponse(_MsgBase):
    """
            Refer to :ref:`definition<gschannelgetsendhresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    # TODO: REMOVE, DEPRECATED
    _tc = MessageTypes.GS_CHANNEL_GET_SENDH_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        UNKNOWN_CHANNEL = 2  #:
        NOT_ATTACHED = 3  #:
        CANT = 4  #:

    def __init__(self, tag, ref, err, sendh=None, err_info='', _tc=None):
        raise NotImplementedError('OBE')
        super().__init__(tag, ref, err)
        if sendh is None:
            sendh = {}

        self.sendh = sendh
        self.err_info = err_info

    @property
    def sendh(self):
        return self._sendh

    @sendh.setter
    def sendh(self, value):
        # FIXME, once we have ChannelSendHandle
        # if isinstance(value, ChannelSendHandle):
        #    self._sendh = value
        # else:
        #    self._sendh = ChannelSendHandle.from_sdict(value)

        self._sendh = value

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.CANT == self.err:
            rv['err_info'] = self.err_info
        elif self.Errors.SUCCESS == self.err:
            rv['sendh'] = self.sendh.get_sdict()

        return rv


class GSChannelGetRecvH(_MsgBase):
    """
            Refer to :ref:`definition<gschannelgetrecvh>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    # TODO: REMOVE, DEPRECATED
    _tc = MessageTypes.GS_CHANNEL_GET_RECVH

    def __init__(self, tag, p_uid, r_c_uid, c_uid, _tc=None):
        raise NotImplementedError('OBE')
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.c_uid = int(c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['c_uid'] = self.c_uid
        return rv


class GSChannelGetRecvHResponse(_MsgBase):
    """
            Refer to :ref:`definition<gschannelgetrecvhresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    # TODO: REMOVE, DEPRECATED
    _tc = MessageTypes.GS_CHANNEL_GET_RECVH_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        UNKNOWN_CHANNEL = 2  #:
        NOT_ATTACHED = 3  #:
        CANT = 4  #:

    def __init__(self, tag, ref, err, recvh=None, err_info='', _tc=None):
        raise NotImplementedError('OBE')
        super().__init__(tag, ref, err)
        if recvh is None:
            recvh = {}
        self.recvh = recvh
        self.err_info = err_info

    @property
    def recvh(self):
        return self._recvh

    @recvh.setter
    def recvh(self, value):
        # FIXME, once we have ChannelRecvHandle
        # if isinstance(value, ChannelRecvHandle):
        #    self._recvh = value
        # else:
        #    self._recvh = ChannelRecvHandle.from_sdict(value)

        self._recvh = value

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.CANT == self.err:
            rv['err_info'] = self.err_info
        elif self.Errors.SUCCESS == self.err:
            rv['recvh'] = self.recvh.get_sdict()

        return rv


class GSNodeList(_MsgBase):
    """
    *type enum*
        GS_NODE_LIST (= 102)

    *purpose*
        Return a list of tuples of ``h_uid`` for all nodes currently registered.

    *fields*
        None additional

    *response*
        GSNodeListResponse

    *see also*
        GSNodeQuery

    *see also*
        refer to the :ref:`cfs` section for additional request message fields
    """

    _tc = MessageTypes.GS_NODE_LIST

    def __init__(self, tag, p_uid, r_c_uid,
                 _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class GSNodeListResponse(_MsgBase):
    """
    *type enum*
        GS_NODE_LIST_RESPONSE (= 103)

    *purpose*
        Responds with a list of ``h_uid`` for all the
        nodes currently registered.

    *fields*
        **hlist**
            - list of nonnegative integers

    *request*
        GSNodeList

    *see also*
        GSNodeQuery
    """

    _tc = MessageTypes.GS_NODE_LIST_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Always succeeds

    def __init__(self, tag, ref, err, hlist=None, _tc=None):
        super().__init__(tag, ref, err)

        if hlist is None:
            self.hlist = []
        else:
            self.hlist = list(hlist)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['hlist'] = self.hlist
        return rv


class GSNodeQuery(_MsgBase):
    """
            Refer to :ref:`definition<gsnodequery>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_NODE_QUERY

    def __init__(self, tag, p_uid:int, r_c_uid:int, name:str='', h_uid=None, _tc=None):

        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        self.name = name
        self.h_uid = h_uid

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['name'] = self.name
        rv['h_uid'] = self.h_uid

        return rv


class GSNodeQueryResponse(_MsgBase):
    """
            Refer to :ref:`definition<gsnodequeryresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.GS_NODE_QUERY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        self.err_info = err_info
        if desc is None:
            desc = {}

        if self.Errors.SUCCESS == self.err:
            self.desc = desc
        elif self.Errors.UNKNOWN == self.err:
            self.err_info = err_info
        else:
            raise NotImplementedError('open enum')

    @property
    def desc(self):
        return self._desc

    @desc.setter
    def desc(self, value):
        if isinstance(value, node_desc.NodeDescriptor):
            self._desc = value
        else:
            self._desc = node_desc.NodeDescriptor.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err:
            rv['desc'] = self.desc.get_sdict()
        elif self.Errors.UNKNOWN == self.err:
            rv['err_info'] = self.err_info
        else:
            raise NotImplementedError('open enum')

        return rv


class GSNodeQueryTotalCPUCount(_MsgBase):
    """
    *type enum*
        GS_NODE_QUERY_TOTAL_CPU_COUNT (= 104)

    *purpose*
        Asks GS to return the total number of CPUS beloging to all of the registered nodes.

    *see also*
        refer to the :ref:`cfs` section for additional request message fields

    *response*
        GSNodeQueryTotalCPUCountResponse
    """

    _tc = MessageTypes.GS_NODE_QUERY_TOTAL_CPU_COUNT

    def __init__(self, tag, p_uid:int, r_c_uid:int, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid

        return rv


class GSNodeQueryTotalCPUCountResponse(_MsgBase):
    """
    *type enum*
        GS_NODE_QUERY_TOTAL_CPU_COUNT_RESPONSE (= 105)

    *purpose*
         Return the total number of CPUS beloging to all of the registered nodes.

    *fields*
        Alternatives on ``err``:

        SUCCESS (= 0)

            The machine descriptor was successfully constructed

            **total_cpus**
                - total number of CPUS beloging to all of the registered nodes.

        UNKNOWN ( = 1)
            An unknown error has occured.

    *request*
        GSNodeQueryTotalCPUCount

    *see also*
        GSNodeQuery
    """

    _tc = MessageTypes.GS_NODE_QUERY_TOTAL_CPU_COUNT_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:

    def __init__(self, tag, ref, err, total_cpus=0, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        self.err_info = err_info

        if self.Errors.SUCCESS == self.err:
            self.total_cpus = total_cpus
        elif self.Errors.UNKNOWN == self.err:
            self.err_info = err_info
        else:
            raise NotImplementedError('open enum')

    @property
    def total_cpus(self):
        return self._total_cpus

    @total_cpus.setter
    def total_cpus(self, value):
        self._total_cpus = int(value)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.SUCCESS == self.err:
            rv['total_cpus'] = self.total_cpus
        elif self.Errors.UNKNOWN == self.err:
            rv['err_info'] = self.err_info
        else:
            raise NotImplementedError('open enum')

        return rv


class AbnormalTermination(_MsgBase):
    """
            Refer to :ref:`definition<abnormaltermination>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.ABNORMAL_TERMINATION

    def __init__(self, tag, err_info='', _tc=None):
        super().__init__(tag)
        self.err_info = err_info

    def get_sdict(self):
        rv = super().get_sdict()
        rv['err_info'] = self.err_info
        return rv

    def __str__(self):
        traceback = str(self.err_info)
        if len(traceback) > 0:
            return str(super()) + f'\n***** Dragon Traceback: *****\n{traceback}***** End of Dragon Traceback: *****\n'

        return str(super())


class ExceptionlessAbort(_MsgBase):
    """
            Refer to :ref:`definition<exceptionlessabort>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.EXCEPTIONLESS_ABORT

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSStarted(_MsgBase):
    """
            Refer to :ref:`definition<gsstarted>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    """Indicates Global Services head process has started."""

    _tc = MessageTypes.GS_STARTED

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSPingSH(_MsgBase):
    """
            Refer to :ref:`definition<gspingsh>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_PING_SH

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSIsUp(_MsgBase):
    """
            Refer to :ref:`definition<gsisup>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_IS_UP

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSPingProc(_MsgBase):
    """
            Refer to :ref:`definition<gspingproc>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_PING_PROC

    def __init__(self, tag, mode=None, argdata=None, _tc=None):
        super().__init__(tag)

        self.mode = process_desc.mk_argmode_from_default(mode)

        self.argdata = argdata

    def get_sdict(self):
        rv = super().get_sdict()

        rv['mode'] = self.mode.value
        rv['argdata'] = self.argdata

        return rv


class GSDump(_MsgBase):
    """
            Refer to :ref:`definition<gsdump>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_DUMP

    def __init__(self, tag, filename, _tc=None):
        super().__init__(tag)
        self.filename = filename

    def get_sdict(self):
        rv = super().get_sdict()
        rv['filename'] = self.filename
        return rv


class GSHeadExit(_MsgBase):
    """
            Refer to :ref:`definition<gsheadexit>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_HEAD_EXIT

    def __init__(self, tag, exit_code=0, _tc=None):
        super().__init__(tag)
        self.exit_code = exit_code

    def get_sdict(self):
        rv = super().get_sdict()
        rv['exit_code'] = self.exit_code
        return rv


class GSTeardown(_MsgBase):
    """
            Refer to :ref:`definition<gsteardown>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_TEARDOWN

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSUnexpected(_MsgBase):
    """
            Refer to :ref:`definition<gsunexpected>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.GS_UNEXPECTED

    def __init__(self, tag, ref, _tc=None):
        super().__init__(tag, ref)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSChannelRelease(_MsgBase):
    """
            Refer to :ref:`definition<gschannelrelease>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.GS_CHANNEL_RELEASE

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class GSHalted(_MsgBase):
    """
            Refer to :ref:`definition<gshalted>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.GS_HALTED

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class SHProcessCreate(_MsgBase):
    """
            Refer to :ref:`definition<shprocesscreate>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

            The initial_stdin is a string which if non-empty is written along with a terminating newline character
            to the stdin of the newly created process.

            The stdin, stdout, and stderr are all either None or an instance of SHChannelCreate to be processed
            by the local services component.

    """

    _tc = MessageTypes.SH_PROCESS_CREATE

    def __init__(self, tag, p_uid, r_c_uid, t_p_uid, exe, args, env=None, rundir='', options=None, initial_stdin='',
                 stdin=None, stdout=None, stderr=None, group=None, user=None, umask=- 1, pipesize=None,
                 stdin_msg=None, stdout_msg=None, stderr_msg=None, pmi_info=None, layout=None, _tc=None):
        super().__init__(tag)

        if options is None:
            options = {}

        if env is None:
            env = {}

        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.t_p_uid = t_p_uid
        self.exe = exe
        self.args = args
        self.env = env
        self.rundir = rundir
        self.options = options
        self.initial_stdin = initial_stdin
        self.stdin = stdin
        self.stdout = stdout
        self.stderr = stderr
        self.group = group
        self.user = user
        self.umask = umask
        self.pipesize = pipesize
        if stdin_msg is None or isinstance(stdin_msg, SHChannelCreate):
            self.stdin_msg = stdin_msg
        else:
            self.stdin_msg = SHChannelCreate.from_sdict(stdin_msg)

        if stdout_msg is None or isinstance(stdout_msg, SHChannelCreate):
            self.stdout_msg = stdout_msg
        else:
            self.stdout_msg = SHChannelCreate.from_sdict(stdout_msg)

        if stderr_msg is None or isinstance(stderr_msg, SHChannelCreate):
            self.stderr_msg = stderr_msg
        else:
            self.stderr_msg = SHChannelCreate.from_sdict(stderr_msg)

        if pmi_info is None:
            self.pmi_info = None
        elif isinstance(pmi_info, dict):
            self.pmi_info = PMIInfo.fromdict(pmi_info)
        elif isinstance(pmi_info, PMIInfo):
            self.pmi_info = pmi_info
        else:
            raise ValueError(f'LS unsupported pmi_info value {pmi_info=}')

        if layout is None:
            self.layout = None
        elif isinstance(layout, dict):
            self.layout = ResourceLayout(**pmi_info)
        elif isinstance(layout, ResourceLayout):
            self.layout = layout
        else:
            raise ValueError(f'LS unsupported layout value {layout=}')

    @property
    def options(self):
        return self._options

    @options.setter
    def options(self, value):

        if isinstance(value, dso.ProcessOptions):
            self._options = value
        else:
            self._options = dso.ProcessOptions.from_sdict(value)

    def get_sdict(self):
        rv = super().get_sdict()

        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['t_p_uid'] = self.t_p_uid
        rv['exe'] = self.exe
        rv['args'] = self.args
        rv['env'] = self.env
        rv['rundir'] = self.rundir
        rv['options'] = self.options.get_sdict()
        rv['initial_stdin'] = self.initial_stdin
        rv['stdin'] = self.stdin
        rv['stdout'] = self.stdout
        rv['stderr'] = self.stderr
        rv['group'] = self.group
        rv['user'] = self.user
        rv['umask'] = self.umask
        rv['pipesize'] = self.pipesize
        rv['stdin_msg'] = (None if self.stdin_msg is None else self.stdin_msg.get_sdict())
        rv['stdout_msg'] = (None if self.stdout_msg is None else self.stdout_msg.get_sdict())
        rv['stderr_msg'] = (None if self.stderr_msg is None else self.stderr_msg.get_sdict())
        rv['pmi_info'] = None if self.pmi_info is None else asdict(self.pmi_info)
        rv['layout'] = None if self.layout is None else asdict(self.layout)
        return rv


class SHProcessCreateResponse(_MsgBase):
    """
            Refer to :ref:`definition<shprocesscreateresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.SH_PROCESS_CREATE_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:

    def __init__(self, tag, ref, err, err_info='', stdin_resp=None, stdout_resp=None, stderr_resp=None, _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info


        if stdin_resp is None or isinstance(stdin_resp, SHChannelCreateResponse):
            self.stdin_resp = stdin_resp
        else:
            self.stdin_resp = SHChannelCreateResponse.from_sdict(stdin_resp)

        if stdout_resp is None or isinstance(stdout_resp, SHChannelCreateResponse):
            self.stdout_resp = stdout_resp
        else:
            self.stdout_resp = SHChannelCreateResponse.from_sdict(stdout_resp)

        if stderr_resp is None or isinstance(stderr_resp, SHChannelCreateResponse):
            self.stderr_resp = stderr_resp
        else:
            self.stderr_resp = SHChannelCreateResponse.from_sdict(stderr_resp)

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.FAIL == self.err:
            rv['err_info'] = self.err_info
        rv['stdin_resp'] = (None if self.stdin_resp is None else self.stdin_resp.get_sdict())
        rv['stdout_resp'] = (None if self.stdout_resp is None else self.stdout_resp.get_sdict())
        rv['stderr_resp'] = (None if self.stderr_resp is None else self.stderr_resp.get_sdict())
        return rv


class SHProcessKill(_MsgBase):
    """
            Refer to :ref:`definition<shprocesskill>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.SH_PROCESS_KILL

    def __init__(self, tag, p_uid, r_c_uid, t_p_uid, sig, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        self.t_p_uid = t_p_uid
        self.sig = sig

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['t_p_uid'] = self.t_p_uid
        rv['sig'] = self.sig
        return rv


class SHProcessKillResponse(_MsgBase):
    _tc = MessageTypes.SH_PROCESS_KILL_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0
        FAIL = 1

    def __init__(self, tag, ref, err, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info

    def get_sdict(self):
        rv = super().get_sdict()
        if self.Errors.FAIL == self.err:
            rv['err_info'] = self.err_info

        return rv


class SHProcessExit(_MsgBase):
    """
        Refer to :ref:`definition<shprocessexit>` and to
        :ref:`Common Fields<cfs>` for a description of
        the message structure.
    """

    _tc = MessageTypes.SH_PROCESS_EXIT

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:

    def __init__(self, tag, p_uid, exit_code=0, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.exit_code = exit_code

    def get_sdict(self):
        rv = super().get_sdict()
        rv['exit_code'] = self.exit_code
        rv['p_uid'] = self.p_uid
        return rv


class SHPoolCreate(_MsgBase):
    """
        Refer to :ref:`definition<shpoolcreate>` and
        :ref:`Common Fields<cfs>` for a description of the message structure.
    """

    _tc = MessageTypes.SH_POOL_CREATE

    def __init__(self, tag, p_uid, r_c_uid, size, m_uid, name, attr='', _tc=None):
        super().__init__(tag)

        self.m_uid = int(m_uid)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.size = int(size)
        self.name = name
        self.attr = attr

    def get_sdict(self):
        rv = super().get_sdict()

        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['size'] = self.size
        rv['m_uid'] = self.m_uid
        rv['name'] = self.name
        rv['attr'] = self.attr
        return rv


class SHPoolCreateResponse(_MsgBase):
    """
        Refer to :ref:`definition<shpoolcreateresponse>` and
        :ref:`Common Fields<cfs>` for a description of the message structure.
    """

    _tc = MessageTypes.SH_POOL_CREATE_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Pool was created
        FAIL = 1  #: Pool was not created

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        if self.Errors.SUCCESS == self.err:
            self.desc = desc
        elif self.err == self.Errors.FAIL:
            self.err_info = err_info
        else:
            raise NotImplementedError('close case')

    def get_sdict(self):
        rv = super().get_sdict()

        if self.err == self.Errors.SUCCESS:
            rv['desc'] = self.desc
        elif self.err == self.Errors.FAIL:
            rv['err_info'] = self.err_info
        else:
            raise NotImplementedError('close case')

        return rv


class SHPoolDestroy(_MsgBase):
    """
        Refer to :ref:`definition<shpooldestroy>` and to
        the :ref:`Common Fields<cfs>` for a description of
        the message structure.
    """

    _tc = MessageTypes.SH_POOL_DESTROY

    def __init__(self, tag, p_uid, r_c_uid, m_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.m_uid = int(m_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['m_uid'] = self.m_uid

        return rv


class SHPoolDestroyResponse(_MsgBase):
    """
        Refer to :ref:`definition<shpooldestroyresponse>` and :ref:`Common Fields<cfs>` for a
        description of the message structure.
    """

    _tc = MessageTypes.SH_POOL_DESTROY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:

    def __init__(self, tag, ref, err, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.FAIL == self.err:
            rv['err_info'] = self.err_info

        return rv


class SHExecMemRequest(_MsgBase):
    """
            Refer to :ref:`definition<shexecmemrequest>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.SH_EXEC_MEM_REQUEST

    @enum.unique
    class KINDS(enum.Enum):
        MEMORY_POOL_REQUEST = 0
        MEMORY_ALLOC_REQUEST = 1

    def __init__(self, tag, p_uid, r_c_uid, kind, request, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.kind = kind
        self.request = request

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['kind'] = self.kind
        rv['request'] = self.request

        return rv


class SHExecMemResponse(_MsgBase):
    """
            Refer to :ref:`definition<shexecmemresponse>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.SH_EXEC_MEM_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:

    def __init__(self, tag, ref, err, err_code=0, err_info='', response=None, _tc=None):
        super().__init__(tag, ref, err)
        self.response = response
        self.err_info = err_info
        self.err_code = err_code

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.FAIL == self.err:
            rv['err_info'] = self.err_info
            rv['err_code'] = self.err_code
        elif self.Errors.SUCCESS == self.err:
            rv['response'] = self.response

        return rv


class SHChannelCreate(_MsgBase):
    """
        Refer to :ref:`definition<shchannelcreate>` and :ref:`Common Fields<cfs>`
        for a description of the message structure.
    """

    _tc = MessageTypes.SH_CHANNEL_CREATE

    def __init__(self, tag, p_uid, r_c_uid, m_uid, c_uid, options=None, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.m_uid = int(m_uid)
        self.c_uid = int(c_uid)

        if options is None:
            options = {}

        self.options = options

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['m_uid'] = self.m_uid
        rv['c_uid'] = self.c_uid
        rv['options'] = self.options.get_sdict()

        return rv

    @property
    def options(self):
        return self._options

    @options.setter
    def options(self, value):
        if isinstance(value, dso.ChannelOptions):
            self._options = value
        else:
            self._options = dso.ChannelOptions.from_sdict(value)


class SHChannelCreateResponse(_MsgBase):
    """
        Refer to :ref:`definition<shchannelcreateresponse>` and :ref:`Common Fields<cfs>` for a
        description of the message structure.
    """

    _tc = MessageTypes.SH_CHANNEL_CREATE_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:

    def __init__(self, tag, ref, err, desc=None, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        self.err_info = err_info
        self.desc = desc

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.SUCCESS == self.err:
            rv['desc'] = self.desc
        else:
            rv['err_info'] = self.err_info

        return rv


class SHChannelDestroy(_MsgBase):
    """
            Refer to :ref:`definition<shchanneldestroy>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.SH_CHANNEL_DESTROY

    def __init__(self, tag, p_uid, r_c_uid, c_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = p_uid
        self.r_c_uid = int(r_c_uid)
        self.c_uid = int(c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['c_uid'] = self.c_uid

        return rv


class SHChannelDestroyResponse(_MsgBase):
    """
            Refer to :ref:`definition<shchanneldestroyresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.SH_CHANNEL_DESTROY_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:

    def __init__(self, tag, ref, err, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.FAIL == self.err:
            rv['err_info'] = self.err_info
        return rv


class SHLockChannel(_MsgBase):
    """
            Refer to :ref:`definition<shlockchannel>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.SH_LOCK_CHANNEL

    def __init__(self, tag, p_uid, r_c_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class SHLockChannelResponse(_MsgBase):
    """
            Refer to :ref:`definition<shlockchannelresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.SH_LOCK_CHANNEL_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        ALREADY = 1  #:

    def __init__(self, tag, ref, err, _tc=None):
        super().__init__(tag, ref, err)

    def get_sdict(self):
        rv = super().get_sdict()

        return rv


class SHAllocMsg(_MsgBase):
    """
            Refer to :ref:`definition<shallocmsg>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_ALLOC_MSG

    def __init__(self, tag, p_uid, r_c_uid, _tc=None):
        super().__init__(tag)
        self.r_c_uid = int(r_c_uid)
        self.p_uid = int(p_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class SHAllocMsgResponse(_MsgBase):
    """
            Refer to :ref:`definition<shallocmsgresponse>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.SH_ALLOC_MSG_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:

    def __init__(self, tag, ref, err, _tc=None):
        super().__init__(tag, ref, err)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class SHAllocBlock(_MsgBase):
    """
            Refer to :ref:`definition<shallocblock>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.SH_ALLOC_BLOCK

    def __init__(self, tag, p_uid, r_c_uid, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        return rv


class SHAllocBlockResponse(_MsgBase):
    """
            Refer to :ref:`definition<shallocblockresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.SH_ALLOC_BLOCK_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        UNKNOWN = 1  #:
        FAIL = 2  #:

    def __init__(self, tag, ref, err, seg_name='',
                 offset=0, err_info='', _tc=None):
        super().__init__(tag, ref, err)

        self.seg_name = seg_name
        self.offset = offset
        self.err_info = err_info

    def get_sdict(self):
        rv = super().get_sdict()

        if self.Errors.SUCCESS == self.err:
            rv['seg_name'] = self.seg_name
            rv['offset'] = self.offset
        elif self.Errors.FAIL == self.err:
            rv['err_info'] = self.err_info

        return rv


class SHChannelsUp(_MsgBase):
    """
            Refer to :ref:`definition<shchannelsup>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.SH_CHANNELS_UP

    def __init__(self, tag, ip_addrs, host_id, host_name, shep_cd, gs_cd, idx=0, _tc=None):
        super().__init__(tag)
        self.ip_addrs = ip_addrs
        self.host_id = host_id
        self.host_name = host_name
        self.shep_cd = shep_cd
        # On the primary node the gs_cd is set to the base64 encoded gs channel descriptor.
        # Otherwise, it is ignored and presumably the empty string.
        self.gs_cd = gs_cd
        self.idx = idx

    def get_sdict(self):
        rv = super().get_sdict()
        rv['ip_addrs'] = self.ip_addrs
        rv['host_id'] = self.host_id
        rv['host_name'] = self.host_name
        rv['shep_cd'] = self.shep_cd
        rv['gs_cd'] = self.gs_cd
        rv['idx'] = self.idx
        return rv


class SHPingGS(_MsgBase):
    """
            Refer to :ref:`definition<shpinggs>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_PING_GS

    def __init__(self, tag, idx=0, node_sdesc=None, _tc=None):
        super().__init__(tag)
        self.idx = idx
        self.node_sdesc = node_sdesc

    def get_sdict(self):
        rv = super().get_sdict()
        rv['idx'] = self.idx
        rv['node_sdesc'] = self.node_sdesc

        return rv

class SHTeardown(_MsgBase):
    """
            Refer to :ref:`definition<shteardown>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_TEARDOWN

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class SHPingBE(_MsgBase):
    """
            Refer to :ref:`definition<shpingbe>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_PING_BE
    EMPTY = B64.bytes_to_str(b'')

    def __init__(self, tag, shep_cd=EMPTY, be_cd=EMPTY, gs_cd=EMPTY,
                 default_pd=EMPTY, inf_pd=EMPTY, _tc=None):
        super().__init__(tag)
        self.shep_cd = shep_cd
        self.be_cd = be_cd
        self.gs_cd = gs_cd
        self.default_pd = default_pd
        self.inf_pd = inf_pd

    def get_sdict(self):
        rv = super().get_sdict()
        rv['shep_cd'] = self.shep_cd
        rv['be_cd'] = self.be_cd
        rv['gs_cd'] = self.gs_cd
        rv['default_pd'] = self.default_pd
        rv['inf_pd'] = self.inf_pd
        return rv


class SHHaltTA(_MsgBase):
    """
            Refer to :ref:`definition<shhaltta>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_HALT_TA

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class SHHaltBE(_MsgBase):
    """
            Refer to :ref:`definition<shhaltbe>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_HALT_BE

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class SHHalted(_MsgBase):
    """
            Refer to :ref:`definition<shhalted>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_HALTED

    def __init__(self, tag, idx=0, _tc=None):
        super().__init__(tag)
        self.idx = idx

    def get_sdict(self):
        rv = super().get_sdict()
        rv['idx'] = self.idx
        return rv


class SHFwdInput(_MsgBase):
    """
            Refer to :ref:`definition<shfwdinput>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_FWD_INPUT

    MAX = 1024

    def __init__(self, tag, p_uid, r_c_uid, t_p_uid=None, input='', confirm=False, _tc=None):
        super().__init__(tag)
        self.p_uid = int(p_uid)
        self.r_c_uid = int(r_c_uid)
        self.t_p_uid = t_p_uid
        if len(input) > self.MAX:
            raise ValueError(f'input limited to {self.MAX} bytes')

        self.input = input
        self.confirm = confirm

    def get_sdict(self):
        rv = super().get_sdict()
        rv['p_uid'] = self.p_uid
        rv['r_c_uid'] = self.r_c_uid
        rv['t_p_uid'] = self.t_p_uid
        rv['input'] = self.input
        rv['confirm'] = self.confirm
        return rv


class SHFwdInputErr(_MsgBase):
    """
            Refer to :ref:`definition<shfwdinputerr>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.SH_FWD_INPUT_ERR

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:

    def __init__(self, tag, ref, err, idx=0, err_info='', _tc=None):
        super().__init__(tag, ref, err)
        self.err_info = err_info
        self.idx = idx

    def get_sdict(self):
        rv = super().get_sdict()
        rv['idx'] = self.idx
        if self.Errors.FAIL == self.err:
            rv['err_info'] = self.err_info

        return rv


class SHFwdOutput(_MsgBase):
    """
            Refer to :ref:`definition<shfwdoutput>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_FWD_OUTPUT

    MAX = 1024

    # TODO: contemplate interface where this is used
    @enum.unique
    class FDNum(enum.Enum):
        STDOUT = 1
        STDERR = 2

    def __init__(self, tag, p_uid, idx, fd_num, data, _tc=None, pid=-1, hostname='NONE'):
        super().__init__(tag)
        self.idx = int(idx)
        if len(data) > self.MAX:
            raise ValueError(f'output data limited to {self.MAX} bytes')

        self.data = data
        self.p_uid = int(p_uid)
        assert fd_num in {self.FDNum.STDOUT.value, self.FDNum.STDERR.value}
        self.fd_num = fd_num
        self.hostname = hostname
        self.pid = pid

    def get_sdict(self):
        rv = super().get_sdict()
        rv['idx'] = self.idx
        rv['data'] = self.data
        rv['p_uid'] = self.p_uid
        rv['fd_num'] = self.fd_num
        rv['hostname'] = self.hostname
        rv['pid'] = self.pid
        return rv

    def __str__(self):
        return f'{super().__str__()}, self.data={self.data!r}, self.p_uid={self.p_uid!r}, self.fd_num={self.fd_num!r}'


class SHDumpState(_MsgBase):
    """
            Refer to :ref:`definition<shdumpstate>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.SH_DUMP

    def __init__(self, tag, filename=None, _tc=None):
        super().__init__(tag)
        self.filename = filename

    def get_sdict(self):
        rv = super().get_sdict()
        rv['filename'] = self.filename
        return rv


class BENodeIdxSH(_MsgBase):
    """
            Refer to :ref:`definition<benodeidxsh>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.BE_NODEINDEX_SH

    def __init__(self, tag, node_idx, host_name=None, ip_addrs=None,
                 primary=None, logger_sdesc=None, _tc=None):
        super().__init__(tag)
        self.node_idx = node_idx
        self.host_name = host_name
        self.ip_addrs = ip_addrs
        self.primary = primary
        self.logger_sdesc = logger_sdesc

    @property
    def logger_sdesc(self):
        return self._logger_sdesc

    @logger_sdesc.setter
    def logger_sdesc(self, value):
        if isinstance(value, str):
            self._logger_sdesc = B64.from_str(value)
        else:
            self._logger_sdesc = value

    def get_sdict(self):
        rv = super().get_sdict()
        rv['node_idx'] = self.node_idx
        rv['host_name'] = self.host_name
        rv['ip_addrs'] = self.ip_addrs
        rv['primary'] = self.primary
        if isinstance(self._logger_sdesc, B64):
            rv['logger_sdesc'] = str(self._logger_sdesc)
        else:
            rv['logger_sdesc'] = self._logger_sdesc
        return rv


class BEPingSH(_MsgBase):
    """
            Refer to :ref:`definition<bepingsh>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.BE_PING_SH

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class BEHalted(_MsgBase):
    """
            Refer to :ref:`definition<behalted>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.BE_HALTED

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class LABroadcast(_MsgBase):
    """
            Refer to :ref:`definition<labroadcast>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.
            The *_tc* value of this message must be set to 68. The :ref:`Launcher Network Front End <launchernet>`
            has this as
            an external dependency.

    """

    _tc = MessageTypes.LA_BROADCAST

    def __init__(self, tag, data, _tc=None):
        super().__init__(tag)
        self.data = data

    def get_sdict(self):
        rv = super().get_sdict()
        rv['data'] = self.data
        return rv


class LAPassThruFB(_MsgBase):
    """
            Refer to :ref:`definition<lapassthrufb>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.LA_PASSTHRU_FB

    def __init__(self, tag, c_uid, data, _tc=None):
        super().__init__(tag)
        self.c_uid = int(c_uid)
        self.data = data

    def get_sdict(self):
        rv = super().get_sdict()
        rv['c_uid'] = self.c_uid
        rv['data'] = self.data
        return rv


class LAPassThruBF(_MsgBase):
    """
            Refer to :ref:`definition<lapassthrubf>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.LA_PASSTHRU_BF

    def __init__(self, tag, data, _tc=None):
        super().__init__(tag)
        self.data = data

    def get_sdict(self):
        rv = super().get_sdict()
        rv['data'] = self.data
        return rv


class LAServerMode(_MsgBase):
    """
            Refer to :ref:`definition<laservermode>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.LA_SERVER_MODE

    def __init__(self, tag, frontend, backend, frontend_args=None, backend_args=None, backend_env=None,
                 backend_run_dir='', backend_user_name='', backend_options=None, _tc=None):
        super().__init__(tag)

        if backend_args is None:
            self.backend_args = []
        else:
            self.backend_args = list(backend_args)

        if frontend_args is None:
            self.frontend_args = []
        else:
            self.frontend_args = list(frontend_args)

        self.frontend = frontend
        self.backend = backend
        self.frontend_args = frontend_args
        self.backend_args = backend_args
        self.backend_env = backend_env
        self.backend_run_dir = backend_run_dir
        self.backend_user_name = backend_user_name
        self.backend_options = backend_options

    def get_sdict(self):
        rv = super().get_sdict()
        rv['frontend'] = self.frontend
        rv['backend'] = self.backend
        rv['frontend_args'] = self.frontend_args
        rv['backend_args'] = self.backend_args
        rv['backend_env'] = self.backend_env
        rv['backend_run_dir'] = self.backend_run_dir
        rv['backend_user_name'] = self.backend_user_name
        rv['backend_options'] = self.backend_options
        return rv


# TODO FIXME: if messages are in this hierarchy they must follow the rules.
#   This one does not; the spec needs fixing too.
class LAServerModeExit(_MsgBase):
    """
            Refer to :ref:`definition<laservermodeexit>` and :ref:`Common Fields<cfs>` for a description
            of the message structure.

    """

    _tc = MessageTypes.LA_SERVER_MODE_EXIT

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #:
        FAIL = 1  #:

    def __init__(self, tag, ref, err, _tc=None):
        # The err field is the exit code of the
        # process exit.
        super().__init__(tag, ref, err)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class LAProcessDict(_MsgBase):
    """
            Refer to :ref:`definition<laprocessdict>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.LA_PROCESS_DICT

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class LAProcessDictResponse(_MsgBase):
    """
            Refer to :ref:`definition<laprocessdictresponse>` and :ref:`Common Fields<cfs>` for a
            description of the message structure.

    """

    _tc = MessageTypes.LA_PROCESS_DICT_RESPONSE

    @enum.unique
    class Errors(enum.Enum):
        SUCCESS = 0  #: Always succeeds

    def __init__(self, tag, ref, err, pdict=None, _tc=None):
        super().__init__(tag, ref, err)

        if pdict is None:
            pdict = {}

        self.pdict = pdict

    def get_sdict(self):
        rv = super().get_sdict()
        rv['pdict'] = self.pdict
        return rv


class LADumpState(_MsgBase):
    """
            Refer to :ref:`definition<ladumpstate>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.LA_DUMP

    def __init__(self, tag, filename=None, _tc=None):
        super().__init__(tag)
        self.filename = filename

    def get_sdict(self):
        rv = super().get_sdict()
        rv['filename'] = self.filename
        return rv


class LAChannelsInfo(_MsgBase):
    """
            Refer to :ref:`definition<lachannelsinfo>` and :ref:`Common Fields<cfs>` for a description of
            the message structure.

    """

    _tc = MessageTypes.LA_CHANNELS_INFO

    def __init__(self, tag, nodes_desc, gs_cd, num_gw_channels, port=dfacts.DEFAULT_TRANSPORT_PORT,
                 transport=str(dfacts.TransportAgentOptions.TCP), _tc=None):
        super().__init__(tag)

        self.gs_cd = gs_cd
        self.transport = dfacts.TransportAgentOptions.from_str(transport)
        self.num_gw_channels = num_gw_channels

        msg = "LAChannelsInfo nodes_desc requires list of dmsg.SHChannelsUp or already arranged dict"
        if isinstance(nodes_desc, list):
            sh_chs = all(isinstance(node_desc, SHChannelsUp) for node_desc in nodes_desc)
            if not sh_chs:
                raise TypeError(msg)
            self.nodes_desc = {node.idx: node_desc.NodeDescriptor(host_name=node.host_name,
                                                                  name=node.host_name,
                                                                  host_id=node.host_id,
                                                                  ip_addrs=node.ip_addrs,
                                                                  shep_cd=node.shep_cd,
                                                                  port=port)
                               for node in nodes_desc}
        elif isinstance(nodes_desc, dict):
            self.nodes_desc = {key: node_desc.NodeDescriptor(name=nodes_desc[key]['name'] if 'name' in nodes_desc[key] else None,
                                                             host_name=nodes_desc[key]['host_name'],
                                                             host_id=nodes_desc[key]['host_id'],
                                                             ip_addrs=nodes_desc[key]['ip_addrs'],
                                                             shep_cd=nodes_desc[key]['shep_cd'])
                               for key in nodes_desc.keys()}
        else:
            raise TypeError(msg)

    def serialize(self):
        j = {'COMPRESSED': base64.b64encode(zlib.compress(json.dumps(self.get_sdict()).encode('utf-8'))).decode('ascii')}
        return json.dumps(j)

    def get_sdict(self):
        rv = super().get_sdict()

        rv['nodes_desc'] = self.nodes_desc.copy()
        for key in self.nodes_desc.keys():
            rv['nodes_desc'][key] = self.nodes_desc[key].get_sdict()
        rv['gs_cd'] = self.gs_cd
        rv['num_gw_channels'] = self.num_gw_channels
        rv['transport'] = str(self.transport)
        return rv


class LoggingMsg(_MsgBase):
    """
            Refer to :ref:`definition<loggingmsg>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.LOGGING_MSG

    def __init__(self, tag, name, msg, time, func, hostname, ip_address, port, service, level, _tc=None):
        super().__init__(tag)
        self.name = name
        self.msg = msg
        self.time = time
        self.func = func
        self.hostname = hostname
        self.ip_address = ip_address
        self.port = port
        self.service = service
        self.level = level

    def get_sdict(self):
        rv = super().get_sdict()
        rv['name'] = self.name
        rv['msg'] = self.msg
        rv['time'] = self.time
        rv['func'] = self.func
        rv['hostname'] = self.hostname
        rv['ip_address'] = self.ip_address
        rv['port'] = self.port
        rv['service'] = self.service
        rv['level'] = self.level

        return rv

    # Extra method to cleanly generate the dictionary needed for logging
    # which requires omitting names that match attributes in LogRecord
    def get_logging_dict(self):
        rv = super().get_sdict()
        rv['time'] = self.time
        rv['hostname'] = self.hostname
        rv['ip_address'] = self.ip_address
        rv['port'] = self.port
        rv['service'] = self.service
        rv['level'] = self.level

        return rv


class LoggingMsgList(_MsgBase):
    """
            Refer to :ref:`definition<loggingmsglist>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.LOGGING_MSG_LIST

    def __init__(self, tag, records, _tc=None):
        super().__init__(tag)
        self.records = records

    @property
    def records(self):
        return self._records

    @records.setter
    def records(self, value):
        if isinstance(value, list):
            self._records = value
        elif (value, dict):
            self._records = [LoggingMsg(**v) for v in value.values()]
        else:
            msg = "LoggingMsgList records attributes requires a list or dict of LoggingMsg"
            raise AttributeError(msg)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['records'] = {i: v.get_sdict() for i, v in enumerate(self._records)}
        return rv


class LogFlushed(_MsgBase):
    """
            Refer to :ref:`definition<logflushed>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.LOG_FLUSHED

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class TAPingSH(_MsgBase):
    """
            Refer to :ref:`definition<tapingsh>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.TA_PING_SH

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class TAHalted(_MsgBase):
    """
            Refer to :ref:`definition<tahalted>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.TA_HALTED

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class TAUp(_MsgBase):
    """
            Refer to :ref:`definition<taup>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

            The test_channels are empty unless the DRAGON_TRANSPORT_TEST environment variable is
            set to some value. When set in the environment, local services will create
            two channels and provide the base64 encoded serialized channel descriptors in this
            test_channels field to the launcher front end which then disseminates them to
            the test program which is started on each node.
    """

    _tc = MessageTypes.TA_UP

    def __init__(self, tag, idx=0, test_channels=[], _tc=None):
        super().__init__(tag)
        self.idx = int(idx)
        self.test_channels = list(test_channels)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['idx'] = self.idx
        rv['test_channels'] = self.test_channels
        return rv


class Breakpoint(_MsgBase):
    _tc = MessageTypes.BREAKPOINT

    def __init__(self, tag, p_uid, index, out_desc, in_desc, _tc=None):
        super().__init__(tag)
        self.in_desc = in_desc
        self.out_desc = out_desc
        self.index = index
        self.p_uid = p_uid

    def get_sdict(self):
        rv = super().get_sdict()
        rv.update({'index': self.index,
                   'p_uid': self.p_uid,
                   'in_desc': self.in_desc,
                   'out_desc': self.out_desc})
        return rv


class BEIsUp(_MsgBase):
    """
            Refer to :ref:`definition<beisup>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.
    """

    _tc = MessageTypes.BE_IS_UP

    def __init__(self, tag, be_ch_desc, host_id, _tc=None):
        super().__init__(tag)
        self.be_ch_desc = be_ch_desc
        self.host_id = host_id

    def get_sdict(self):
        rv = super().get_sdict()
        rv['be_ch_desc'] = self.be_ch_desc
        rv['host_id'] = self.host_id
        return rv


class FENodeIdxBE(_MsgBase):
    """
            Refer to :ref:`definition<fenodeidxbe>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.
    """

    _tc = MessageTypes.FE_NODE_IDX_BE

    def __init__(self, tag, node_index, _tc=None):
        super().__init__(tag)
        self.node_index = int(node_index)

    def get_sdict(self):
        rv = super().get_sdict()
        rv['node_index'] = self.node_index
        return rv


class HaltLoggingInfra(_MsgBase):
    """
            Refer to :ref:`definition<haltlogginginfra>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.HALT_LOGGING_INFRA

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class HaltOverlay(_MsgBase):
    """
            Refer to :ref:`definition<haltoverlay>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """

    _tc = MessageTypes.HALT_OVERLAY

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class OverlayHalted(_MsgBase):
    """
            Refer to :ref:`definition<overlayhalted>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """
    _tc = MessageTypes.OVERLAY_HALTED

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class BEHaltOverlay(_MsgBase):
    """
            Refer to :ref:`definition<behaltoverlay>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """
    _tc = MessageTypes.BE_HALT_OVERLAY

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class LAHaltOverlay(_MsgBase):
    """
            Refer to :ref:`definition<lahaltoverlay>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """
    _tc = MessageTypes.LA_HALT_OVERLAY

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class OverlayPingBE(_MsgBase):
    """
            Refer to :ref:`definition<overlaypingbe>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """
    _tc = MessageTypes.OVERLAY_PING_BE

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class OverlayPingLA(_MsgBase):
    """
            Refer to :ref:`definition<overlaypingla>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """
    _tc = MessageTypes.OVERLAY_PING_LA

    def __init__(self, tag, _tc=None):
        super().__init__(tag)

    def get_sdict(self):
        rv = super().get_sdict()
        return rv


class LAExit(_MsgBase):
    """
            Refer to :ref:`definition<laexit>` and :ref:`Common Fields<cfs>` for a description of the
            message structure.

    """
    _tc = MessageTypes.LA_EXIT

    def __init__(self, tag, sigint=False, _tc=None):
        super().__init__(tag)
        self.sigint = sigint

    def get_sdict(self):
        rv = super().get_sdict()
        rv['sigint'] = self.sigint
        return rv


all_message_classes = [GSProcessCreate,
                       GSProcessCreateResponse,
                       GSProcessList,
                       GSProcessListResponse,
                       GSProcessQuery,
                       GSProcessQueryResponse,
                       GSProcessKill,
                       GSProcessKillResponse,
                       GSProcessJoin,
                       GSProcessJoinResponse,
                       GSChannelCreate,
                       GSChannelCreateResponse,
                       GSChannelList,
                       GSChannelListResponse,
                       GSChannelQuery,
                       GSChannelQueryResponse,
                       GSChannelDestroy,
                       GSChannelDestroyResponse,
                       GSChannelJoin,
                       GSChannelJoinResponse,
                       GSChannelDetach,
                       GSChannelDetachResponse,
                       GSChannelGetSendH,
                       GSChannelGetSendHResponse,
                       GSChannelGetRecvH,
                       GSChannelGetRecvHResponse,
                       AbnormalTermination,
                       GSStarted,
                       GSPingSH,
                       GSIsUp,
                       GSHeadExit,
                       GSChannelRelease,
                       GSHalted,
                       SHProcessCreate,
                       SHProcessCreateResponse,
                       SHProcessKill,
                       SHProcessExit,
                       SHChannelCreate,
                       SHChannelCreateResponse,
                       SHChannelDestroy,
                       SHChannelDestroyResponse,
                       SHLockChannel,
                       SHLockChannelResponse,
                       SHAllocMsg,
                       SHAllocMsgResponse,
                       SHAllocBlock,
                       SHAllocBlockResponse,
                       SHChannelsUp,
                       SHPingGS,
                       SHHalted,
                       SHFwdInput,
                       SHFwdInputErr,
                       SHFwdOutput,
                       SHDumpState,
                       GSTeardown,
                       SHTeardown,
                       SHPingBE,
                       BEPingSH,
                       BENodeIdxSH,
                       TAPingSH,
                       SHHaltTA,
                       TAHalted,
                       SHHaltBE,
                       BEHalted,
                       TAUp,
                       GSPingProc,
                       GSDump,
                       LABroadcast,
                       LAChannelsInfo,
                       LAPassThruBF,
                       LAPassThruFB,
                       LAServerMode,
                       LAServerModeExit,
                       LAProcessDict,
                       LAProcessDictResponse,
                       LADumpState,
                       GSPoolList,
                       GSPoolListResponse,
                       GSPoolCreate,
                       GSPoolCreateResponse,
                       GSPoolDestroy,
                       GSPoolDestroyResponse,
                       GSPoolQuery,
                       GSPoolQueryResponse,
                       SHPoolCreate,
                       SHPoolCreateResponse,
                       SHPoolDestroy,
                       SHPoolDestroyResponse,
                       SHExecMemRequest,
                       SHExecMemResponse,
                       GSUnexpected,
                       SHProcessKillResponse,
                       Breakpoint,
                       GSProcessJoinList,
                       GSProcessJoinListResponse,
                       GSNodeList,
                       GSNodeListResponse,
                       GSNodeQuery,
                       GSNodeQueryResponse,
                       GSNodeQueryTotalCPUCount,
                       GSNodeQueryTotalCPUCountResponse,
                       LoggingMsg,
                       LoggingMsgList,
                       LogFlushed,
                       BEIsUp,
                       FENodeIdxBE,
                       HaltOverlay,
                       HaltLoggingInfra,
                       OverlayPingBE,
                       OverlayPingLA,
                       LAHaltOverlay,
                       BEHaltOverlay,
                       OverlayHalted,
                       ExceptionlessAbort,
                       LAExit,
                       GSGroupCreate,
                       GSGroupCreateResponse,
                       GSGroupList,
                       GSGroupQuery,
                       GSGroupListResponse,
                       GSGroupCreateResponse,
                       GSGroupQueryResponse,
                       GSGroupKill,
                       GSGroupKillResponse,
                       GSGroupDestroy,
                       GSGroupDestroyResponse,
                       GSGroupAddTo,
                       GSGroupAddToResponse,
                       GSGroupRemoveFrom,
                       GSGroupRemoveFromResponse,
                       GSGroupCreateAddTo,
                       GSGroupCreateAddToResponse,
                       GSGroupDestroyRemoveFrom,
                       GSGroupDestroyRemoveFromResponse]

mt_dispatch = {cls._tc.value: cls for cls in all_message_classes}


def parse(jstring, restrict=None):

    try:
        sdict = json.loads(jstring)
    except TypeError as e:
        raise TypeError(f'The message "{jstring}" could not be parsed.') from e

    # if a compressed message, decompress to get the service message
    if 'COMPRESSED' in sdict:
        jstring = zlib.decompress(base64.b64decode(sdict['COMPRESSED']))
        sdict = json.loads(jstring)

    typecode = sdict['_tc']

    if restrict:
        assert typecode in restrict

    return mt_dispatch[typecode].from_sdict(sdict)


# types of all the messages that global services can receive
all_gs_messages = {GSProcessCreate, GSProcessList, GSProcessQuery,
                   GSProcessKill, GSProcessJoin, GSChannelCreate,
                   GSChannelList, GSChannelQuery, GSChannelDestroy,
                   GSChannelJoin, GSChannelDetach, GSChannelGetSendH,
                   GSChannelGetRecvH, GSChannelRelease, SHProcessCreateResponse,
                   SHProcessExit, SHChannelCreateResponse,
                   SHChannelDestroyResponse, SHLockChannelResponse,
                   SHAllocMsgResponse, SHAllocBlockResponse,
                   SHPingGS, GSTeardown, GSDump, GSPoolList, GSPoolCreate, GSPoolQuery,
                   GSPoolDestroy, SHProcessKillResponse,
                   GSPoolListResponse, GSPoolDestroyResponse, GSPoolCreateResponse,
                   GSPoolQueryResponse, GSTeardown,
                   SHPoolCreateResponse, SHPoolDestroyResponse, GSProcessJoinList,
                   GSNodeList, GSNodeQuery, GSNodeQueryTotalCPUCount,
                   GSGroupList, GSGroupCreate, GSGroupQuery,
                   GSGroupListResponse, GSGroupDestroy, GSGroupAddTo,
                   GSGroupRemoveFrom, GSGroupCreateAddTo,
                   GSGroupKill, GSGroupDestroyRemoveFrom}

src/dragon/infrastructure/messages.py
DragonHPC/dragon
import time

import dragon.channels as dch
import dragon.managed_memory as dmm


class MinConnection:
    """Minimal one way Connection style class based on channels

        DEPRECATED

        Only trying to get this to work for local channels now.

        Also, this only works for sending newline-free strings around, intended
        to be handled as infrastructure
        messages, and there is no attempt to avoid copying the data
        multiple times.
    """

    def __init__(self, the_pool, the_channel, reading: bool = False,
                 sleepy_poll_interval_sec=0.001):
        """Construct from a channel descriptor

        :param reading: is this a read-side connection
        :type reading: bool

        :param the_channel: handle to initialize stuff from
        :type the_channel: dragon.channels.Channel
        """

        assert isinstance(the_channel, dch.Channel)
        self.the_channel = the_channel

        assert isinstance(the_pool, dmm.MemoryPool)
        self.the_pool = the_pool

        self._is_read_direction = bool(reading)

        if self._is_read_direction:
            self.handle = self.the_channel.recvh()
        else:
            self.handle = self.the_channel.sendh()

        self.handle.open()
        self.opened = True

        self._stashed_msg = None

        self.sleepy_poll_interval_sec = sleepy_poll_interval_sec

    def __del__(self):
        # unclear whether detaching from or destroying the
        # underlying channel is correct.  I think it is not
        # correct to do in this object for the time being.

        try:
            self.close()
        except Exception as e:
            pass  # No penalty for attempting to close if already closed.

    def send(self, data):
        assert (not self._is_read_direction)
        assert isinstance(data, str)
        msg_bytes = data.encode().strip().replace(b'\n', b' ') + b'\n'
        the_msg = dch.Message.create_alloc(self.the_pool, len(msg_bytes))
        mmv = the_msg.bytes_memview()
        mmv[:] = msg_bytes

        sent_successfully = False
        while not sent_successfully:
            try:
                self.handle.send(the_msg)
                sent_successfully = True
            except dch.ChannelSendError as cse:
                if cse.ex_code == dch.ChannelSendEx.Errors.CH_FULL:
                    time.sleep(self.sleepy_poll_interval_sec)

        the_msg.destroy()

    @staticmethod
    def _extract(msg):
        return bytes(msg.bytes_memview()).decode()

    def recv(self):
        """Blocking read from channel.

        This blocking read does a sleepy poll on the channel.

        If poll was called ahead of time, the message may already
        be present, and is returned in that case.

        :return: The string received from the channel.
        """

        assert self._is_read_direction

        # is there something leftover from polling before?
        # if so return that

        if self._stashed_msg is not None:
            result = self._extract(self._stashed_msg)
            self._stashed_msg.destroy()
            self._stashed_msg = None
            return result

        # otherwise sleepy poll

        acquired = False
        result = None

        while not acquired:
            try:
                the_msg = self.handle.recv()
                acquired = True
                result = self._extract(the_msg)
                the_msg.destroy()
            except dch.ChannelRecvError as cre:
                if cre.ex_code == dch.ChannelRecvEx.Errors.CH_EMPTY:
                    time.sleep(self.sleepy_poll_interval_sec)
                else:
                    raise cre

        return result

    def poll(self, timeout=0):
        """Polls to see whether a message can be received.

        Note this only makes any sense if there is only one entity
        reading from a given channel, just like a pipe.  The object
        may read and cache a message from the channel and return
        this on subsequent recv calls.

        This is basically a hack to get a good consistent interface with
        other objects being used in the infrastructure between production
        and test situations.

        :param timeout: None means infinite timeout
        :type timeout: float or None
        :return: True or False according to whether there is a message there
        or the timeout has expired.

        """
        if self._stashed_msg is not None:
            return True

        if timeout is not None:
            remaining_timeout = timeout
        else:
            remaining_timeout = -1

        got_something = False

        while timeout is None or remaining_timeout >= 0:
            try:
                self._stashed_msg = self.handle.recv()
                got_something = True
                break
            except dch.ChannelRecvError as cre:
                if cre.ex_code == dch.ChannelRecvEx.Errors.CH_EMPTY:
                    time.sleep(self.sleepy_poll_interval_sec)
                else:
                    raise cre

            # imprecise but it avoids doing a time system call
            remaining_timeout = remaining_timeout - self.sleepy_poll_interval_sec

        return got_something

    def close(self):
        if self.opened:
            self.handle.close()
            self.opened = False

    @property
    def is_read(self):
        return self._is_read_direction

    @property
    def is_write(self):
        return not self._is_read_direction


class MinBConnection(MinConnection):
    def send(self, data):
        assert (not self._is_read_direction)
        if not isinstance(data, bytes):
            # Forcibly convert or die; helpful when handed a memoryview as data.
            data = bytes(data)
        msg_bytes = data # .strip().replace(b'\n', b' ') + b'\n'
        the_msg = dch.Message.create_alloc(self.the_pool, len(msg_bytes))
        mmv = the_msg.bytes_memview()
        mmv[:] = msg_bytes

        sent_successfully = False
        while not sent_successfully:
            try:
                self.handle.send(the_msg)
                sent_successfully = True
            except dch.ChannelSendError as cse:
                if cse.ex_code == dch.ChannelSendEx.Errors.CH_FULL:
                    time.sleep(self.sleepy_poll_interval_sec)

        the_msg.destroy()

    @staticmethod
    def _extract(msg):
        return bytes(msg.bytes_memview())

src/dragon/infrastructure/minconnection.py
DragonHPC/dragon
""" A class that uniquely defines properties of a managed hardware node for infrastructure communication.

This is a stub at the moment, and should be extended for dynamic hardware
allocations/elasticity aka The Cloud (TM) and more complex node properties
(accelerators, networks, encryption).
"""

import enum
import re
import os
from socket import gethostname, socket, AF_INET, SOCK_STREAM
from typing import Optional

from .facts import DEFAULT_TRANSPORT_NETIF, DEFAULT_OVERLAY_NETWORK_PORT, DEFAULT_PORT_RANGE
from ..utils import host_id as get_host_id
from .util import port_check


class NodeDescriptor:
    """Globally available description of a 'managed node'."""

    @enum.unique
    class State(enum.IntEnum):
        NULL = enum.auto()
        DISCOVERABLE = enum.auto()
        PENDING = enum.auto()
        ACTIVE = enum.auto()
        ERROR = enum.auto()

    # TODO: this is a stub (PE-42397). How do we deal with:
    #   * networks and their encryption ?
    #   * hetergeneous federated systems -> "cloud"
    #   * heterogenous networks betweens these systems -> "cloud"

    # TODO: Look for opportunities to integrate this further with LS and LA
    # during startup. See dragon.launcher.node_desc.py

    def __init__(
        self,
        state: State = State.NULL,
        h_uid: Optional[int] = None,
        name: str = "",
        ip_addrs: Optional[list[str]] = None,
        port: int = None,
        num_cpus: int = 0,
        physical_mem: int = 0,
        is_primary: bool = False,
        host_id: int = None,
        shep_cd: str = '',
        host_name: str = ''
    ):
        self.h_uid = h_uid
        self.name = name
        self.state = state
        self.num_cpus = num_cpus
        self.physical_mem = physical_mem
        self.is_primary = is_primary

        self.host_name = host_name
        self.shep_cd = shep_cd
        self.host_id = host_id

        # Not a very accurate measure since we don't know when a policy group is done,
        #  but it gives some manner of tracking for block application
        # TODO: This might be useful later when we can GC finished policy jobs
        self.num_policies = 0

        # Get a list of available CPUs on the node
        self.cpu_devices = list(os.sched_getaffinity(0))

        # Get a list of available GPUs on the node
        # TODO: Currently all I'm finding is to use a tensorflow library
        #       Maybe exec a bash command? lspci | grep ' VGA '
        # NOTE: Apparently there is a py-lspci that might do this more easily for us

        if port is not None:
            self.ip_addrs = [f'{ip_addr}:{port}' for ip_addr in ip_addrs]
        else:
            self.ip_addrs = ip_addrs

    def __str__(self):
        return f"name:{self.name}, host_id:{self.host_id} at {self.ip_addrs}, state:{self.state.name}"


    @classmethod
    def make_for_current_node(cls, name: Optional[str] = None, ip_addrs: Optional[list[str]] = None, is_primary: bool = False):
        """Create a serialized node descriptor for the node this Shepherd is running on.
        Can only be used to send to GS, as h_uid will be None in the descriptor.

        :param name: hostname of this node, defaults to None
        :type name: str, optional
        :param ip_addrs: List of the IP addresses for this node, defaults to ["127.0.0.1"]
        :type ip_addrs: list[str], optional
        :param is_primary: denote if this is the primary node running GS, defaults to False
        :type is_primary: bool, optional
        :return: serialized Node descruptor
        :rtype: dict
        """

        state = cls.State.ACTIVE

        huid = get_host_id() # will become h_uid in GS

        if name is None:
            name = f"Node-{huid}"

        if ip_addrs is None:
            ip_addrs = ["127.0.0.1"]

        num_cpus = os.cpu_count()
        physical_mem = os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES")
        host_name = gethostname()

        desc = cls(
            state=state, name=name, ip_addrs=ip_addrs, num_cpus=num_cpus,
            physical_mem=physical_mem, is_primary=is_primary, host_id=huid, host_name=host_name
        )

        return desc

    @classmethod
    def get_local_node_network_conf(cls,
                                    network_prefix: str = DEFAULT_TRANSPORT_NETIF,
                                    port_range: tuple[int, int] = (DEFAULT_OVERLAY_NETWORK_PORT, DEFAULT_OVERLAY_NETWORK_PORT+DEFAULT_PORT_RANGE)):
        """Return NodeDescriptor with IP for given network prefix, hostname, and host ID

        :param network_prefix: network prefix used to find IP address of this node. Defaults to DEFAULT_TRANSPORT_NETIF
        :type network_prefix: str, optional

        :param port_range: Port range to use for communication. Defaults to (DEFAULT_OVERLAY_NETWORK_PORT, DEFAULT_OVERLAY_NETWORK_PORT+DEFAULT_PORT_RANGE)
        :type port_range: tuple[int, int], optional

        :raises RuntimeError:  RuntimeError: Unable to find network prefix matching requested

        :return: Filled with local network info for node of execution
        :rtype: NodeDescriptor
        """
        from dragon.transport.ifaddrs import getifaddrs, InterfaceAddressFilter

        if type(port_range) is int:
            port_range = (port_range, port_range+DEFAULT_PORT_RANGE)
        hostname = gethostname()
        ifaddr_filter = InterfaceAddressFilter()
        ifaddr_filter.af_inet(inet6=False)  # Disable IPv6 for now
        ifaddr_filter.up_and_running()
        try:
            ifaddrs = list(filter(ifaddr_filter, getifaddrs()))
        except OSError:
            raise
        if not ifaddrs:
            _msg = 'No network interface with an AF_INET address was found'
            raise RuntimeError(_msg)

        try:
            re_prefix = re.compile(network_prefix)
        except (TypeError, NameError):
            _msg = "expected a string regular expression for network interface network_prefix"
            raise RuntimeError(_msg)

        ifaddr_filter.clear()
        ifaddr_filter.name_re(re.compile(re_prefix))
        ip_addrs = [ifa['addr']['addr'] for ifa in filter(ifaddr_filter, ifaddrs)]
        if not ip_addrs:
            _msg = f'No IP addresses found for {hostname} matching regex pattern: {network_prefix}'
            raise ValueError(_msg)

        ip_addr = ip_addrs[0]

        # There is a small chance (very small) that the port could appear to be available
        # and then in the next second or so, something else could grab it. In that case
        # binding to the port could fail. But this would be no different than other port
        # binding applications and on the next try it should be fine since it will then find
        # a new port. Also, predictably getting the same port on each run could be beneficial
        # for debugging. However, if randomness of port was really desired, then rand_range
        # could be used below in place of range.

        for port in range(port_range[0], port_range[1]):
            if port_check((ip_addr, port)):
                node_info = cls(state=NodeDescriptor.State.ACTIVE,
                                name=hostname,
                                host_name=hostname,
                                ip_addrs=ip_addrs,
                                host_id=get_host_id(),
                                port=port)

                return node_info

        raise RuntimeError(f'Could not find available port for IP address={ip_addr} in port range {port_range}')

    @property
    def sdesc(self):
        return self.get_sdict()

    def get_sdict(self):

        rv = {
            "state": int(self.state),
            "h_uid": self.h_uid,
            "name": self.name,
            "host_name": self.host_name,
            "is_primary": self.is_primary,
            "ip_addrs": self.ip_addrs,
            "host_id": self.host_id,
            "num_cpus": self.num_cpus,
            "physical_mem": self.physical_mem,
            "shep_cd": self.shep_cd,
        }

        return rv

    @classmethod
    def from_sdict(cls, sdict):
        sdict["state"] = NodeDescriptor.State(sdict["state"])
        return NodeDescriptor(**sdict)

src/dragon/infrastructure/node_desc.py
DragonHPC/dragon
"""Launch parameters for Dragon infrastructure and managed processes

Dragon processes that get launched need a number of parameters explaining
the size of the allocation, which node they are on, and other information
needed to bootstrap communication to the infrastructure itself.

It isn't reasonable to provide these on the command line, so they are instead
passed through environment variables managed by this module.
"""
import base64
import binascii
import os
import shlex

import dragon.infrastructure.facts as dfacts
import dragon.dtypes as dtypes
import dragon.utils as du

class ParmError(Exception):
    pass


class TypedParm:
    """Defines a parameter and how it is to be checked.

    Attributes:
        name: string name of parameter, as env variable comes out as DRAGON_name
        cast: callable to cast the string parameter to what it should be
        check: had better be True, call on output of cast
        default: default value for parameter; default/default is None
    """

    def __init__(self, name, cast, check, default=None):
        self.check = check
        self.cast = cast
        self.name = name
        self.default = default


def nocast(x):
    return x


# noinspection PyUnusedLocal
def nocheck(dummy):
    return True


def nonnegative(x):
    return x >= 0


def positive(x):
    return x > 0


def check_mode(mode):
    return mode in {dfacts.CRAYHPE_MULTI_MODE, dfacts.SINGLE_MODE, dfacts.TEST_MODE}

def check_pool(pool):
    return pool in {dfacts.PATCHED, dfacts.NATIVE}


def cast_wait_mode(wait_mode):
    if isinstance(wait_mode, dtypes.WaitMode):
        return wait_mode
    if wait_mode == "IDLE_WAIT":
        return dtypes.IDLE_WAIT
    if wait_mode == "SPIN_WAIT":
        return dtypes.SPIN_WAIT
    if wait_mode == "ADAPTIVE_WAIT":
        return dtypes.ADAPTIVE_WAIT
    if wait_mode == dtypes.IDLE_WAIT.value:
        return dtypes.IDLE_WAIT
    if wait_mode == dtypes.SPIN_WAIT.value:
        return dtypes.SPIN_WAIT
    if wait_mode == dtypes.ADAPTIVE_WAIT.value:
        return dtypes.ADAPTIVE_WAIT
    if wait_mode == str(dtypes.IDLE_WAIT.value):
        return dtypes.IDLE_WAIT
    if wait_mode == str(dtypes.SPIN_WAIT.value):
        return dtypes.SPIN_WAIT
    if wait_mode == str(dtypes.ADAPTIVE_WAIT.value):
        return dtypes.ADAPTIVE_WAIT
    if wait_mode == str(dtypes.IDLE_WAIT):
        return dtypes.IDLE_WAIT
    if wait_mode == str(dtypes.SPIN_WAIT):
        return dtypes.SPIN_WAIT
    if wait_mode == str(dtypes.ADAPTIVE_WAIT):
        return dtypes.ADAPTIVE_WAIT
    raise ValueError(f"The provided wait mode must be one of IDLE_WAIT, SPIN_WAIT, or ADAPTIVE_WAIT. The value was {wait_mode}")



def cast_return_when_mode(return_when):
    if isinstance(return_when, dtypes.ReturnWhen):
        return return_when
    if return_when == "WHEN_IMMEDIATE":
        return dtypes.WHEN_IMMEDIATE
    if return_when == "WHEN_BUFFERED":
        return dtypes.WHEN_BUFFERED
    if return_when == "WHEN_DEPOSITED":
        return dtypes.WHEN_DEPOSITED
    if return_when == "WHEN_RECEIVED":
        return dtypes.WHEN_RECEIVED
    if return_when == str(dtypes.WHEN_IMMEDIATE.value):
        return dtypes.WHEN_IMMEDIATE
    if return_when == str(dtypes.WHEN_BUFFERED.value):
        return dtypes.WHEN_BUFFERED
    if return_when == str(dtypes.WHEN_DEPOSITED.value):
        return dtypes.WHEN_DEPOSITED
    if return_when == str(dtypes.WHEN_RECEIVED.value):
        return dtypes.WHEN_RECEIVED
    if return_when == dtypes.WHEN_IMMEDIATE.value:
        return dtypes.WHEN_IMMEDIATE
    if return_when == dtypes.WHEN_BUFFERED.value:
        return dtypes.WHEN_BUFFERED
    if return_when == dtypes.WHEN_DEPOSITED.value:
        return dtypes.WHEN_DEPOSITED
    if return_when == dtypes.WHEN_RECEIVED.value:
        return dtypes.WHEN_RECEIVED
    if return_when == str(dtypes.WHEN_IMMEDIATE):
        return dtypes.WHEN_IMMEDIATE
    if return_when == str(dtypes.WHEN_BUFFERED):
        return dtypes.WHEN_BUFFERED
    if return_when == str(dtypes.WHEN_DEPOSITED):
        return dtypes.WHEN_DEPOSITED
    if return_when == str(dtypes.WHEN_RECEIVED):
        return dtypes.WHEN_RECEIVED

    raise ValueError(
        f"The provided return mode must be one of WHEN_IMMEDIATE, WHEN_BUFFERED, WHEN_DEPOSITED, or WHEN_RECEIVED. The value was {return_when}"
    )


def check_base64(strdata):
    decoded_ok = True
    try:
        base64.b64decode(strdata, validate=True)
    except binascii.Error:
        decoded_ok = False

    return decoded_ok


class LaunchParameters:
    """Launch Parameters for Dragon processes.

    This class contains the basic runtime parameters for Dragon infrastructure
    and managed processes.

    The parameter names are available in the environment of every Dragon managed
    process with the prefix `DRAGON_`. E.g. `DRAGON_MY_PUID` for `MY_PUID`. Setting
    some of these parameters will alter runtime behaviour. E.g. setting
    `DRAGON_DEBUG=1` in the environment before starting the runtime, will setup the
    runtime in debug mode and result in verbose log files.

    The parameter names are available as attributes with lower cased names in
    Python.

    The following launch parameters are defined

    :param MODE: Startup mode of the infrastructure, one of `'crayhpe-multi'`, `'single'`, `'test'`, defaults to `'test'`.
    :type MODE: str
    :param INDEX: Internal non-negative index of the current node, defaults to 0.
    :type INDEX: int
    :param DEFAULT_PD: Base64 encoded serialized descriptor of the default user memory pool, defaults to ''.
    :type DEFAULT_PD: str
    :param INF_PD: Base64 encoded serialized descriptor of the default infrastructure memory pool, defaults to ''.
    :type INF_PD: str
    :param LOCAL_SHEP_CD: Base64 encoded serialized descriptor of the channel to send to the local services on the same node, defaults to ''.
    :type LOCAL_SHEP_CD: str
    :param LOCAL_BE_CD: Base64 encoded serialized descriptor of the channel to send to the launcher backend on the same node, defaults to ''.
    :type LOCAL_BE_CD:  str
    :param GS_CD: Base64 encoded serialized descriptor of the channel to send to global services, defaults to ''.
    :type GS_CD: str
    :param GS_RET_CD: Base64 encoded serialized descriptor of the channel to receive from global services, defaults to ''.
    :type GS_RET_CD: str
    :param SHEP_RET_CD: Base64 encoded serialized descriptor of the channel to receive from the local services on the same node, defaults to ''.
    :type SHEP_RET_CD: str
    :param DEFAULT_SEG_SZ: Size of the default user managed memory pool in bytes, defaults to 2**32.
    :type DEFAULT_SEG_SZ: int
    :param INF_SEG_SZ: Size of the default infrastructure managed memory pool in bytes, defaults to 2**30.
    :type INF_SEG_SZ: int
    :param MY_PUID: Unique process ID (`puid`) given to this process by global services, defaults to 1.
    :type MY_PUID: int
    :param TEST: if in test mode, defaults to 0.
    :type TEST: int
    :param DEBUG: if in debug mode, defaults to 0.
    :type DEBUG: int
    :param BE_CUID: Unique channel ID (`cuid`) of the launcher backend on this node, defaults to 2**60.
    :type BE_CUID: int
    :param INF_WAIT_MODE: Default wait mode of infrastructure objects. One of `'IDLE_WAIT'`, `'SPIN_WAIT'`, or `'ADAPTIVE_WAIT'`. Defaults to `'IDLE_WAIT'`.
    :type INF_WAIT_MODE: str
    :param USER_WAIT_MODE: Default wait mode of user objects. One of `'IDLE_WAIT'`, `'SPIN_WAIT'`, , or `'ADAPTIVE_WAIT'`. defaults to `'ADAPTIVE_WAIT'`.
    :type USER_WAIT_MODE: str
    :param INF_RETURN_WHEN_MODE: Default return mode for infrastructure objects. One of `'WHEN_IMMEDIATE'`, `'WHEN_BUFFERED'`, `'WHEN_DEPOSITED'`, `'WHEN_RECEIVED'`, defaults to `'WHEN_BUFFERED'`.
    :type INF_RETURN_WHEN_MODE: str
    :param USER_RETURN_WHEN_MODE: Default return mode for user objects. One of `'WHEN_IMMEDIATE'`, `'WHEN_BUFFERED'`, `'WHEN_DEPOSITED'`, `'WHEN_RECEIVED'`, defaults to `'WHEN_DEPOSITED'`.
    :type USER_RETURN_WHEN_MODE: str
    :param GW_CAPACITY: Positive capacity of gateway channels, defaults to 2048.
    :type GW_CAPACITY: int
    :param PMOD_COMMUNICATION_TIMEOUT: Timeout in seconds for PMOD communication with child MPI processes, defaults to 30.
    :type PMOD_COMMUNICATION_TIMEOUT: int
    :param BASEPOOL: Default implementation of multiprocessing pool. One of `'NATIVE'` or `'PATCHED'`, defaults to `'NATIVE'`.
    :type BASEPOOL: str
    """

    @classmethod
    def init_class_vars(cls):

        PARMS = [
            TypedParm(name=dfacts.MODE, cast=str, check=check_mode, default=dfacts.TEST_MODE),
            TypedParm(name=dfacts.INDEX, cast=int, check=nonnegative, default=0),
            TypedParm(name=dfacts.DEFAULT_PD, cast=str, check=check_base64, default=""),
            TypedParm(name=dfacts.INF_PD, cast=str, check=check_base64, default=""),
            TypedParm(name=dfacts.LOCAL_SHEP_CD, cast=str, check=check_base64, default=""),
            TypedParm(name=dfacts.LOCAL_BE_CD, cast=str, check=check_base64, default=""),
            TypedParm(name=dfacts.GS_RET_CD, cast=str, check=check_base64, default=""),
            TypedParm(name=dfacts.SHEP_RET_CD, cast=str, check=check_base64, default=""),
            TypedParm(name=dfacts.GS_CD, cast=str, check=check_base64, default=""),
            TypedParm(
                name=dfacts.DEFAULT_SEG_SZ,
                cast=int,
                check=positive,
                default=int(dfacts.DEFAULT_SINGLE_DEF_SEG_SZ),
            ),
            TypedParm(
                name=dfacts.INF_SEG_SZ, cast=int, check=positive, default=int(dfacts.DEFAULT_SINGLE_INF_SEG_SZ)
            ),
            TypedParm(name=dfacts.MY_PUID, cast=int, check=positive, default=1),
            TypedParm(name=dfacts.TEST, cast=int, check=nonnegative, default=0),
            TypedParm(name=dfacts.DEBUG, cast=int, check=nonnegative, default=0),
            TypedParm(name=dfacts.BE_CUID, cast=int, check=positive, default=dfacts.BASE_BE_CUID),
            TypedParm(
                name=dfacts.INF_WAIT_MODE,
                cast=cast_wait_mode,
                check=nocheck,
                default=dfacts.INFRASTRUCTURE_DEFAULT_WAIT_MODE,
            ),
            TypedParm(
                name=dfacts.USER_WAIT_MODE,
                cast=cast_wait_mode,
                check=nocheck,
                default=dfacts.USER_DEFAULT_WAIT_MODE,
            ),
            TypedParm(
                name=dfacts.INF_RETURN_WHEN_MODE,
                cast=cast_return_when_mode,
                check=nocheck,
                default=dfacts.INFRASTRUCTURE_DEFAULT_RETURN_WHEN_MODE,
            ),
            TypedParm(
                name=dfacts.USER_RETURN_WHEN_MODE,
                cast=cast_return_when_mode,
                check=nocheck,
                default=dfacts.USER_DEFAULT_RETURN_WHEN_MODE,
            ),
            TypedParm(name=dfacts.GW_CAPACITY, cast=int, check=positive, default=dfacts.GW_DEFAULT_CAPACITY),
            TypedParm(name=dfacts.NUM_GW_CHANNELS_PER_NODE, cast=int, check=nonnegative, default=dfacts.DRAGON_DEFAULT_NUM_GW_CHANNELS_PER_NODE),
            TypedParm(name=dfacts.HSTA_MAX_EJECTION_MB, cast=int, check=positive, default=8),
            TypedParm(name=dfacts.HSTA_MAX_GETMSG_MB, cast=int, check=positive, default=8),
            TypedParm(name=dfacts.PMOD_COMMUNICATION_TIMEOUT, cast=int, check=positive, default=30),
            TypedParm(name=dfacts.BASEPOOL, cast=str, check=check_pool, default=dfacts.DEFAULT_POOL),
        ]

        env = os.environ

        try:
            num_gateways_env_var = dfacts.PREFIX + dfacts.NUM_GW_CHANNELS_PER_NODE_VAR
            num_gateways = int(env[num_gateways_env_var]) if num_gateways_env_var in env else dfacts.DRAGON_DEFAULT_NUM_GW_CHANNELS_PER_NODE
            assert(num_gateways >= 0)
        except:
            raise ValueError(f'The value of the environment variable {dfacts.NUM_GW_CHANNELS_PER_NODE} must be valid integer >= 1. It was {env[dfacts.NUM_GW_CHANNELS_PER_NODE]}')

        cls.gw_env_vars = frozenset([f'GW{x+1}' for x in range(num_gateways)])

        for gw_env_var in cls.gw_env_vars:
            PARMS.append(TypedParm(gw_env_var, cast=str, check=nocheck, default=""))

        cls.PARMS = PARMS

        cls.NODE_LOCAL_PARAMS = \
            frozenset([dfacts.DEFAULT_PD, dfacts.INF_PD, dfacts.LOCAL_SHEP_CD, dfacts.LOCAL_BE_CD, dfacts.BE_CUID]) | \
            cls.gw_env_vars


    def __init__(self, **kwargs):
        for pt in LaunchParameters.PARMS:
            attr_name = pt.name.lower()
            if attr_name in kwargs:
                setattr(self, attr_name, kwargs[attr_name])
            else:
                setattr(self, attr_name, None)

    @classmethod
    def _ev_name_from_pt(cls, ptname: str):
        if ptname in cls.gw_env_vars:
            return dfacts.PREFIX + ptname

        return dfacts.env_name(ptname)

    @staticmethod
    def _attr_name_from_pt(ptname):
        return ptname.lower()

    @classmethod
    def all_parm_names(cls):
        return {cls._attr_name_from_pt(pt.name) for pt in cls.PARMS}

    @classmethod
    def remove_node_local_evars(cls, env):
        for e in cls.NODE_LOCAL_PARAMS:
            ev_name = cls._ev_name_from_pt(e)
            if ev_name in env:
                del env[ev_name]

    @staticmethod
    def from_env(init_env=None):
        if init_env is None:
            init_env = os.environ

        evd = {}

        for pt in LaunchParameters.PARMS:
            ev_name = LaunchParameters._ev_name_from_pt(pt.name)
            attr_name = LaunchParameters._attr_name_from_pt(pt.name)
            if ev_name in init_env:
                try:
                    val = pt.cast(init_env[ev_name])
                except Exception as ex:
                    raise ValueError(
                        f"The value of environment variable {ev_name} was not the correct type. It was {init_env[ev_name]}. The error was: {repr(ex)}"
                    )

                if not pt.check(val):
                    msg = f"{ev_name} = {val} failed check {pt.check.__name__}"
                    raise ParmError(msg)

                evd[attr_name] = val
            else:
                evd[attr_name] = pt.default

        return LaunchParameters(**evd)

    def env(self):
        """Returns dictionary of environment variables for this object.

        An entry is created for parameters that are not None and are listed in
        LaunchParameters.Parms, with the correctly prefixed name of the
        environment variable and so on.

        The reason for this option is to let a process or actor construct
        the launch parameters it wants to have for another process
        it is getting ready to start.
        """
        rv = {}
        for pt in LaunchParameters.PARMS:
            attr_name = LaunchParameters._attr_name_from_pt(pt.name)
            val = getattr(self, attr_name, None)

            if val is not None:
                ev_name = LaunchParameters._ev_name_from_pt(pt.name)
                rv[ev_name] = str(val)

        return rv

    def set_num_gateways_per_node(self, num_gateways=dfacts.DRAGON_DEFAULT_NUM_GW_CHANNELS_PER_NODE):
        # Configure the number of gateways per node as an
        # environment variable.
        num_gateways_env_var = dfacts.PREFIX+dfacts.NUM_GW_CHANNELS_PER_NODE_VAR
        assert(num_gateways >= 0), f'The number of gateways must be greater or equal to 0. It was {num_gateways}'
        os.environ[num_gateways_env_var] = str(num_gateways)
        this_process.num_gw_channels_per_node = num_gateways
        self.init_class_vars()

LaunchParameters.init_class_vars()
this_process = LaunchParameters.from_env()


class Policy:
    """Used to encapsulate policy decisions.

    An instance of this class can be used to configure the policy decisions
    made within code as shown. Right now it includes the wait mode. But it
    is intended to include other configurable policy decisions in the future
    like affinity.

    There are two constant instances of Policy that are created as
    POLICY_INFRASTRUCTURE and POLICY_USER as defaults for both
    infrastructure and user program resources.

    """

    def __init__(self, wait_mode, return_from_send_when):
        self._wait_mode = wait_mode
        self._return_from_send_when = return_from_send_when

    def __getstate__(self):
        return (self._wait_mode, self._return_from_send_when)

    def __setstate__(self, state):
        (wait_mode, return_from_send_when) = state
        self._wait_mode = wait_mode
        self._return_from_send_when = return_from_send_when

    def __repr__(self):
        return f"{self.__class__.__name__}({self._wait_mode, self._return_from_send_when})"

    def __str__(self):
        return repr(self)

    @property
    def wait_mode(self):
        return self._wait_mode

    @property
    def return_when(self):
        return self._return_from_send_when


POLICY_INFRASTRUCTURE = Policy(this_process.inf_wait_mode, this_process.inf_return_when_mode)
POLICY_USER = Policy(this_process.user_wait_mode, this_process.user_return_when_mode)

src/dragon/infrastructure/parameters.py
DragonHPC/dragon
import enum
from dataclasses import asdict, dataclass, field

@dataclass
class Policy:
    """
    Dataclass to describe what node, affinity, and pattern to distrbiute a Policy to
    """

    class Placement(enum.IntEnum):
        """
        Which node to assign a policy to.

        Local and Anywhere will be useful later for multi-system communication
        Right now Placement will have little effect unless HOST_NAME or HOST_ID
        are used, which will try to place a policy on the specified node
        """
        LOCAL = -5 # Local to the current system of nodes
        ANYWHERE = -4 # What it says on the tin
        HOST_NAME = -3 # Tells evaluator to check host_name in policy and apply to that node
        HOST_ID = -2 # Tells evaluator to check host_id in policy and apply to that node
        DEFAULT = -1 # For now, defaults to ANYWHERE

    class Distribution(enum.IntEnum):
        """
        Pattern to use to distribute policies across nodes
        """
        ROUNDROBIN = enum.auto()
        BLOCK = enum.auto()
        DEFAULT = enum.auto()

    class Device(enum.IntEnum):
        """
        Which type of device the affinity policy will apply to
        """
        CPU = enum.auto()
        GPU = enum.auto()
        DEFAULT = enum.auto()

    # TODO: The NodeDescriptor may need to hold its devices in a dictionary so we can count how many devices have been used to do some kind of rudamentary load-balancing
    class Affinity(enum.IntEnum):
        """
        Device Affinity distribution
        """
        ANY = -5
        # ROUNDROBIN = -4 # TODO: Not implemented
        # BLOCK = -3 # TODO: Not implemented
        SPECIFIC = -2
        DEFAULT = -1 # Same as ANY

    # TODO: Not implemented
    class WaitMode(enum.IntEnum):
        """Channel WaitMode type"""
        IDLE = enum.auto()
        SPIN = enum.auto()
        DEFAULT = enum.auto()

    placement : Placement = Placement.DEFAULT
    host_name : str = "" # To be populated by caller in use with Placement.HOST_NAME
    host_id : int = -1 # To be populated by caller in use with Placement.HOST_ID
    distribution : Distribution = Distribution.DEFAULT
    device : Device = Device.DEFAULT
    affinity : Affinity = Affinity.DEFAULT
    specific_affinity : list[int] = field(default_factory=list) # To be populated by called in use with Affinity.SPECIFIC
    wait_mode : WaitMode = WaitMode.DEFAULT # TODO (For channels)
    refcounted : bool = True # TODO (Apply refcounting for this resource)

    @property
    def sdesc(self):
        return self.get_sdict()

    def get_sdict(self):
        rv = {
            "placement" : self.placement,
            "host_name" : self.host_name,
            "host_id" : self.host_id,
            "distribution" : self.distribution,
            "device" : self.device,
            "affinity" : self.affinity,
            "specific_affinity" : self.specific_affinity,
            "wait_mode" : self.wait_mode,
            "refcounted" : self.refcounted,
        }

        return rv

    @classmethod
    def from_sdict(cls, sdict):
        return Policy(**sdict)


# Default policy to merge incoming policies against
GS_DEFAULT_POLICY = Policy(
    placement=Policy.Placement.ANYWHERE,
    host_name="",
    host_id=-1,
    distribution=Policy.Distribution.ROUNDROBIN,
    device=Policy.Device.CPU,
    affinity=Policy.Affinity.ANY,
    specific_affinity=[],
    wait_mode=Policy.WaitMode.IDLE,
    refcounted=True
)
src/dragon/infrastructure/policy.py
DragonHPC/dragon
"""A class that uniquely defines properties of a managed memory pool for infrastructure communication.
"""
import enum
from ..infrastructure import parameters as dparm
from .. import utils as du


class PoolDescriptor:
    """Globally available descriptior of a pool.

    Attributes:

        m_uid : nonnegative int
            unique memory pool id

        name : str
            string indicating pool name that servers as an identifier
            to query for the pool

        sdesc : bytes-like object
            serialized library descriptor

        node : nonnegative int
            which node this memory pool is associated with

        state : cls.State enum member
            current state of this pool
    """

    @enum.unique
    class State(enum.Enum):
        """Current state of the pool."""
        PENDING = enum.auto()
        ACTIVE = enum.auto()
        DEAD = enum.auto()

    def __init__(self, m_uid=0, name='', node=0, state=None, sdesc=''):
        """
            sdesc - serialized library descriptor
        """
        self.m_uid = int(m_uid)
        self.name = name
        self.sdesc = sdesc
        self.node = int(node)

        if state is None:
            self.state = self.State.PENDING
        else:
            if isinstance(state, self.State):
                self.state = state
            elif isinstance(state, int):
                self.state = self.State(state)
            else:
                raise NotImplementedError('unknown state init')

    @property
    def sdesc(self):
        """Returns the serialized library descriptor of the pool."""
        return self._sdesc

    @sdesc.setter
    def sdesc(self, value):
        """Populate the serialized library descriptor of the pool.

        :param value: value to be populated as the serialzied descriptor
        :type value: either bytes (or) string passed through the environment
        """
        if isinstance(value, str):
            self._sdesc = du.B64.str_to_bytes(value)
        else:
            self._sdesc = value

    def get_sdict(self):
        """Collect the entire information of the pool.

        :return: A dictionary with all key-value pairs of the available description of the pool.
        :rtype: Dictionary
        """
        rv = {'sdesc': du.B64.bytes_to_str(self.sdesc),
              'name': self.name,
              'm_uid': self.m_uid,
              'node': self.node,
              'state': self.state.value}
        return rv

    def __str__(self):
        return f'{self.m_uid}:{self.name} on {self.node}: {self.state.name}'

    @classmethod
    def from_sdict(cls, sdict):
        """Returns the descriptor of the memory pool."""
        return PoolDescriptor(**sdict)


class PoolOptions:
    """Options object for how a Pool is supposed to be created.
    """
    class Placement(enum.Enum):
        """Enum for placement strategy of the pool for its creation."""
        LOCAL = 0
        DEFAULT = 1
        SPECIFIC = 2

    def __init__(self, placement=0, target_node=0, sattr=''):
        """
            sattr - serialized library attributes

            placement - enum for placement strategy

            target_node - ignored if placement request isn't for a specific node
        """
        self.placement = self.Placement(placement)
        self.target_node = int(target_node)
        self.sattr = sattr

    def get_sdict(self):
        """Collect the different attributes/options of the memory pool.

        :return: A dictionary with key-value pairs of the pool creation options.
        :rtype: Dictionary
        """
        rv = {'sattr': self.sattr,
              'target_node': self.target_node,
              'placement': self.placement.value}

        return rv

    @staticmethod
    def from_sdict(sdict):
        """Returns the options of the memory pool."""
        return PoolOptions(**sdict)

src/dragon/infrastructure/pool_desc.py
DragonHPC/dragon
""" A class that uniquely defines properties of a managed Process for infrastructure communication.
"""
import enum

@enum.unique
class ArgMode(enum.Enum):
    """
        Enum for mode of delivering the arguments, either process will be using the channel (or)
        to deliver the argdata directly
    """
    NONE = 0  # no args
    PYTHON_IMMEDIATE = 1  # args are found in GSPingProc argdata members in pickled form
    PYTHON_CHANNEL = 2  # argdata in GSPingProc is a recipe for a channel descriptor, attach and read from it


def mk_argmode_from_default(mode):
    """Populate and return the mode of arguments as per the ArgMode

    :return: Mode of delivering the arguments for the launching process
    :rtype: ArgMode Enum
    """
    if isinstance(mode, ArgMode):
        mode = mode
    elif mode is None:
        mode = ArgMode.NONE
    elif isinstance(mode, int):
        mode = ArgMode(mode)
    else:
        raise NotImplementedError('unknown mode initializer')

    return mode


class ProcessDescriptor:
    """Globally available description of a managed process.

    Attributes:

        ecode : int or None
            exit code of the process, or None if still running

        node : nonnegative int
            which node of the allocation this process is/was on

        name : str
            user assigned unique name of the process

        p_uid : nonnegative int
            unique process id

        p_p_uid : nonnegative int
            parent process unique id

        live_children : set of nonnegative int

        gs_ret_cuid : nonnegative int
            c_uid of the global services return channel

        shep_ret_cuid : nonnegative int
            c_uid of the local shepherd return channel (currently not used)

        state : ProcessDescriptor.State enum member
            current run state of this process

        stdin_sdesc : a serialized channel where the stdin of the process can be provided.

        stdout_sdesc : a serialized channel where the stdout of the process is written.

        stderr_sdesc : a serialized channel where the stderr of the process is written.

    """

    @enum.unique
    class State(enum.Enum):
        """Current state of the process."""
        PENDING = enum.auto()
        ACTIVE = enum.auto()
        DEAD = enum.auto()

    def __init__(self, *, p_uid, p_p_uid, name, node, live_children=None,
                 gs_ret_cuid=None, shep_ret_cuid=None, state=None, ecode=None,
                 stdin_sdesc=None, stdout_sdesc=None, stderr_sdesc=None):
        self.p_p_uid = p_p_uid
        self.ecode = ecode
        self.node = int(node)
        self.name = name
        self.p_uid = int(p_uid)
        self.gs_ret_cuid = gs_ret_cuid
        self.shep_ret_cuid = shep_ret_cuid
        self.stdin_sdesc = stdin_sdesc
        self.stdout_sdesc = stdout_sdesc
        self.stderr_sdesc = stderr_sdesc
        self.stdin_conn = None
        self.stdout_conn = None
        self.stderr_conn = None


        if live_children is None:
            self.live_children = set()
        else:
            self.live_children = set(live_children)

        if state is None:
            self.state = self.State.PENDING
        else:
            if isinstance(state, self.State):
                self.state = state
            elif isinstance(state, int):
                self.state = self.State(state)
            else:
                raise NotImplementedError('unknown state init')

    def __str__(self):
        rv = f'{self.p_uid}:{self.name} on {self.node}: {self.state.name}'

        if self.state is self.State.DEAD:
            rv += f' exit {self.ecode}'

        return rv

    def get_sdict(self):
        """Collect the entire information of the process.

        :return: A dictionary with all key-value pairs of the available information of the process.
        :rtype: Dictionary
        """
        rv = {'node': self.node,
              'name': self.name,
              'p_uid': self.p_uid,
              'p_p_uid': self.p_p_uid,
              'live_children': list(self.live_children),
              'state': self.state.value}

        if self.state == self.State.DEAD:
            rv['ecode'] = self.ecode

        if self.gs_ret_cuid is not None:
            rv['gs_ret_cuid'] = self.gs_ret_cuid

        if self.shep_ret_cuid is not None:
            rv['shep_ret_cuid'] = self.shep_ret_cuid

        rv['stdin_sdesc'] = self.stdin_sdesc
        rv['stdout_sdesc'] = self.stdout_sdesc
        rv['stderr_sdesc'] = self.stderr_sdesc

        return rv

    @classmethod
    def from_sdict(cls, sdict):
        """Returns the descriptor of the launched process."""
        return ProcessDescriptor(**sdict)

# TODO: document what these options actually mean !
class ProcessOptions:
    """Open type, placeholder, for options on how to launch a process

    Attributes:

    argdata
        argument data, bytes to deliver to the launching process

    mode
        mode of delivering the arguments, either process will be using the channel (or)
        to deliver the argdata directly

    make_inf_channels
        bool, default False, whether process should start making the infrastructure channels
        when they are needed, or just start the process directly.
    """

    def __init__(self, make_inf_channels=False, mode=None, argdata=None):
        self.argdata = argdata
        self.mode = mk_argmode_from_default(mode)
        self.make_inf_channels = make_inf_channels

    def get_sdict(self):
        """Collect the different attributes/options of launching the process.

        :return: A dictionary with key-value pairs of the process options.
        :rtype: Dictionary
        """
        return {'mode': self.mode.value,
                'argdata': self.argdata,
                'make_inf_channels': self.make_inf_channels}

    @staticmethod
    def from_sdict(sdict):
        """Returns the options of the launched process."""
        return ProcessOptions(**sdict)

src/dragon/infrastructure/process_desc.py
DragonHPC/dragon
"""A quick and dirty scheme for doing single node channels by themselves.


"""
import atexit
import logging
import os
import subprocess

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.infrastructure.connection as dic
import dragon.infrastructure.parameters as dparm
import dragon.utils as du

_STANDALONE_POOL_EV_NAME = 'DRAGON_STANDALONE_POOL'
_STANDALONE_POOL = None
_MY_CHANNELS = set()
_IS_INIT = False
_ID_BASE = None

LOG = logging.getLogger('connection')


def destroy_standalone():
    global _STANDALONE_POOL
    _STANDALONE_POOL.destroy()
    _STANDALONE_POOL = None
    os.environ.pop(_STANDALONE_POOL_EV_NAME)


def destroy_my_channels():
    global _MY_CHANNELS

    for chan in _MY_CHANNELS:
        try:
            chan.destroy()
        except dch.ChannelError as ce:
            logging.exception('destroy of {} failed: {}'.format(chan.cuid, ce))

    _MY_CHANNELS.clear()


def get_going(pool_size=2 ** 30, pool_attr=None, pool_uid=None):
    global _STANDALONE_POOL
    if _STANDALONE_POOL_EV_NAME in os.environ:
        # get the pool descriptor from the environment
        pool_desc = du.B64.str_to_bytes(os.environ[_STANDALONE_POOL_EV_NAME])
        _STANDALONE_POOL = dmm.MemoryPool.attach(pool_desc)
    else:
        # make the pool descriptor
        username = os.environ.get('USER', str(os.getuid()))
        pool_name = f'sa_{username}_{os.getpid()!s}'
        if pool_uid is None:
            pool_uid = os.getpid()

        _STANDALONE_POOL = dmm.MemoryPool(pool_size, pool_name, pool_uid)
        os.environ[_STANDALONE_POOL_EV_NAME] = du.B64.bytes_to_str(_STANDALONE_POOL.serialize())

        # whoever was the first to make the standalone pool has the
        # responsibility for cleaning up the pool.
        # here, it just cleans it up; slightly kinder might be
        # to go into a polling loop to verify nothing is in the pool.
        # But: for use in testbenches this should probably be ok.
        atexit.register(destroy_standalone)

    # fake up some uniqueness for c_uid we make from this process.
    # the PID is probably not the very worst choice.
    global _ID_BASE
    _ID_BASE = os.getpid() * (2**20)
    atexit.register(destroy_my_channels)

    global _IS_INIT
    _IS_INIT = True


def get_new_id():
    global _ID_BASE
    tmp = _ID_BASE
    _ID_BASE += 1
    return tmp


def Pipe(duplex=True):
    global _IS_INIT
    global _MY_CHANNELS

    if not _IS_INIT:
        get_going()

    first_chan = dch.Channel(_STANDALONE_POOL, get_new_id())
    _MY_CHANNELS.add(first_chan)

    if duplex:
        second_chan = dch.Channel(_STANDALONE_POOL, get_new_id())
        _MY_CHANNELS.add(second_chan)

        first = dic.Connection(inbound_initializer=first_chan, outbound_initializer=second_chan)
        second = dic.Connection(inbound_initializer=second_chan, outbound_initializer=first_chan)
    else:
        first = dic.Connection(inbound_initializer=first_chan)
        second = dic.Connection(outbound_initializer=first_chan)

    return first, second

src/dragon/infrastructure/standalone_conn.py
DragonHPC/dragon
"""Python infrastructure utilities.

This file is where to put useful internal utilities and wrappers and stuff.

It's also to solve the problem of 'I don't know what the right long term way is
to do something, but I need a version of it now to make progress'.  The
initial solution here makes all the uses consistent and we can refactor more
easily if it needs to change.

Rules of engagement:

Anyone can put anything into this file.

Nothing in here is user facing.

Once there are a few things doing similar kinds of things in this file,
they need to be broken out into another file.
"""
from collections.abc import Iterable
import fcntl
import heapq
import io
from itertools import chain
import os
import re
import selectors
import sys
import time
import logging
from warnings import warn
from .parameters import this_process
from socket import socket, AF_INET, SOCK_STREAM
from .facts import FIRST_PUID


def _get_access_modes(f):
    """Returns (readable, writeable) pair of bools corresponding to the access
    mode of the specified file-like object (or file descriptor).
    """
    mode = fcntl.fcntl(f, fcntl.F_GETFL) & os.O_ACCMODE
    readable = writeable = mode == os.O_RDWR
    return (readable or mode == os.O_RDONLY, writeable or mode == os.O_WRONLY)


class NewlineStreamWrapper:
    """Helper class for sending json objects through streams.

        This class sends newline-delimited newline-free strings
        through streams with an interface similar to multiprocessing.Connection.

        It is used for sending around JSON-encoded infrastructure
        messages through streams and in testing.

        TODO: consider whether this belongs in base code.  It feels more
        like something that should only be in a test bench.
    """

    def __init__(self, stream, read_intent=None, write_intent=None):
        """
            read_intent and write_intent are used to ensure
            clients are using the right operations on the stream.
        """
        try:
            readable, writeable = _get_access_modes(stream)
        except (TypeError, io.UnsupportedOperation):
            # Fallback to original defaults instead of raising the ValueError
            # most likely to be raised by sel.register() below.
            readable, writeable = True, True
        self.read_intent = bool(readable if read_intent is None else read_intent)
        self.write_intent = bool(writeable if write_intent is None else write_intent)
        if self.read_intent and not readable:
            warn(f'Intend to read a stream not open for reading: {stream}')
        if self.write_intent and not writeable:
            warn(f'Intend to write a stream not open for writing: {stream}')
        self.stream = stream

    def send(self, data):
        """Perform the write operation of the given data into the stream.
        Assert that the write_intent is set, ensuring the client is using the
        right operation on the stream.

        :param data: data to be written into the stream
        :type data: string
        """
        try:
            assert self.write_intent, 'sending to a read wrap'
            assert isinstance(data, str)
            msg_str = data.strip().replace('\n', ' ') + '\n'

            # try to tolerate if the stream is a text mode stream
            if isinstance(self.stream, io.TextIOBase):
                self.stream.write(msg_str)
            else:
                self.stream.write(msg_str.encode())

            self.stream.flush()
        except ValueError:
            # If we get a ValueError exception the underlying
            # stream is closed. This means things are going away
            # anyway. So nothing to do in this case.
            pass

    def recv(self):
        """Perform the read operation on the stream. Assert that the read_intent
        is set, ensuring the client is using the right operation on the stream.
        """
        assert self.read_intent, 'receiving from a write wrap'
        line = self.stream.readline()

        # try to tolerate if the stream is a text stream
        if isinstance(line, str):
            stuff = line
        else:
            stuff = line.decode()

        return stuff

    def poll(self, timeout=0):
        """ Poll for the read I/O events on the registered file objects of type stream.
        Collect the selector if already set and its associated events, else the default
        selector and register with the read event, to wait upon. Also, assert that the
        ride_intent is set, ensuring the client is using the right operation on the stream.

        :param timeout: If timeout > 0, specifies the maximum wait time in seconds
        :type timeout: int, defaults to 0
        :return: True if read I/O event is ready on the selected stream
        :rtype: boolean
        """
        assert self.read_intent, "requires read intent"

        # To support SSH launch, see if there's a poll method provided by stream. The
        # ssh launch does provide one:
        try:
            return self.stream.poll(timeout=timeout)
        except AttributeError:
            pass


        try:
            sel = self.selector
        except AttributeError:
            sel = self.selector = selectors.DefaultSelector()
            sel.register(self.stream, selectors.EVENT_READ)
        events = sel.select(timeout)
        for key, mask in events:
            if key.fileobj is not self.stream:
                continue
            if mask & selectors.EVENT_READ:
                return True
        return False

    def close(self):
        """Close thes selector and the associated stream"""
        try:
            sel = self.selector
        except AttributeError:
            pass
        else:
            sel.close()
        self.stream.close()


class PriorityMultiMap:
    """Class combining a multimap with a priority queue to allow timeouts.

        Correct behavior relies on keys never being reused.

        This is used in Global Services to manage requests for timeout
        notification from multiple sources.

        TODO: this needs some unit tests
        TODO: consider reimplementing in terms of a collections.Counter object!
    """

    def __init__(self):
        self.internal_items = dict()
        self.dead_keys = set()
        self.timeout_keys = set()
        self.timeout_pq = list()
        self.deadlines = dict() # key: (key,value), value: deadline

    def put(self, key, value, timeout=None):
        """Store (key, value) pair to the map (internal items), used by the global requests
        to manage timeout requests with the associated deadlines of multiple sources.
        Can be used for all channels, pools, and processes.

        :param key: identifier of the requested source(typically name)
        :type key: string
        :param value: value of the identifier(typically id)
        :type value: int
        :param timeout: used to evaluate the deadline of the (key, value) pair, if provided
        :type timeout: int, defaults to Nonde
        """
        if key in self.internal_items:
            self.internal_items[key].add(value)
        else:
            self.internal_items[key] = {value}

        if timeout is not None:
            deadline = time.time() + timeout
            self.deadlines[(key, value)] = deadline
            heapq.heappush(self.timeout_pq, (deadline, key, value))
            self.timeout_keys.add(key)

            if key in self.dead_keys:
                self.dead_keys.remove(key)

    def time_left(self, key, value):
        """Provides the deadline (if stored) for the given (key, value) pair.

        :param key: identifier of the source
        :type key: string
        :param value: value of the identifier
        :type value: int
        :return: timeout/deadline of the given (key, value) pair
        :rtype: int
        """
        if (key, value) in self.deadlines:
            return self.deadlines[(key, value)]

    def get(self, key):
        """Provides the value for the stored key in the map

        :param key: identifier of the source
        :type key: string
        :return: value of the identifier
        :rtype: int
        """
        return self.internal_items[key]

    def next_deadline(self):
        """Returns the timeout/deadline of the front item in the timeout priority queue, if present

        :return: value of the deadline
        :rtype: int
        """
        if self.timeout_pq:
            return self.timeout_pq[0][0]
        else:
            return None

    def get_timed_out(self):
        """Remove all the requests in the timeout priority queue, taking the opportunity to
        clear the associated deadlines and the dead keys

        :return: List of all the (key, value) tuples of the timed out requests
        :rtype: List
        """
        rv = []
        now = time.time()
        while self.timeout_pq and self.timeout_pq[0][0] <= now:
            key = self.timeout_pq[0][1]
            if key not in self.dead_keys:
                rv.append(self.timeout_pq[0][1:])

            if self.deadlines and (self.timeout_pq[0][1], self.timeout_pq[0][2]) in self.deadlines:
                del self.deadlines[(self.timeout_pq[0][1], self.timeout_pq[0][2])]

            heapq.heappop(self.timeout_pq)

        # take opportunity to flush dead keys
        if not self.timeout_pq:
            self.dead_keys.clear()

        return rv

    def remove(self, k):
        """Removes the source from the map, with all the values of it placed in the queue,
        along with the timeout/deadline management of the source

        :param k: identifier of the object/source(typically name)
        :type k: string
        """
        if self.deadlines:
            if k in self.internal_items and self.internal_items[k]:
                for v in self.internal_items[k]:
                    if (k, v) in self.deadlines:
                        del self.deadlines[(k, v)]

        if k in self.internal_items:
            del self.internal_items[k]
        if k in self.timeout_keys:
            self.dead_keys.add(k)
            self.timeout_keys.remove(k)

    def remove_one(self, k, v):
        """Removes the request of a particular source placed in the queue,
        along with removing the timeout/deadline management of the associated source, if set.

        :param k: identifier of the source(typically name)
        :type k: string
        :param v: value of the identifier(typically id)
        :type v: int
        """
        if v in self.internal_items[k]:
            self.internal_items[k].remove(v)

            if not self.internal_items[k]:
                self.remove(k)

        if self.deadlines:
            if (k, v) in self.deadlines:
                del self.deadlines[(k, v)]

    def keys(self):
        """Helps with all the keys of the global requests placed in the priority map

        :return: List of all the keys of the queue items
        :rtype: List
        """
        return self.internal_items.keys()

    def values(self):
        """Helps with all the values of the global requests placed in the priority map

        :return: List of all the values of the queue items
        :rtype: List
        """
        return self.internal_items.values()

    def items(self):
        """Helps with all the items(key, value) pairs of the
            global requests placed in the priority map

        :return: List of all the tuples(key, value) of the queue items
        :rtype: List
        """
        return self.internal_items.items()

    def __contains__(self, key):
        return key in self.internal_items


# TODO: this needs a better name
class AbsorbingChannel:
    """Object that absorbs all sends, remembering the last one; useful in GS"""

    def __init__(self):
        self.msg = None

    def send(self, msg):
        self.msg = msg


def survey_dev_shm():
    """Looks at what is in /dev/shm owned by current user

    :return: set of filenames owned by current user in /dev/shm
    """

    my_id = os.getuid()
    with os.scandir('/dev/shm') as it:
        # during a certain shutdown race condition, FileNotFoundError may be raised.
        try:
            rv = {entry.name for entry in it if entry.stat().st_uid == my_id}
            return rv
        except FileNotFoundError:
            pass

def compare_dev_shm(previous):
    """Warns if there are files owned by current user in /dev/shm not previously seen

    :param previous: set of previously seen filenames
    :return: None
    """

    now = survey_dev_shm()

    # Due to a race condition, it is possible that None is returned.
    if now is None:
        return

    remaining = now - previous

    if remaining:
        print(f'warning: {len(remaining)} leftover files in /dev/shm: ')

    for fn in remaining:
        print(fn)


def route(msg_type, routing_table, metadata=None):
    """ Decorator routing adapter.

    This is a function decorator used to accumulate handlers for a particular
    kind of message into a routing table indexed by type, used
    by server classes processing infrastructure messages.

    The routing table is usually a class attribute and the message type
    is typically an infrastructure message.  The metadata, if used,
    is typically information used to check that the state the server object
    is in, is appropriate for that type of message.

    :param msg_type: the type of an infrastructure message class
    :param routing_table: dict, indexed by msg_type with values a tuple (function, metadata)
    :param metadata: metadata to check before function is called, usage dependent
    :return: the function.
    """

    def decorator_route(f):
        assert msg_type not in routing_table
        routing_table[msg_type] = (f, metadata)
        return f

    return decorator_route


# TODO: PE-37739, the full on client debug hookup story.
def mk_fifo_debugger(basename, *, override_bph=False, quiet=False):
    dbg_in_fn = basename + '_dbg_in'
    dbg_out_fn = basename + '_dbg_out'
    os.mkfifo(dbg_in_fn)
    os.mkfifo(dbg_out_fn)
    if not quiet:
        print(f'fifos made: in: {dbg_in_fn} out: {dbg_out_fn}; waiting on open')
    sys.stdout.flush()
    in_fifo = open(dbg_in_fn, 'r')
    out_fifo = open(dbg_out_fn, 'w')
    import pdb
    debugger = pdb.Pdb(stdin=in_fifo, stdout=out_fifo)

    if override_bph:
        sys.breakpointhook = debugger.set_trace

    return pdb.Pdb(stdin=in_fifo, stdout=out_fifo)


def to_str(x):
    """Convert anything to a string"""
    if x is None:
        return ''
    if hasattr(x, 'decode'):
        return x.decode()
    return str(x)


def to_str_iter(seq):
    """Iterate over sequence yielding string conversion of each item.

    >>> list(_to_str_iter(['hello', 'hello'.encode(), Path('/hello'), None, 0, False]))
    ['hello', 'hello', '/hello', '', '0', 'False']
    >>> list(_to_str_iter('hello'))
    ['hello']
    >>> list(_to_str_iter('hello'.encode()))
    ['hello']
    >>> list(_to_str_iter(pathlib.Path('/hello')))
    ['/hello']
    >>> list(_to_str_iter(None))
    Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
    File "<stdin>", line 30, in _to_str_iter
    TypeError: 'NoneType' object is not iterable
    >>> list(_to_str_iter(0))
    Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
    File "<stdin>", line 7, in _to_str_iter
    TypeError: 'int' object is not iterable
    >>> list(_to_str_iter(False))
    Traceback (most recent call last):
    File "<stdin>", line 1, in <module>
    File "<stdin>", line 7, in _to_str_iter
    TypeError: 'bool' object is not iterable
    """
    if isinstance(seq, (str, bytes, os.PathLike)):
        yield to_str(seq)
    else:
        yield from map(to_str, seq)


def range_expr(s: str, prog: re.Pattern = re.compile(r'(?P<start>\d+)(?:-(?P<stop>\d+)(?::(?P<step>\d+))?)?$')) -> Iterable[int]:
    """Return iterable corresponding to the given range expression of
    the form START[-STOP[:STEP]][,...].

    NOTE: Range expressions are intended to be compatible with the --cpu-list
    expressions from taskset(1).

    >>> tuple(range_expr('0-10:2'))
    (0, 2, 4, 6, 8, 10)
    >>> tuple(range_expr('0-10:2,0-10:3'))
    (0, 2, 4, 6, 8, 10, 0, 3, 6, 9)
    >>> set(range_expr('0-10:2,0-10:3'))
    {0, 2, 3, 4, 6, 8, 9, 10}
    """
    iters = []
    for expr in s.split(','):
        m = prog.match(expr)
        if m is None:
            raise ValueError(f'Invalid range: {expr}')
        args = list(map(int, filter(None, m.groups())))
        if len(args) == 1:
            # Singular value specified
            args.append(args[0]+1,)
        else:
            # Add one to stop value, to make the range inclusive
            args[1] += 1
        iters.append(range(*args))
    return chain.from_iterable(iters)


def enable_logging(level=logging.DEBUG):
    logging.basicConfig(stream=sys.stdout, level=level)


def user_print(*args, **kwargs):
    if this_process.my_puid >= FIRST_PUID:
        kwargs['flush'] = True
        print(*args, **kwargs)


def port_check(ip_port):
    s = socket(AF_INET, SOCK_STREAM)
    try:
        s.bind(ip_port)
    except Exception:
        if s is not None:
            s.close()
        return False
    else:
        s.close()
        return True


def get_host_info(network_prefix) -> tuple[str, str, list[str]]:
    """Return username, hostname, and list of IP addresses."""
    from dragon.transport.ifaddrs import getifaddrs, InterfaceAddressFilter
    from ..dlogging.util import DragonLoggingServices as dls

    log = logging.getLogger(dls.LA_BE).getChild('get_host_info')
    _user = os.environ.get('USER', str(os.getuid()))

    ifaddr_filter = InterfaceAddressFilter()
    ifaddr_filter.af_inet(inet6=False)  # Disable IPv6 for now
    ifaddr_filter.up_and_running()
    try:
        ifaddrs = list(filter(ifaddr_filter, getifaddrs()))
    except OSError:
        log.exception('Failed to get network interface addresses')
        raise
    if not ifaddrs:
        _msg = 'No network interface with an AF_INET address was found'
        log.error(_msg)
        raise RuntimeError(_msg)

    try:
        re_prefix = re.compile(network_prefix)
    except (TypeError, NameError):
        _msg = "expected a string regular expression for network interface network_prefix"
        log.error(_msg)
        raise RuntimeError(_msg)

    ifaddr_filter.clear()
    ifaddr_filter.name_re(re.compile(re_prefix))
    ip_addrs = [ifa['addr']['addr'] for ifa in filter(ifaddr_filter, ifaddrs)]
    if not ip_addrs:
        _msg = f'No IP addresses found matching regex pattern: {network_prefix}'
        log.error(_msg)
        raise RuntimeError(_msg)
    log.info(f"Found IP addresses: {','.join(ip_addrs)}")

    return _user, ip_addrs

src/dragon/infrastructure/util.py
DragonHPC/dragon
import enum

from . import launch_selector, network_config, launchargs

@enum.unique
class ConfigOutputType(enum.Enum):
    JSON = 'json'
    YAML = "yaml"

    def __str__(self):
        return self.value


@enum.unique
class WLM(enum.Enum):
    SLURM = 'slurm'
    PBS_PALS = 'pbs+pals'

    def __str__(self):
        return self.value

    @staticmethod
    def from_str(s):
        try:
            return WLM(s)
        except KeyError:
            raise ValueError()

src/dragon/launcher/__init__.py
DragonHPC/dragon
import os
import logging
import threading
import re
import subprocess
from enum import Enum
from functools import total_ordering

from ..utils import B64
from ..managed_memory import MemoryPool, DragonPoolError, DragonMemoryError
from ..channels import Channel, register_gateways_from_env, ChannelError, discard_gateways

from . import util as dlutil
from ..transport.overlay import start_overlay_network

from ..dlogging.util import setup_dragon_logging, setup_BE_logging, detach_from_dragon_handler
from ..dlogging.util import DragonLoggingServices as dls
from ..dlogging.logger import DragonLogger, DragonLoggingError

from ..infrastructure import facts as dfacts
from ..infrastructure import messages as dmsg
from ..infrastructure.node_desc import NodeDescriptor
from ..infrastructure.connection import Connection, ConnectionOptions
from ..infrastructure.parameters import POLICY_INFRASTRUCTURE
from ..infrastructure.util import NewlineStreamWrapper, route


def _get_host_info(network_prefix) -> tuple[str, str, list[str]]:
    """Return username, hostname, and list of IP addresses."""
    from dragon.transport.ifaddrs import getifaddrs, InterfaceAddressFilter

    log = logging.getLogger(dls.LA_BE).getChild('get_host_info')
    _user = os.environ.get('USER', str(os.getuid()))

    ifaddr_filter = InterfaceAddressFilter()
    ifaddr_filter.af_inet(inet6=False)  # Disable IPv6 for now
    ifaddr_filter.up_and_running()
    try:
        ifaddrs = list(filter(ifaddr_filter, getifaddrs()))
    except OSError:
        log.exception('Failed to get network interface addresses')
        raise
    if not ifaddrs:
        _msg = 'No network interface with an AF_INET address was found'
        log.error(_msg)
        raise RuntimeError(_msg)

    try:
        re_prefix = re.compile(network_prefix)
    except (TypeError, NameError):
        _msg = "expected a string regular expression for network interface network_prefix"
        log.error(_msg)
        raise RuntimeError(_msg)

    ifaddr_filter.clear()
    ifaddr_filter.name_re(re_prefix)
    ip_addrs = [ifa['addr']['addr'] for ifa in filter(ifaddr_filter, ifaddrs)]
    if not ip_addrs:
        _msg = f'No IP addresses found matching regex pattern: {network_prefix}'
        log.error(_msg)
        raise RuntimeError(_msg)
    log.info(f"Found IP addresses: {','.join(ip_addrs)}")

    return _user, ip_addrs


@total_ordering
class BackendState(Enum):
    """Enumerated states of Dragon LauncherBackEnd
    """

    INITIALIZING = 0
    OVERLAY_STARTING = 1
    OVERLAY_UP = 2
    OVERLAY_THREADS_UP = 3
    LS_STARTING = 4
    LS_UP = 5
    TA_UP = 6
    BACKEND_UP = 7
    TEARDOWN = 8
    BACKEND_DOWN = 9

    def __lt__(self, other):
        if self.__class__ is other.__class__:
            return self.value < other.value

    def __gt__(self, other):
        if self.__class__ is other.__class__:
            return self.value > other.value

    def __eq__(self, other):
        if self.__class__ is other.__class__:
            return self.value == other.value


class LauncherBackEnd:
    """State manager for Dragon BackEnd

    Returns:
        LauncherBackEnd: instance of class
    """

    _DTBL = {}  # dispatch router for msgs, keyed by type of message

    def __init__(self, transport_test_env, network_prefix):

        self._state = BackendState.INITIALIZING
        self._abnormally_terminating = False

        self.transport_test_env = bool(transport_test_env)
        self.network_prefix = network_prefix

        # The la_be_stdin queue is the connection between this main thread
        # (i.e. the launcher backend) and the OverlayNet down the tree thread.
        # The down the tree thread writes to this queue when communicating
        # a message sent from the front end to this queue.
        self.la_be_stdin = dlutil.SRQueue()

        # The la_be_stdout queue is the connection between this main thread
        # (i.e. the launcher backend) and the OverlayNet up the tree thread.
        # This thread writes to this queue when it wants to send something
        # to the launcher front end.
        self.la_be_stdout = dlutil.SRQueue()

        # This shutdown value is used to control when to bring down the OverlayNet
        # and its threads for routing values up and down the tree.
        self._shutdown = threading.Event()
        self._logging_shutdown = threading.Event()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):

        log = logging.getLogger('__exit__')
        self._cleanup()

        # If we were exiting under pretense of an abnormal termination, raise that error
        log.debug('cleaned up backend')
        if self._abnormally_terminating:
            log.debug('raising RuntimeError')
            raise RuntimeError("Abnormal exit detected")

    def _cleanup(self):
        """General purpose clean-up all the crap in this object method"""

        log = logging.getLogger(dls.LA_BE).getChild('_cleanup')
        # handle server threads. Assume it's abnormal because if they're not closed
        # before here, they should be
        log.debug('cleaning up threads')
        self._close_threads(abnormal=True)

        # Make sure our LS is down
        try:
            log.debug('waiting on LS')
            self.ls_proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            log.debug('killing LS')
            self.ls_proc.kill()
        except AttributeError:
            pass

        log.debug('closing connections')
        self._close_conns_chs_pools()

        log.debug('backend down')
        self._state = BackendState.BACKEND_DOWN

    def _close_threads(self, abnormal=False):
        """Join on the launcher backend threads"""

        if self._state == BackendState.BACKEND_DOWN:
            return

        log = logging.getLogger(dls.LA_BE).getChild('_close_threads')
        halt_overlay_msg = dmsg.HaltOverlay(tag=dlutil.next_tag())
        halt_logging_msg = dmsg.HaltLoggingInfra(tag=dlutil.next_tag())

        try:
            if self.recv_msgs_from_overlaynet_thread.is_alive():
                if abnormal:
                    log.info('abnormal closing of recv overlay thread')
                    self._shutdown.set()
                    self.infra_in_bd.send(halt_overlay_msg.serialize())
                self.recv_overlaynet_thread.join()
        except AttributeError:
            pass
        log.info('recv_overlaynet_thread joined')

        try:
            if self.send_msgs_to_overlaynet_thread.is_alive():
                if abnormal:
                    self.la_be_stdout.send(halt_overlay_msg.serialize())
                self.recv_overlaynet_thread.join()
        except AttributeError:
            pass
        log.info('send_overlaynet_thread joined')

        try:
            if self.send_log_to_overlaynet_thread.is_alive():
                # Send it a message to make sure it can get out
                self._logging_shutdown.set()
                self.dragon_logger.put(halt_logging_msg.serialize())
            self.send_log_to_overlaynet_thread.join()
        except AttributeError:
            pass
        log.info('send_logs_to_overlaynet_thread joined')

        try:
            if self.channel_monitor_thread.is_alive():
                if abnormal:
                    self.la_queue_bd.send(halt_overlay_msg.serialize())
                self.recv_overlaynet_thread.join()
        except AttributeError:
            pass
        log.info('channel_monitor_thread joined')

    def _close_conns_chs_pools(self):
        """Clean up all communication-based data structures:"""

        if self._state == BackendState.BACKEND_DOWN:
            return

        log = logging.getLogger(dls.LA_BE).getChild('_close_conns')
        try:
            log.debug('local_inout close')
            self.local_inout.close()
        except Exception:
            pass

        try:
            log.debug('local_in close')
            self.local_in.close()
        except Exception:
            pass

        try:
            log.debug('local_in_be close')
            self.local_in_be.close()
        except Exception:
            pass

        try:
            log.debug('local_out close')
            self.local_out.close()
        except Exception:
            pass

        try:
            log.debug('ls_queue close')
            self.ls_queue.close()
        except Exception:
            pass

        try:
            log.debug('la_queue close')
            self.la_queue.close()
        except Exception:
            pass

        try:
            log.debug('la_queue_bd close')
            self.la_queue_bd.close()
        except Exception:
            pass

        try:
            log.debug('gs_queue close')
            self.gs_queue.close()
        except Exception:
            pass

        try:
            log.debug('be_inbound destroy')
            self.be_inbound.destroy()
        except Exception:
            pass

        try:
            log.debug('local_ch_in destroy')
            self.local_ch_in.destroy()
        except Exception:
            pass

        try:
            log.debug('local_ch_out destroy')
            self.local_ch_out.destroy()
        except Exception:
            pass

        try:
            log.debug('gw_ch destroy')
            self.gw_ch.destroy()
            discard_gateways()
            try:
                del os.environ[dfacts.GW_ENV_PREFIX+str(dfacts.DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE)]
            except KeyError:
                pass
            log.debug('gw closed')
        except Exception:
            pass

        try:
            log.debug('ls_stdin close')
            self.ls_stdin.close()
        except AttributeError:
            pass

        try:
            log.debug('ls_stdout close')
            self.ls_stdout.close()
        except Exception:
            pass

        try:
            log.debug('dragon_logger destroy')
            self.dragon_logger.destroy()
        except (DragonLoggingError, AttributeError):
            pass

        try:
            log.debug('be_mpool destroy')
            self.be_mpool.destroy()
        except Exception:
            pass

    def _handle_abnormal_gsteardown(self, msg):
        '''Send teardown messages through appropriate queues based on backend state'''
        log = logging.getLogger(dls.LA_BE).getChild('_handle_abnormal_gsteardown')

        if self._state < BackendState.OVERLAY_THREADS_UP:
            log.debug(f'sending {type(msg)} via frontend channel rather than send_overlay thread')
            self.infra_out.send(msg)
        else:
            log.debug(f'sending {type(msg)} message to send_overlay thread')
            self.la_be_stdout.send(msg)

    def _abnormal_termination(self, msg):
        """Clean up the backend when we've found ourselves in an indeterminate state
        """
        log = logging.getLogger('Abnormal termination')

        self._abnormally_terminating = True

        log.debug('Begin handling of abnormal termination')

        # Handle the situation where our runtime services didn't get all the way up, but the Frontend
        # signaled a GSTeardown for whatever reason
        if self._state >= BackendState.OVERLAY_UP:

            # in this case we weren't able to route a GSTeardown to GS. So, we need to act on its
            # behalf in the teardown and just let LS kill it later.
            if isinstance(msg, dmsg.GSTeardown):
                self._handle_abnormal_gsteardown(dmsg.GSHalted(tag=dlutil.next_tag()).serialize())

            elif isinstance(msg, dmsg.SHHaltTA):
                self._handle_abnormal_gsteardown(dmsg.TAHalted(tag=dlutil.next_tag()).serialize())

            elif isinstance(msg, dmsg.SHTeardown):
                self._handle_abnormal_gsteardown(dmsg.SHHaltBE(tag=dlutil.next_tag()).serialize())

            elif isinstance(msg, dmsg.BEHalted):
                pass

    def _close_overlay_comms(self):

        log = logging.getLogger(dls.LA_BE).getChild('_close_overlay_comms')

        log.info('shutting down overlay tree agent')
        self.local_inout.send(dmsg.BEHaltOverlay(tag=dlutil.next_tag()).serialize())
        ta_halted = dlutil.get_with_blocking(self.local_inout)
        assert isinstance(ta_halted, dmsg.OverlayHalted)
        self.msg_log.info(f'received {ta_halted} from local TCP agent')
        try:
            self.tree_proc.wait(timeout=3)
        except subprocess.TimeoutExpired:
            self.tree_proc.kill()
        log.info("overlay tree down")

        # Clean up the resources needed for the
        # infrastructure transport agent
        log.info('destroying tcp channels')
        self.infra_out.close()
        log.debug('infra out closed')
        self.infra_in.close()
        log.debug('infra in closed')

        self.dragon_logger.destroy()
        log.debug('dragon_logger closed')

        # Close rest of infrastructure
        self.gw_ch.destroy()
        discard_gateways()
        log.debug('gateway closed')
        try:
            del os.environ[dfacts.GW_ENV_PREFIX+str(dfacts.DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE)]
        except KeyError:
            pass
        self.local_inout.close()
        log.debug('local_inout closed')
        self.local_ch_in.destroy()
        log.debug('local_ch_in destroyed')
        self.local_ch_out.destroy()
        log.debug('local_ch_out destroyed')
        self.be_inbound.destroy()
        log.debug('be_inbound destroyed')
        self.be_mpool.destroy()
        log.debug('be_mpool destroyed')

        log.info('overlay infra down')

    def _start_localservices(self):
        log = logging.getLogger(dls.LA_BE).getChild('start_localservices')

        log.info("standing up local services")
        try:
            ls_args = [dfacts.PROCNAME_LS]
            ls_proc = subprocess.Popen(args=ls_args, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
        except Exception as ex:
            log.fatal('BE failed to stand up LS')
            raise RuntimeError('Local Services launch failed from launcher backend') from ex
        log.debug("LA backend started LS")

        # Return the file descriptors to make mocking easier
        return ls_proc, ls_proc.stdin.fileno(), ls_proc.stdout.fileno()

    def run_startup(self,
                    arg_ip_addr: str,
                    arg_host_id: str,
                    frontend_sdesc: B64,
                    level=logging.INFO, fname=None):
        """Begin bringup of backend services

        :param arg_ip_addr: IP address of frontened with port appended
        :type arg_ip_addr: str
        :param arg_host_id: host ID of frontend
        :type arg_host_id: str
        :param frontend_sdesc: serialized channel descriptor for comms to frontend
        :type frontend_sdesc: B64
        :param level: logging level, defaults to logging.INFO
        :type level: int, optional
        :param fname: filename to log to, defaults to None
        :type fname: str, optional
        """
        log = logging.getLogger(dls.LA_BE).getChild('run_startup')

        net_conf = NodeDescriptor.get_local_node_network_conf()
        self.host_id = int(net_conf.host_id)
        self.hostname = net_conf.name

        conn_options = ConnectionOptions(min_block_size=2 ** 16)
        conn_policy = POLICY_INFRASTRUCTURE
        try:
            local_cuid_in = dfacts.be_local_cuid_in_from_hostid(self.host_id)
            local_cuid_out = dfacts.be_local_cuid_out_from_hostid(self.host_id)
            gw_cuid = dfacts.be_gw_cuid_from_hostid(self.host_id)
            be_cuid = dfacts.be_fe_cuid_from_hostid(self.host_id)

            # Create my memory pool
            self.be_mpool = MemoryPool(int(dfacts.DEFAULT_BE_OVERLAY_TRANSPORT_SEG_SZ),
                                       f'{os.getuid()}_{os.getpid()}_{self.host_id}' + dfacts.DEFAULT_POOL_SUFFIX,
                                       dfacts.be_pool_muid_from_hostid(self.host_id))
            puid, mpool_fname = MemoryPool.serialized_uid_fname(self.be_mpool.serialize())
            log.info(f'be_mpool has uid {puid} and file {mpool_fname}')

            # Create my receiving channels:
            self.be_inbound = Channel(self.be_mpool, be_cuid)
            self.local_ch_in = Channel(self.be_mpool, local_cuid_in)
            self.local_ch_out = Channel(self.be_mpool, local_cuid_out)
            self.local_inout = Connection(inbound_initializer=self.local_ch_in,
                                          outbound_initializer=self.local_ch_out,
                                          options=conn_options,
                                          policy=conn_policy)
            self.infra_in = Connection(inbound_initializer=self.be_inbound,
                                       options=conn_options,
                                       policy=conn_policy)

            # Create a backdoor into sending messages to infra_in in case I need
            # to tell the thread using it to shutdown:
            self.infra_in_bd = Connection(outbound_initializer=self.be_inbound,
                                          options=conn_options,
                                          policy=conn_policy)

            # Create a gateway channel for myself
            self.gw_ch = Channel(self.be_mpool, gw_cuid)

            # Connect to frontend
            be_outbound = Channel.attach(frontend_sdesc.decode(), mem_pool=self.be_mpool)
            conn_options = ConnectionOptions(default_pool=self.be_mpool, min_block_size=2 ** 16)
            self.infra_out = Connection(outbound_initializer=be_outbound,
                                        options=conn_options,
                                        policy=conn_policy)
            self.infra_out.ghost = True

        except (ChannelError, DragonPoolError,
                DragonMemoryError) as init_err:
            log.fatal('could not create resources')
            raise RuntimeError('infrastructure transport resource creation failed') from init_err

        self._state = BackendState.OVERLAY_STARTING
        log.info('Channels created from backend memory pool')

        # Set gateway in my environment and register it
        encoded_ser_gw = B64(self.gw_ch.serialize())
        encoded_ser_gw_str = str(encoded_ser_gw)
        os.environ[dfacts.GW_ENV_PREFIX + str(dfacts.DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE)] = encoded_ser_gw_str
        register_gateways_from_env()

        # start my transport agent
        # Add my own host_id and ip_addr to frontend's
        host_ids = [arg_host_id, str(self.host_id)]
        ip_addrs = [arg_ip_addr, net_conf.ip_addrs[0]]  # it includes the port
        log.debug(f'standing up tcp agent with gw: {encoded_ser_gw_str}, host_ids={host_ids}, and ip_addrs={ip_addrs}')
        self.dragon_logger = setup_dragon_logging(node_index=0)

        try:
            self.tree_proc = start_overlay_network(ch_in_sdesc=B64(self.local_ch_out.serialize()),
                                                   ch_out_sdesc=B64(self.local_ch_in.serialize()),
                                                   log_sdesc=B64(self.dragon_logger.serialize()),
                                                   host_ids=host_ids,
                                                   ip_addrs=ip_addrs,
                                                   frontend=False,
                                                   env=os.environ)
        except Exception as e:
            log.fatal(f'transport agent launch failed on node with host_id {self.host_id}')
            raise RuntimeError(f'transport agent launch failed on node with host_id {self.host_id}') from e

        # Wait on a message from the transport agent telling me it's good to go.
        log.info('Channel tree initializing....')
        ping_back = dlutil.get_with_blocking(self.local_inout)
        assert isinstance(ping_back, dmsg.OverlayPingBE)
        log.debug(f'comm tree initialized with {type(ping_back)} with pid {self.tree_proc.pid}')

        # Send BEIsUP msg to FE
        # Send my serialized descriptor and host_id to the frontend
        be_ch_desc = str(B64(self.be_inbound.serialize()))
        be_up_msg = dmsg.BEIsUp(tag=dlutil.next_tag(), be_ch_desc=be_ch_desc,
                                host_id=self.host_id)
        self.infra_out.send(be_up_msg.serialize())
        log.info(f'sent BEIsUp to the frontend, with host_id = {self.host_id}')

        # Receive my node_index from the frontend - FENodeIdxBE msg
        fe_node_idx_msg = dlutil.get_with_blocking(self.infra_in)
        assert isinstance(fe_node_idx_msg, dmsg.FENodeIdxBE), 'la_be node_index from fe expected'
        self._state = BackendState.OVERLAY_UP

        log.info('la_be recv FENodeIdxBE')
        self.node_idx = fe_node_idx_msg.node_index
        if self.node_idx < 0:
            log.error(f'node_idx = {self.node_idx} | invalid node id, error starting LA BE')
        log.debug(f'my node index is {self.node_idx}')

        # This starts the "down the tree" router.
        log.debug('starting frontend monitor')
        recv_msgs_from_overlaynet_thread_args = (self.la_be_stdin,)
        self.recv_msgs_from_overlaynet_thread = threading.Thread(name='Frontend Monitor',
                                                                 target=self.receive_messages_from_overlaynet,
                                                                 args=recv_msgs_from_overlaynet_thread_args, daemon=False)
        self.recv_msgs_from_overlaynet_thread.start()

        # Start the "up the tree" router
        send_msgs_to_overlaynet_thread_args = (self.la_be_stdout,)
        self.send_msgs_to_overlaynet_thread = threading.Thread(name='Backend Monitor',
                                                               target=self.send_messages_to_overlaynet,
                                                               args=send_msgs_to_overlaynet_thread_args, daemon=False)
        self.send_msgs_to_overlaynet_thread.start()
        self._state = BackendState.OVERLAY_THREADS_UP

        logger_sdesc = self.dragon_logger.serialize()
        log.debug("setting up logging again")
        level, fname = setup_BE_logging(dls.LA_BE, logger_sdesc=B64(logger_sdesc), fname=fname)
        log = logging.getLogger(dls.LA_BE).getChild('run_startup')

        self.is_primary = self.node_idx == 0
        log.debug(f'primary node status = {self.is_primary}')

        # Add the dragon channel as a handler to our internal logging calls
        log = logging.getLogger(dls.LA_BE).getChild('run_startup')
        log.info(f'start in pid {os.getpid()}, pgid {os.getpgid(0)}')
        log.info(f'OverlayNet be start complete. La_BE running on node {self.node_idx}')
        log.debug('testing debug level')

        # This starts the "logging" router.
        send_log_to_overlaynet_thread_args = (self.dragon_logger, level)
        self.send_log_to_overlaynet_thread = threading.Thread(name='Logging Monitor',
                                                              target=self.send_log_msgs_to_overlaynet,
                                                              args=send_log_to_overlaynet_thread_args, daemon=False)

        self.send_log_to_overlaynet_thread.start()

        # Now start local services....
        self._state = BackendState.LS_STARTING
        self.ls_proc, stdin_fd, stdout_fd = self._start_localservices()
        self.ls_stdin = NewlineStreamWrapper(os.fdopen(stdin_fd, 'wb', buffering=0), read_intent=False, write_intent=True)
        self.ls_stdout = NewlineStreamWrapper(os.fdopen(stdout_fd, 'rb', buffering=0), read_intent=True, write_intent=False)

        # Let local services know we're up and give it everything it needs to get going
        _user, ip_addrs = _get_host_info(self.network_prefix)  # get ip_addrs without the port

        be_node_idx_msg = dmsg.BENodeIdxSH(tag=dlutil.next_tag(), node_idx=self.node_idx,
                                           host_name=self.hostname, ip_addrs=ip_addrs,
                                           primary=self.is_primary, logger_sdesc=B64(logger_sdesc))

        self.ls_stdin.send(be_node_idx_msg.serialize())
        log.info(f'sent BENodeIdxSH(node_idx={self.node_idx}) - m2.1 -- ip_addrs = {ip_addrs}')

        # Wait for SHPingBE on the incoming Posix Message Queue. This tells us the
        # infrastructure channels have been created. Then we know we can attach to them.
        sh_ping_be_msg = dlutil.get_with_blocking(self.ls_stdout)
        assert isinstance(sh_ping_be_msg, dmsg.SHPingBE), 'la_be ping from ls expected'
        log.info('la_be recv SHPingBE - m3')


        # switch to comms with local services over channels and proceed with bring up
        self.ls_channel = Channel.attach(B64.from_str(sh_ping_be_msg.shep_cd).decode())
        self.la_channel = Channel.attach(B64.from_str(sh_ping_be_msg.be_cd).decode())
        self.ls_queue = Connection(outbound_initializer=self.ls_channel, policy=POLICY_INFRASTRUCTURE)
        self.la_queue = Connection(inbound_initializer=self.la_channel, policy=POLICY_INFRASTRUCTURE)
        self.la_queue_bd = Connection(outbound_initializer=self.la_channel, policy=POLICY_INFRASTRUCTURE)
        log.info('la_be attached to ls created channels - a7')

        # Start the channel monitor thread
        channel_monitor_args = (self.la_queue, self.la_be_stdout)
        self.channel_monitor_thread = threading.Thread(name='BE Channel Monitor',
                                                       target=self.be_channel_monitor,
                                                       args=channel_monitor_args, daemon=False)
        self.channel_monitor_thread.start()

        # Now send the BEPingSH message to Local Services on its channel to confirm
        # that this code has switched over to using the infra channels.
        self.ls_queue.send(dmsg.BEPingSH(tag=dlutil.next_tag()).serialize())
        log.info('la_be sent BEPingSH - m4')

        log.debug('Exiting backend startup...')

    def receive_messages_from_overlaynet(self, la_be_stdin: dlutil.SRQueue):
        """This service forwards anything sent down from the frontend launcher to the backend main
           thread for processing.

        Args:
           la_be_stdin (dlutil.SRQueue): The connection to the main thread to route messages to.
        """
        log = logging.getLogger(dls.LA_BE).getChild('recv_msgs_from_overlaynet')
        running = True
    # send msg from callback to parent
        while running:
            try:
                fe_msg = dlutil.get_with_blocking(self.infra_in)
                if isinstance(fe_msg, dmsg.HaltOverlay):
                    running = False
                    break
                la_be_stdin.send(fe_msg.serialize())
                log.info(f'received {type(fe_msg)} from Launcher FE and sent to Launcher BE')
                if isinstance(fe_msg, dmsg.BEHalted):
                    log.debug("Exiting receive_messages_from_overlaynet")
                    running = False
            except Exception as ex:  # pylint: disable=broad-except
                log.critical(f'Caught exception in router: {ex}')
                # Send an ExceptionlessAbort in order to get an AbnormalTermiation
                # to the frontend
                self.la_be_stdout.send(dmsg.ExceptionlessAbort(tag=dlutil.next_tag()).serialize())

        log.info('exiting receive_messages_from_overlaynet thread')

    def send_messages_to_overlaynet(self, la_be_stdout: dlutil.SRQueue):
        """This service forwards any messages writen to la_be_stdout to the front end launcher.

        Args:
            la_be_stdout (dlutil.SRQueue): SimpleQueue used for sending one information
            message from main thread to the frontend via the overlay network
        """
        self.to_overlaynet_log = logging.getLogger(dls.LA_BE).getChild('send_msgs_to_overlaynet')

        try:
            # Just recv messages and send them where they need to
            # go until told to stop
            while not self._shutdown.is_set():
                # recv msg from parent to send to frontend
                msg = dlutil.get_with_blocking(la_be_stdout)

                if isinstance(msg, dmsg.HaltOverlay):
                    self._shutdown.set()
                    self.to_overlaynet_log.info(f"Being told by BE server to exit overlay send loop: {type(msg)}")

                # if it's be halted, get out of here. we got it just to break out of our blocking
                # recv
                if isinstance(msg, dmsg.BEHalted):
                    self.to_overlaynet_log.info(f"Leaving msgs to overlaynet loop due to {type(msg)}")
                    break

                # Handle the case that we're able to catch an abort and need to send
                # that to the frontend but then assume the frontend will no longer talk to us.
                if isinstance(msg, dmsg.ExceptionlessAbort):
                    self.infra_out.send(dmsg.AbnormalTermination(tag=dlutil.next_tag()).serialize())
                    self._abnormally_terminating = True
                    self.to_overlaynet_log.info('forwarded AbnormalTermination to frontend via ExceptionlessAbort')

                # Do any specific to the message handling then forward it along
                self.to_overlaynet_log.info(f"received {type(msg)}")
                if type(msg) in LauncherBackEnd._DTBL:
                    self.infra_out.send(msg.serialize())
                    self.to_overlaynet_log.info(f"forwarded {type(msg)} to frontend")

            self.to_overlaynet_log.info("exiting thread routine")

        except Exception as err:  # pylint: disable=broad-except
            self.to_overlaynet_log.exception(f'OverlayNet BE thread failure: {err}')

        self.to_overlaynet_log.info('setting shutdown flag')
        self._shutdown.set()
        self.to_overlaynet_log.info('exiting send_messages_to_overlaynet')

    def send_log_msgs_to_overlaynet(self, my_dragon_logger: DragonLogger, level: int):
        """This service forwards any log messages received on the Dragon logging channel to
        the OverlayNet queue to be sent to the front end launcher.

        Args:
            logger_sdesc (B64): Dragon handle used to attach to the logging channel
            level (int): minimum log priority of messages to forward
        """
        log = logging.getLogger(dls.LA_BE).getChild('forward_logging_to_FE')
        try:
            # Set timeout to None which allows better interaction with the GIL
            while not self._logging_shutdown.is_set():
                try:
                    serialized_msg = my_dragon_logger.get(level, timeout=None)
                    if isinstance(serialized_msg, dmsg.HaltLoggingInfra):
                        break
                    elif isinstance(serialized_msg, dmsg.LoggingMsg):
                        self.infra_out.send(serialized_msg)
                except Exception:
                    pass
            log.debug("exiting send logs loop")
        except Exception as ex:  # pylint: disable=broad-except
            log.warning(f'Caught exception {ex} in logging thread')

    def be_channel_monitor(self, la_in, la_be_stdout):
        """This function (running in a thread) monitors the backend channel
        and any communication coming into it from the rest of the node
        (i.e. local services or global services)

        Args:
            la_in (Connection): communication channel receiving messages
            la_be_stdout (SRQueue): main queue of messages to process
        """
        log = logging.getLogger(dls.LA_BE).getChild('be_channel_monitor')
        running = True
        while running:
            try:
                try:
                    msg = dlutil.get_with_blocking(la_in)
                except Exception:
                    # Don't deliver abort if we're already in the middle of teardown.
                    if self._state < BackendState.TEARDOWN:
                        msg = dmsg.ExceptionlessAbort(tag=dlutil.next_tag())
                    else:
                        # Don't throw anything here as we're in the middle of teardown and want that
                        # to continue and ideally complete. But do set the flag to raise an exception at exit
                        log.debug('Caught an abnormal termination during teardown. Setting to flag to raise Exception at exit.')
                        self._abnormally_terminating = True
                        continue

                if isinstance(msg, dmsg.HaltOverlay):
                    break

                if isinstance(msg, dmsg.SHChannelsUp):
                    self._state = BackendState.LS_UP

                la_be_stdout.send(msg.serialize())
                log.debug(f'be_channel_monitor - forwarded {type(msg)}')
                if isinstance(msg, dmsg.SHHaltBE):
                    log.info('setting logging for shutdown')
                    running = False
                    self._logging_shutdown.set()

                    log.info('detaching from dragon logging')
                    detach_from_dragon_handler(dls.LA_BE)
                    self.dragon_logger.put(dmsg.HaltLoggingInfra(tag=dlutil.next_tag()).serialize())

            except Exception as ex:
                log.critical(f'Caught exception in be_channel_monitor thread: {ex}')

        log.debug("exiting channel monitor loop")

    def run_msg_server(self):
        """Handle messages from dragon services
        """
        self.msg_log = logging.getLogger(dls.LA_BE).getChild('msg_server')

        while not self._shutdown.is_set():
            msg = dlutil.get_with_blocking(self.la_be_stdin)
            self.msg_log.info(f"JUST received {type(msg)}")
            if hasattr(msg, 'r_c_uid'):
                msg.r_c_uid = dfacts.launcher_cuid_from_index(self.node_idx)
            self.msg_log.info(f"received {type(msg)}")
            if type(msg) in LauncherBackEnd._DTBL:
                self._DTBL[type(msg)][0](self, msg=msg)
            else:
                self.msg_log.warning(f'unexpected msg type: {repr(msg)}')

        # Join on all service threads
        self.msg_log.debug("joining on threads...")
        self._close_threads()

        self.msg_log.debug("la_be joining on LS child proc")
        self.ls_proc.wait(timeout=None)
        self.msg_log.info('localservices closed from LA_BE perspective')
        self.msg_log.debug("la_be server exiting cleanly")

        # Close infrasructure communication
        self._close_overlay_comms()

    @route(dmsg.GSProcessCreate, _DTBL)
    def handle_gs_process_create(self, msg: dmsg.GSProcessCreate):
        self.msg_log.info('received a GSProcess Create in the Launcher Backend and forwarded to GS.')
        self.gs_queue.send(msg.serialize())

    @route(dmsg.GSProcessCreateResponse, _DTBL)
    def handle_gs_process_create_resp(self, msg: dmsg.GSProcessCreateResponse):
        pass

    @route(dmsg.SHFwdOutput, _DTBL)
    def handle_sh_fwd_output(self, msg: dmsg.SHFwdOutput):
        pass

    @route(dmsg.GSHeadExit, _DTBL)
    def handle_gs_head_exit(self, msg: dmsg.GSHeadExit):
        self.msg_log.info('m3.1 la_be transmitted GSHeadExit to la_fe')

    @route(dmsg.GSTeardown, _DTBL)
    def handle_gs_teardown(self, msg: dmsg.GSTeardown):

        # If we've been told to teardown before infrastructure is fully up,
        # we may be in trouble, so guard against that
        self.msg_log.info('m4.1 primary la_be received GSTeardown')
        try:
            self.gs_queue.send(msg.serialize())
            self._state = BackendState.TEARDOWN
        except (AttributeError, ConnectionError):
            self.msg_log.info('Received a GSTeardown before backend fully up. Will work towards cleaning backend and exit')
            self._abnormal_termination(msg)

    @route(dmsg.SHHaltTA, _DTBL)
    def handle_sh_halt_ta(self, msg: dmsg.SHHaltTA):
        self.msg_log.info('m7.1 la_be received SHHaltTA')
        try:
            if self._state >= BackendState.LS_UP:
                self.ls_queue.send(msg.serialize())
                self.msg_log.info('m7.2 la_be forwarded SHHaltTA to LS')
                if not self.is_primary:
                    self._state = BackendState.TEARDOWN
            else:
                raise ConnectionError('Cannot communicate over LS channel reliably')
        except (AttributeError, ConnectionError):
            self.msg_log.info('Received a SHHaltTA before backend fully up. Will work towards cleaning backend and exit')
            self._abnormal_termination(msg)

    @route(dmsg.SHTeardown, _DTBL)
    def handle_sh_teardown(self, msg: dmsg.SHTeardown):
        self.msg_log.info('m11.1 la_be received SHTeardown')
        self.msg_log.info('m11.2 la_be forwarded SHTeardown to LS')
        self.ls_queue.send(msg.serialize())

    @route(dmsg.BEHalted, _DTBL)
    def handle_be_halted(self, msg: dmsg.BEHalted):
        self.msg_log.info("m14 la_be received dmsg.BEHalted")
        self._shutdown.set()
        self.la_be_stdout.send(msg.serialize())

        # This is the dmsg.BEHalted
        self.ls_stdin.send(msg.serialize())
        self.msg_log.debug("la_be forwarded m14 BEHalted")

    @route(dmsg.AbnormalTermination, _DTBL)
    def handle_abnormal_term(self, msg: dmsg.AbnormalTermination):
        self.msg_log.info("la_be received dmsg.AbnormalTermination")

    @route(dmsg.SHProcessCreate, _DTBL)
    def handle_sh_proc_crate(self, msg: dmsg.SHProcessCreate):
        self.msg_log.info('In Transport Test Mode: Sending SHProcessCreate directly to local services.')
        self.ls_queue.send(msg.serialize())

    @route(dmsg.SHChannelsUp, _DTBL)
    def handle_sh_channels_up(self, msg: dmsg.SHChannelsUp):
        self.to_overlaynet_log.info('la_be recv SHChannelsUp - m5.1')
        self.to_overlaynet_log.info('la_be sent SHChannelsUp to OverlayNet be - m5.2')

    @route(dmsg.LAChannelsInfo, _DTBL)
    def handle_la_channels_info(self, msg: dmsg.LAChannelsInfo):
        # Wait for the LAChannelsInfo Message to be sent
        # The LAChannelsInfo gave us the serialized descriptor of
        # global services channel, so now we can attach to that channel.
        if self.is_primary and (not self.transport_test_env):
            self.gs_channel = Channel.attach(B64.from_str(msg.gs_cd).decode())
            self.gs_queue = Connection(outbound_initializer=self.gs_channel, policy=POLICY_INFRASTRUCTURE)
            self.msg_log.debug(f'Created gs_queue: {self.gs_queue}')

        # Forward the LAChannelsInfo Message to Local Services
        self.ls_queue.send(msg.serialize())
        self.msg_log.info('la_be sent LAChannelsInfo to ls - m5.3')

    @route(dmsg.TAUp, _DTBL)
    def handle_ta_up(self, msg: dmsg.TAUp):
        self._state = BackendState.TA_UP
        self.to_overlaynet_log.info('la_be recvd Ping response from TA - m8.1')
        if not self.is_primary:
            self._state = BackendState.BACKEND_UP
            self.to_overlaynet_log.info('non-primary backend node is up')

    @route(dmsg.GSIsUp, _DTBL)
    def handle_gs_is_up(self, msg: dmsg.GSIsUp):
        self.to_overlaynet_log.info('la_be received GSIsUp ping - m12.1')
        self._state = BackendState.BACKEND_UP

    @route(dmsg.GSHalted, _DTBL)
    def handle_gs_halted(self, msg: dmsg.GSHalted):
        self.to_overlaynet_log.info('m6.1 primary la_be received GSHalted')

    @route(dmsg.TAHalted, _DTBL)
    def handle_ta_halted(self, msg: dmsg.TAHalted):
        self.to_overlaynet_log.info('m10.1 la_be received TAHalted')
        self.to_overlaynet_log.info('m10.2 la_be forwarded TAHalted to FE')

    @route(dmsg.SHHaltBE, _DTBL)
    def handle_sh_halt_be(self, msg: dmsg.SHHaltBE):
        self.to_overlaynet_log.info('m12 la_be received SHHaltBE')

    @route(dmsg.SHProcessExit, _DTBL)
    def handle_sh_process_exit(self, msg: dmsg.SHProcessExit):
        pass

    @route(dmsg.SHProcessCreateResponse, _DTBL)
    def handle_sh_proc_create_response(self, msg: dmsg.SHProcessCreateResponse):
        pass

src/dragon/launcher/backend.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import sys
import logging
import socket

from .backend import LauncherBackEnd
from .launchargs import NETWORK_HELP

from ..utils import B64
from ..infrastructure.facts import PROCNAME_LA_BE, DEFAULT_TRANSPORT_NETIF, TRANSPORT_TEST_ENV
from ..dlogging.util import setup_BE_logging
from ..dlogging.util import DragonLoggingServices as dls
from . import util as dlutil

LOGBASE = 'launcher_multi_be'


def main(transport_test_env: bool = False):

    from argparse import ArgumentParser
    from distutils.util import strtobool
    from dragon.utils import set_procname
    from dragon import _patch_multiprocessing

    _patch_multiprocessing()

    set_procname(PROCNAME_LA_BE)

    # Set up the logging level. If debug, start logging immediately
    level, fname = setup_BE_logging(dls.LA_BE)
    log = logging.getLogger(dls.LA_BE).getChild('main')
    log.debug("in multi_be")
    log.debug(f"sys.argv = {sys.argv}")

    parser = ArgumentParser(description='Run Dragon backend')

    parser.add_argument('--ip-addr', metavar='FRONTEND_IP', dest='ip_addrs',
                        type=str, help="IP address to connect to frontend")
    parser.add_argument('--host-id', dest='host_ids', type=str, help="Host ID of frontend")
    parser.add_argument('--frontend-sdesc', dest='frontend_sdesc', type=B64.from_str,
                        help="File descriptor for communication to frontend")
    parser.add_argument('--transport-test', action='store_true', help="Run in transport test mode")
    parser.add_argument('--network-prefix', dest='network_prefix', type=str, help=NETWORK_HELP)

    parser.set_defaults(
        transport_test=bool(strtobool(os.environ.get(TRANSPORT_TEST_ENV, str(transport_test_env)))),
        network_prefix=DEFAULT_TRANSPORT_NETIF
    )
    args = parser.parse_args()

    with LauncherBackEnd(args.transport_test, args.network_prefix) as be_server:
        try:
            be_server.run_startup(args.ip_addrs, args.host_ids, args.frontend_sdesc, level, fname)
            be_server.run_msg_server()

        except Exception as err:  # pylint: disable=broad-except
            log.exception(f'la_be {socket.gethostname} main exception: {err}')
            raise


if __name__ == "__main__":
    main()

src/dragon/launcher/dragon_multi_be.py
DragonHPC/dragon
#!/usr/bin/env python3

"""Simple multi-node dragon infrastructure startup"""
import os
import logging

from .frontend import LauncherFrontEnd, LAUNCHER_FAIL_EXIT
from .launchargs import get_args as get_cli_args

from ..utils import set_procname, set_host_id, host_id
from ..dlogging.util import setup_FE_logging, DragonLoggingServices as dls

from ..infrastructure.facts import PROCNAME_LA_FE, FRONTEND_HOSTID

def main():

    args_map = get_cli_args()

    setup_FE_logging(log_device_level_map=args_map['log_device_level_map'],
                     basename='dragon', basedir=os.getcwd())
    log = logging.getLogger(dls.LA_FE).getChild('main')
    log.info(f'start in pid {os.getpid()}, pgid {os.getpgid(0)}')

    # Before doing anything set my host ID
    set_host_id(FRONTEND_HOSTID)
    log.info(f"set host id to {FRONTEND_HOSTID}, and return {host_id()}")

    for key, value in args_map.items():
        if value is not None:
            log.info(f'args_map: {key}: {value}')

    with LauncherFrontEnd(args_map=args_map) as fe_server:
        try:
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()
        except Exception as err:
            log.exception(f'Error in launcher frontend: {err}')
            return LAUNCHER_FAIL_EXIT

    log.info("exiting frontend")


if __name__ == "__main__":
    set_procname(PROCNAME_LA_FE)
    ecode = main()
    exit(ecode)

src/dragon/launcher/dragon_multi_fe.py
DragonHPC/dragon
#!/usr/bin/env python3
"""Simple single node dragon infrastructure startup"""
import time

import os
import sys

import logging
import shutil
import threading

import dragon.channels as dch

import dragon.localservices.local_svc as dsls
import dragon.globalservices.server as dgs
import dragon.infrastructure.debug_support as dds
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.util as dutil
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.process_desc as pdesc
import dragon.launcher.launchargs as launchargs
import dragon.dlogging.util as dlog
import dragon.launcher.util as dlutil
import dragon.utils as du

# TODO: consider ways to bypass serialization when it isn't really needed

LAUNCHER_FAIL_EXIT = 1

# general amount of patience we have for an expected message
# in startup or teardown before we assume something has gone wrong
TIMEOUT_PATIENCE = 1  # seconds, 1 second.

# time to yield to let other things happen esp in unexpected shutdowns.
TIMEOUT_YIELD = 0.200  # seconds, 200 milliseconds

LOGBASE = 'launcher'


def ls_start(ls_args):
    log = logging.getLogger(LOGBASE).getChild('ls_start')
    try:
        dsls.single(**ls_args)
        log.info('normal exit')
    except RuntimeError:
        log.exception('fail exit')


def build_stdmsg(msg, arg_map, is_stdout=True):
    if arg_map['no_label']:
        return f'{msg.data}'

    msg_str = ""
    if is_stdout:
        msg_str += "[stdout: "
    else:
        msg_str += "[stderr: "

    if arg_map['verbose_label']:
        msg_str += f"PID {msg.pid} @ {msg.hostname}]"
    elif arg_map['basic_label']:
        msg_str += f"Dragon PID {msg.p_uid}]"

    msg_str += f" {msg.data}"
    return msg_str

def output_monitor(la_in):
    arg_map = launchargs.get_args() # Get args once to check decoration options
    while True:
        msg = dmsg.parse(la_in.recv())
        if isinstance(msg, dmsg.SHFwdOutput):
            if msg.fd_num == msg.FDNum.STDOUT.value:
                msg_str = build_stdmsg(msg, arg_map, True)
                sys.stdout.write(msg_str)
                sys.stdout.flush()
            elif msg.fd_num == msg.FDNum.STDERR.value:
                msg_str = build_stdmsg(msg, arg_map, False)
                sys.stderr.write(msg_str)
                sys.stderr.flush()
        elif isinstance(msg, dmsg.GSProcessCreateResponse):
            if msg.err != msg.Errors.SUCCESS:
                print('+++ head proc did not start')
                return LAUNCHER_FAIL_EXIT
        elif isinstance(msg, dmsg.GSHeadExit):
            print('+++ head proc exited, code {}'.format(msg.exit_code))
            return msg.exit_code
        elif isinstance(msg, dmsg.Breakpoint):
            log = logging.getLogger(LOGBASE).getChild('breakpoint')
            log.info(f'p_uid {msg.p_uid} node {msg.index}')
            dds.handle_breakpoint(msg)
        else:
            print('unexpected message: {}'.format(msg))


def shutdown_monitor(la_in):
    while la_in.poll(timeout=TIMEOUT_PATIENCE):
        msg = dmsg.parse(la_in.recv())
        if isinstance(msg, dmsg.SHFwdOutput):
            if msg.fd_num == msg.FDNum.STDOUT.value:
                sys.stdout.write(msg.data)
                sys.stdout.flush()
            elif msg.fd_num == msg.FDNum.STDERR.value:
                sys.stderr.write(msg.data)
                sys.stderr.flush()
        else:
            return msg

    raise TimeoutError()


def main():
    arg_map = launchargs.get_args()

    dlog.setup_FE_logging(log_device_level_map=arg_map['log_device_level_map'],
                          basename='dragon', basedir=os.getcwd())

    log = logging.getLogger(LOGBASE).getChild('main')
    log.info(f'start in pid {os.getpid()}, pgid {os.getpgid(0)}')

    start_msg = dlutil.mk_head_proc_start_msg()

    ls_stdin = dlutil.SRQueue()
    ls_stdout = dlutil.SRQueue()
    ls_args = {'ls_stdin': ls_stdin, 'ls_stdout': ls_stdout}

    ls_thread = threading.Thread(name='local services', target=ls_start,
                                 args=(ls_args,), daemon=True)

    shm_status = dutil.survey_dev_shm()

    try:  # ls startup
        ls_thread.start()
        ls_stdin.send(dmsg.BENodeIdxSH(tag=dlutil.next_tag(), node_idx=0).serialize())
        be_ping = dmsg.parse(ls_stdout.recv())
        assert isinstance(be_ping, dmsg.SHPingBE)

        ls_in_ch = dch.Channel.attach(du.B64.str_to_bytes(be_ping.shep_cd))
        la_in_ch = dch.Channel.attach(du.B64.str_to_bytes(be_ping.be_cd))
        gs_in_ch = dch.Channel.attach(du.B64.str_to_bytes(be_ping.gs_cd))

        ls_in_wh = dconn.Connection(outbound_initializer=ls_in_ch, policy=dparm.POLICY_INFRASTRUCTURE)
        la_in_rh = dconn.Connection(inbound_initializer=la_in_ch, policy=dparm.POLICY_INFRASTRUCTURE)
        gs_in_wh = dconn.Connection(outbound_initializer=gs_in_ch, policy=dparm.POLICY_INFRASTRUCTURE)

        ls_in_wh.send(dmsg.BEPingSH(tag=dlutil.next_tag()).serialize())
        ch_up = dlutil.get_with_timeout(la_in_rh)

        assert isinstance(ch_up, dmsg.SHChannelsUp)
    except (AssertionError, dch.ChannelError, TimeoutError) as err:
        log.exception('ls startup')
        print(f'ls startup failed:\n{err}')
        dutil.compare_dev_shm(shm_status)
        return LAUNCHER_FAIL_EXIT

    try:  # gs startup
        gs_thread = threading.Thread(name='global services', target=dgs.single_thread,
                                     args=(ls_in_wh,), daemon=True)

        gs_thread.start()
        gs_up = dlutil.get_with_timeout(la_in_rh)
        assert isinstance(gs_up, dmsg.GSIsUp)
        gs_in_wh.send(start_msg.serialize())
    except (AssertionError, TimeoutError) as err:
        log.exception('gs startup')
        print(f'gs startup failed:\n{err}')
        dutil.compare_dev_shm(shm_status)
        return LAUNCHER_FAIL_EXIT

    try:
        exit_code = output_monitor(la_in_rh)
    except KeyboardInterrupt:
        print('KeyboardInterrupt - going to teardown')
        log.warning('KeyboardInterrupt - tearing down')
        exit_code = LAUNCHER_FAIL_EXIT

    try:
        gs_in_wh.send(dmsg.GSTeardown(tag=dlutil.next_tag()).serialize())
        gs_halt = shutdown_monitor(la_in_rh)
        time.sleep(TIMEOUT_YIELD)
        ls_in_wh.send(dmsg.SHTeardown(tag=dlutil.next_tag()).serialize())
        be_halt = shutdown_monitor(la_in_rh)
        time.sleep(TIMEOUT_YIELD)
        ls_stdin.send(dmsg.BEHalted(tag=dlutil.next_tag()).serialize())
        sh_halt = dmsg.parse(ls_stdout.recv())

        if not isinstance(gs_halt, dmsg.GSHalted):
            log.warning(f'expected GSHalted got {gs_halt}')

        if not isinstance(be_halt, dmsg.SHHaltBE):
            log.warning(f'expected SHHaltBE got {be_halt}')

        if not isinstance(sh_halt, dmsg.SHHalted):
            log.warning(f'expected SHHalted got {sh_halt}')
    except TimeoutError:
        log.exception('teardown error')

    try:
        gs_thread.join(timeout=TIMEOUT_PATIENCE)
        ls_thread.join(timeout=TIMEOUT_PATIENCE)
    except TimeoutError:
        log.warning('infrastructure thread hang')
        print('warning: infrastructure thread hang')

    dutil.compare_dev_shm(shm_status)

    return exit_code


if __name__ == '__main__':
    ecode = main()
    exit(ecode)

src/dragon/launcher/dragon_single.py
DragonHPC/dragon
import ctypes
import os
import sys
import logging
import threading
import signal
import subprocess
from enum import Enum
from shlex import quote

from functools import total_ordering

from ..utils import B64
from ..channels import Channel, ChannelError, ChannelEmpty, register_gateways_from_env, discard_gateways
from ..managed_memory import MemoryPool, DragonPoolError, DragonMemoryError
from ..transport.overlay import start_overlay_network

from ..dlogging.util import DragonLoggingServices as dls
from ..dlogging.util import _get_dragon_log_device_level, LOGGING_OUTPUT_DEVICE_DRAGON_FILE
from ..dlogging.logger import DragonLogger, DragonLoggingError

from ..infrastructure.util import route
from ..infrastructure.parameters import POLICY_INFRASTRUCTURE, this_process
from ..infrastructure.connection import Connection, ConnectionOptions
from ..infrastructure.node_desc import NodeDescriptor
from ..infrastructure import facts as dfacts
from ..infrastructure import messages as dmsg
from ..infrastructure.messages import AbnormalTerminationError

from . import util as dlutil
from .network_config import NetworkConfig
from .launchargs import parse_hosts
from .wlm import WLM
from .wlm.ssh import SSHSubprocessPopen

LAUNCHER_FAIL_EXIT = 1


class LauncherImmediateExit(Exception):

    def __init__(self, msg=''):
        self._msg = msg

    def __str__(self):
        return f'{self._msg}'

    def __repr__(self):
        return f"{str(__class__)}({repr(self._msg)})"


class SigIntImmediateExit(Exception):

    def __init__(self, msg=''):
        self._msg = msg

    def __str__(self):
        return f'{self._msg}'

    def __repr__(self):
        return f"{str(__class__)}({repr(self._msg)})"


@total_ordering
class FrontendState(Enum):
    """Enumerated states of Dragon LauncherFrontEnd"""

    NET_CONFIG = 0
    OVERLAY_STARTING = 1
    OVERLAY_UP = 2
    STARTUP = 3
    BACKEND_LAUNCHED = 4
    STOOD_UP = 5
    APP_EXECUTION = 6
    TEARDOWN = 7
    ABNORMAL_TEARDOWN = 8
    LAUNCHER_DOWN = 9

    def __lt__(self, other):
        if self.__class__ is other.__class__:
            return self.value < other.value
        else:
            return False

    def __gt__(self, other):
        if self.__class__ is other.__class__:
            return self.value > other.value
        else:
            return False

    def __eq__(self, other):
        if self.__class__ is other.__class__:
            return self.value == other.value
        else:
            return False

    def __ne__(self, other):
        if self.__class__ is other.__class__:
            return not self.value == other.value
        else:
            return True


class LauncherFrontEnd():
    """State manager for Dragon Frontend

    Returns:
        LauncherFrontEnd: instance of class
    """
    _DTBL = {}  # dispatch router for msgs, keyed by type of message
    _STBL = {}  # dispatch router for SIGINT handlign, keyed by current state
    _TRDOWN = {}  # dispatch router teardown errors, keyed by FrontendState

    _BE_LAUNCHARGS = f'{dfacts.PROCNAME_LA_BE}' \
                     ' --ip-addr {ip_addr}' \
                     ' --host-id {host_id}' \
                     ' --frontend-sdesc {frontend_sdesc}' \
                     ' --network-prefix {network_prefix}'

    _STATE = None

    def __init__(self, args_map, sigint_trigger=None):

        log = logging.getLogger(dls.LA_FE).getChild('LauncherFrontEnd init')
        log.info(f'start in pid {os.getpid()}, pgid {os.getpgid(0)}')

        # Discover whether the transport service test mode is being requested
        # or not.
        self.transport_test_env = os.environ.get(dfacts.TRANSPORT_TEST_ENV) is not None

        try:
            pals_lib = ctypes.cdll.LoadLibrary('libpals.so')
            self.pals_lib_present = pals_lib != None
            del pals_lib
        except:
            self.pals_lib_present = False

        # This running value is used to control when to bring down the OverlayNet service threads
        self._shutdown = threading.Event()
        self._abnormal_termination = threading.Event()

        self.args_map = args_map
        self.nnodes = args_map.get('node_count', 0)
        self.network_prefix = args_map.get('network_prefix', dfacts.DEFAULT_TRANSPORT_NETIF)
        self.port = args_map.get('port', dfacts.DEFAULT_TRANSPORT_PORT)
        self._config_from_file = args_map.get('network_config', None)
        self.transport = args_map.get('transport')

        # If using SSH, confirm we have enough info do that:
        self._wlm = args_map.get('wlm', None)
        self.hostlist = args_map.get('hostlist', None)
        self.hostfile = args_map.get('hostfile', None)

        # Don't default wlm in argparse. We can better do its error/case handling outside of
        # argparse/make it easier to test
        if self._wlm is None:
            self._wlm = dlutil.detect_wlm()

        # If using SSH, confirm we have enough info do that:
        if self._wlm == WLM.SSH and self._config_from_file is None:
            self.hostlist = parse_hosts(self.hostlist, self.hostfile)

        if self._wlm is WLM.SSH and self.transport is not dfacts.TransportAgentOptions.TCP:
            msg = '''
When using SSH to execute dragon jobs, the TCP transport agent is the only allowed agent.
Please resubmit your dragon launch command with the `--transport tcp` option set.
'''
            print(msg)
            sys.exit(LAUNCHER_FAIL_EXIT)

        # Variety of other state trackers:
        self._sigint_count = 0
        self._sigint_timeout = 5.0
        self._bumpy_exit = threading.Event()

        self._proc_create_resps = 0
        self._proc_exits = 0
        self._tas_halted = 0
        self._sh_halt_be = 0
        self._gs_head_exit_received = False
        self._gs_process_create_resp_received = False

        # Objects that need tracked for abnormal termincation
        self._orig_sigint = None
        self.fe_mpool = None
        self.fe_inbound = None
        self.conn_in = None
        self.conn_in_bd = None
        self.conn_outs = None
        self.local_ch_in = None
        self.local_ch_out = None
        self.local_inout = None
        self.gw_ch = None
        self.dragon_logger = None
        self.over_proc = None
        self.recv_overlaynet_thread = None
        self.send_overlaynet_thread = None
        self.recv_logs_from_overlaynet_thread = None

        # Int trigger for raising SIGINT at various points of
        # bringup for unit testing purposes
        self._sigint_trigger = sigint_trigger

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):

        self._cleanup()

    def _kill_backend(self):
        '''Simple function for transmitting SIGKILL to backend with helpful message'''

        if self.wlm_proc.returncode is None:
            self._bumpy_exit.set()
            self.wlm_proc.kill()

    def _takedown_abnormal_backend(self, sigint=False):
        '''Function to iteratively take down the backend with one thread to enable easy abort'''
        log = logging.getLogger(dls.LA_FE).getChild('_takedown_abnormal_backend')

        if not sigint:
            print("Abnormal Exit detected. Attempting to clean up over next minute....", flush=True)

        # This is a chance we try to use the route decorated functions which all assume self.msg_log
        # exits. It doesn't necessarily:
        self.msg_log = logging.getLogger(dls.LA_FE).getChild('_takedown_abnormal_backend')

        # Get the threads down first
        self._close_threads(abnormal=True)

        # Do a quick "should I proceed with a teardown loop logic"
        # node index 0 is the primary node
        if FrontendState.STOOD_UP <= self._STATE < FrontendState.TEARDOWN:
            self._STATE = FrontendState.ABNORMAL_TEARDOWN
            log.debug(f'Updated state to {self._STATE}')
            try:
                self.conn_outs[0].send(dmsg.GSTeardown(tag=dlutil.next_tag()).serialize())
            except Exception:
                raise TimeoutError
            log.info('got GSTeardown out')

            while True:
                msg = dlutil.get_with_timeout(self.conn_in, timeout=self._sigint_timeout)

                if isinstance(msg, dmsg.BEHalted):
                    log.debug('Breaking due to BEHalted')
                    break
                elif type(msg) in LauncherFrontEnd._DTBL:
                    self._DTBL[type(msg)][0](self, msg=msg)
                else:
                    log.debug('raising runtime error due to unknown message')
                    raise RuntimeError

    def _wait_on_wlm_proc_exit(self, timeout=30):
        '''Short function for doing a wait on wlm subprocess. Allows for easier mocking'''

        return self.wlm_proc.wait(timeout)

    def _cleanup_abnormal_state(self, sigint=False):
        """Attempt to clean up backend in case it hasn't"""

        log = logging.getLogger(dls.LA_FE).getChild('_cleanup_abnormal_state')
        log.debug('Entered cleanup routine with a still up backend. Trying to teardown...')

        # If we got here, make sure we raise an exception on exit
        if not sigint:
            self._abnormal_termination.set()
        else:
            log.debug(f'executing abnormal cleanup due to SIGINT rather than AbnormalTermination: {self._STATE}')

        # Set an alarm just in case
        try:
            self._takedown_abnormal_backend(sigint=sigint)
        except (subprocess.TimeoutExpired, RuntimeError, TimeoutError):
            pass

        try:
            log.debug('waiting on wlm_proc exit')
            self._wait_on_wlm_proc_exit(timeout=self._sigint_timeout)
        except (subprocess.TimeoutExpired, RuntimeError, TimeoutError):
            log.debug('forcibly kill wlm proc')
            self._kill_backend()

        self._close_comm_overlaynet(abnormal=True)

        log.debug('Finished abnormal state cleanup')

        log.debug('Sending a message to break msg server out of its loop, just in case')
        self.la_fe_stdin.send(dmsg.LAExit(sigint=sigint, tag=dlutil.next_tag()).serialize())
        log.debug('LAExit sent')
        self._STATE = FrontendState.LAUNCHER_DOWN

        # if during teardown we hit an error, clean ourselves up
        if self._bumpy_exit.is_set():
            self._dragon_cleanup_bumpy_exit()

    def _cleanup(self, sigint=False):
        '''clean up any memory resources'''
        log = logging.getLogger(dls.LA_FE).getChild('_cleanup')

        log.debug('beginning cleanup')

        # Close threads best as possible, assume abnormal
        # because we want to make sure they're down at this point
        self._close_threads(abnormal=True)

        # If we're here due to an abnormal termination, we need
        # to make a good faith effort to teardown the backend
        log.debug(f"current state: {self._STATE}")
        if FrontendState.LAUNCHER_DOWN > self._STATE > FrontendState.STARTUP:
            log.debug("Attemping to take down backend before killing it all off")
            self._cleanup_abnormal_state(sigint=sigint)

        try:
            if self.wlm_proc.returncode is None:
                log.debug('killing workload manager backend proc')
                self.wlm_proc.kill()
        except Exception:
            pass

        self._close_comm_overlaynet()

        # if during teardown we hit an error, clean ourselves up
        if self._bumpy_exit.is_set():
            self._dragon_cleanup_bumpy_exit()

        log.debug('Exiting _cleanup')
        # And raise error
        if self._abnormal_termination.is_set():
            log.debug("Raising Abnormal RuntimeError")
            raise RuntimeError('Abnormal exit detected')

    def _wait_on_overlay_init(self):
        """Wait until overlay network returns OverlayPingLA signaling it is up"""
        log = logging.getLogger('overlay_startup')
        log.info('Channel tree initializing...')
        ping_back = dlutil.get_with_blocking(self.local_inout)
        assert isinstance(ping_back, dmsg.OverlayPingLA)
        log.debug(f'recvd msg type {type(ping_back)}')
        self._STATE = FrontendState.OVERLAY_UP

    def _close_overlay(self, abnormal=False):
        """Close the overlay process"""
        log = logging.getLogger('close_overlay')
        if self._STATE >= FrontendState.OVERLAY_UP:
            try:
                self.local_inout.send(dmsg.LAHaltOverlay(tag=dlutil.next_tag()).serialize())
                log.debug('sent halt overlay signal ')

                # If we're in an abnormal exit, make sure to not block
                try:
                    if abnormal:
                        overlay_halted = dlutil.get_with_timeout(self.local_inout, timeout=self._sigint_timeout)
                    else:
                        overlay_halted = dlutil.get_with_blocking(self.local_inout)
                except TimeoutError:
                    log.debug('timeout on recv dmsg.OverlayHalted')
                else:
                    log.debug(f'received {overlay_halted} from local TCP agent')
                    assert isinstance(overlay_halted, dmsg.OverlayHalted)

            except (AttributeError, ConnectionError):
                pass

        # Make sure the overlay is down before continuing
        try:
            self.over_proc.wait(timeout=self._sigint_timeout)
        except subprocess.TimeoutExpired:
            log.debug('timeout on overlay exit. Killing')
            self.over_proc.kill()
        except AttributeError:
            pass

    def _close_threads(self, abnormal=False):
        log = logging.getLogger(dls.LA_FE).getChild('close_threads')
        halt_overlay_msg = dmsg.HaltOverlay(tag=dlutil.next_tag())
        halt_logging_msg = dmsg.HaltLoggingInfra(tag=dlutil.next_tag())
        if self.recv_overlaynet_thread is not None:
            if self.recv_overlaynet_thread.is_alive():
                if abnormal:
                    log.info('abnormal closing of recv overlay thread')
                    self._shutdown.set()
                    self.conn_in_bd.send(halt_overlay_msg.serialize())
                self.recv_overlaynet_thread.join()
        log.info('recv_overlaynet_thread joined')

        # Test a rapid sigint trigger
        try:
            if self._sigint_trigger == 8:
                log.debug("raising signal 2nd time")
                self._sigint_trigger = None
                signal.raise_signal(signal.SIGINT)
        except TypeError:
            pass

        if self.send_overlaynet_thread is not None:
            if self.send_overlaynet_thread.is_alive():
                if abnormal:
                    self._shutdown.set()
                    self.send_msg_to_overlaynet("A", halt_overlay_msg)
                self.send_overlaynet_thread.join()
        log.info('send_overlaynet_thread joined')
        if self.recv_logs_from_overlaynet_thread is not None:
            if self.recv_logs_from_overlaynet_thread.is_alive():
                # Send it a message to make sure it can get out
                self._shutdown.set()
                self.dragon_logger.put(halt_logging_msg.serialize())
            self.recv_logs_from_overlaynet_thread.join()
        log.info('recv_logs_from_overlaynet_thread joined')

    def _close_comm_overlaynet(self, abnormal=False):

        log = logging.getLogger(dls.LA_FE).getChild('_close_comm_overlaynet')

        # Close the overlay tree agent
        try:
            log.info('shutting down frontend overlay tree agent')
            self._close_overlay(abnormal=abnormal)

        except Exception:
            if abnormal:
                log.info('killing overlay on frontend')
                self.over_proc.kill()
            else:
                raise

        try:
            self.conn_in.close()
        except Exception:
            pass

        try:
            self.conn_in_bd.close()
        except Exception:
            pass

        log.debug('closing overlaynet channels and gateway')
        try:
            for conn_out in self.conn_outs.values():
                try:
                    conn_out.close()
                except Exception:
                    pass
        except Exception:
            pass

        try:
            if self.gw_ch is not None:
                log.debug('closing comms to overlay TCP agent and mpool')
                self.gw_ch.destroy()
                discard_gateways()
                try:
                    del os.environ[dfacts.GW_ENV_PREFIX+str(dfacts.DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE)]
                except KeyError:
                    pass
                log.debug('gw closed')
        except ChannelError:
            pass

        try:
            self.local_inout.close()
        except Exception:
            pass
        log.debug('local_input closed')

        try:
            if self.local_ch_in is not None:
                self.local_ch_in.destroy()
                log.debug('local_ch_in closed')
        except ChannelError:
            pass

        try:
            if self.local_ch_out is not None:
                self.local_ch_out.destroy()
                log.debug('local_ch_out closed')
        except ChannelError:
            pass

        try:
            if self.fe_inbound is not None:
                self.fe_inbound.destroy()
                log.debug('fe_inbound closed')
        except ChannelError:
            pass

        try:
            if self.fe_mpool is not None:
                self.fe_mpool.destroy()
                log.debug('fe_mp closed')
        except (DragonMemoryError, DragonPoolError):
            pass
        log.info('comms overlay network closed')

    def _set_quick_teardown(self):
        '''Set up conditions to tear down runtime as quickly as reasonable'''
        print("Quickly tearing down Dragon runtime due to KeyboardInterrupt.", flush=True)
        self._bumpy_exit.set()
        self._sigint_timeout = 0.01

    def _dragon_cleanup_bumpy_exit(self):
        '''Helper function to launch `dragon-cleanup` at end of teardown to clean up all procs and mem'''
        print("Detected an abnormal exit. Will attempt to clean up Dragon resources...", flush=True)
        subprocess.run(args=['dragon-cleanup', '0'])

        # Make sure we don't manage to call this more than once
        self._bumpy_exit.clear()

    def sigint_handler(self, *args):
        """Handler for SIGINT signals for graceful teardown
        """
        log = logging.getLogger('sigint_handler')
        log.debug('Entered sigint handler')
        self.sigint_log = logging.getLogger('sigint_route')
        self._sigint_count = self._sigint_count + 1

        print("KeyboardInterrupt detected. Attempting to clean up...", flush=True)

        # If the user has triggered SIGINT more than once just get out of here
        # very quickly.
        if self._sigint_count > 1:
            self._set_quick_teardown()

        if self._STATE.value in LauncherFrontEnd._STBL:
            self._STBL[self._STATE.value][0](self)
        else:
            log.warning(f'SIGINT detected in {self._STATE} and no routing exists for it')

    @route(FrontendState.NET_CONFIG.value, _STBL)
    def _sigint_net_config(self):
        # if we're here, we should be able to just exit
        self.sigint_log.debug('SIGINT caught during NET_CONFIG. Executing default SIGINT behavior')
        self._close_comm_overlaynet()
        raise KeyboardInterrupt

    @route(FrontendState.OVERLAY_STARTING.value, _STBL)
    def _sigint_overlay_starting(self):
        # Wait on receipt that the Overlay is up so we can
        # calmly tear it down
        self.sigint_log.debug('SIGINT caught during OVERYLAY_STARTING. Waiting for Overlay to come up and then exit')
        self._wait_on_overlay_init()

        # Close everything else:
        self._close_comm_overlaynet()
        raise KeyboardInterrupt

    @route(FrontendState.OVERLAY_UP.value, _STBL)
    def _sigint_overlay_up(self):
        # Close overlay and then exit
        self.sigint_log.debug('SIGINT caught during OVERYLAY_UP. Closing overlay and exiting')

        # Close everything else:
        self._close_threads(abnormal=True)
        self.sigint_log.debug('closed threads')

        self._close_comm_overlaynet()
        self.sigint_log.debug('cleaned comms')
        raise KeyboardInterrupt

    @route(FrontendState.STARTUP.value, _STBL)
    @route(FrontendState.BACKEND_LAUNCHED.value, _STBL)
    def _sigint_startup(self):
        # If in startup, get through it first
        self.sigint_log.debug("Raising SIGINT caught in bringup.")
        self._cleanup(sigint=True)
        self.sigint_log.debug('cleaned up. Raising KeyboardInterrupt')
        raise KeyboardInterrupt

    @route(FrontendState.STOOD_UP.value, _STBL)
    @route(FrontendState.APP_EXECUTION.value, _STBL)
    def _sigint_stood_up(self):
        self.sigint_log.debug("Doing teardown due to SIGINT in signal handler")
        self.start_sigint_teardown()

    @route(FrontendState.TEARDOWN.value, _STBL)
    @route(FrontendState.ABNORMAL_TEARDOWN.value, _STBL)
    @route(FrontendState.LAUNCHER_DOWN.value, _STBL)
    def _sigint_teardown(self):
        # Assume if we've caught a signal during teardown it's because
        # the user wants us to get out of here quickly.
        self.sigint_log.debug("Caught sigint during teardown. Will begin quick teardown")
        self._set_quick_teardown()
        self.start_sigint_teardown()

    def start_sigint_teardown(self):
        """SIGINT teardown init

        Sends GSTeardown message if SIGINT is registered by user
        """
        # Set an alarm just in case
        log = logging.getLogger(dls.LA_FE).getChild('sigint teardown')
        log.debug('SIGINT received. Tearing down runtime....')
        self._cleanup_abnormal_state(sigint=True)

    def _launch_backend(self,
                        nnodes: int,
                        nodelist: list[str],
                        fe_ip_addr: str,
                        fe_host_id: str,
                        frontend_sdesc: str,
                        network_prefix: str):
        """Launch backend with selected wlm"""
        log = logging.getLogger(dls.LA_FE).getChild('_launch_backend')
        try:
            if self._wlm is WLM.SSH:
                be_args = self._BE_LAUNCHARGS.format(ip_addr=fe_ip_addr,
                                                     host_id=fe_host_id,
                                                     frontend_sdesc=frontend_sdesc,
                                                     network_prefix=quote(network_prefix)).split()
            else:
                be_args = self._BE_LAUNCHARGS.format(ip_addr=fe_ip_addr,
                                                     host_id=fe_host_id,
                                                     frontend_sdesc=frontend_sdesc,
                                                     network_prefix=network_prefix).split()
            if self.transport_test_env:
                be_args.append('--transport-test')

        except Exception:
            raise RuntimeError("Unable to construct backend launcher arg list")

        the_env = dict(os.environ)
        the_env['DRAGON_NETWORK_CONFIG'] = self.net.compress()


        # TODO: The differentiation between the SSH path vs. other paths
        #       is clunky. Ideally, this could be abstracted to make the
        #       if/else disappear

        # Syntax is largely same for anything launched with a true WLM
        if self._wlm is not WLM.SSH:

            try:
                wlm_args = dlutil.get_wlm_launch_args(args_map=self.args_map,
                                                      nodes=(nnodes, nodelist),
                                                      wlm=self._wlm)
                args = wlm_args + be_args
                log.info(f'launch be with {args}')


                wlm_proc = subprocess.Popen(args,
                                            stdin=subprocess.DEVNULL,
                                            stdout=subprocess.DEVNULL,
                                            env=the_env,
                                            start_new_session=True)
            except Exception:
                raise RuntimeError("Unable to launch backend using workload manager launch")

        # in case of SSH, we need to loop over each host to launch our backend and
        # the bundle up all the processes into a monitoring object that allows
        # the frontend server to query only one object for its monitoring needs
        else:
            try:
                popen_dict = {}

                for host in nodelist:
                    args = dlutil.get_wlm_launch_args(args_map=self.args_map,
                                                      hostname=host,
                                                      wlm=self._wlm)
                    args = args + be_args
                    log.info(f"SSH launch config: {args}")

                    popen_dict[host] = subprocess.Popen(
                                           args=args,
                                           stdin=subprocess.DEVNULL,
                                           stdout=subprocess.DEVNULL,
                                           env=the_env,
                                           start_new_session=True
                                       )
                wlm_proc = SSHSubprocessPopen(popen_dict)

            except Exception:
                raise RuntimeError("Unable to launch backend via SSH loop")

        return wlm_proc

    def run_startup(self):
        """Complete bring up of runtime services
        """
        log = logging.getLogger(dls.LA_FE).getChild('run_startup')


        # This is set here for the overlay network.
        this_process.set_num_gateways_per_node(dfacts.DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE)

        # setup sigint handling to proceed in a manner that works for us
        self._STATE = FrontendState.NET_CONFIG

        try:
            self._orig_sigint = signal.signal(signal.SIGINT, self.sigint_handler)
            log.debug("got signal handling in place")
        except ValueError:
            # this error is thrown if we are running inside a child thread
            # which we do for unit tests. So pass on this
            log.debug("Unable to do signal handling outside of main thread")

        self.la_fe_stdin = dlutil.OverlayNetLaFEQueue()
        self.la_fe_stdout = dlutil.LaOverlayNetFEQueue()

        # Get the node config for the backend
        log.debug('Getting the node config for the backend.')
        if self._config_from_file is not None:
            try:
                log.info(f"Acquiring network config from file {self._config_from_file}")
                self.net = NetworkConfig.from_file(self._config_from_file)
            except Exception:
                raise RuntimeError("Unable to acquire backend network configuration from input file.")
        else:
            try:
                log.info("Acquiring network config via WLM queries")
                # This sigint trigger is -2 and -1 cases
                self.net = NetworkConfig.from_wlm(workload_manager=self._wlm,
                                                  port=dfacts.DEFAULT_OVERLAY_NETWORK_PORT,
                                                  network_prefix=dfacts.DEFAULT_TRANSPORT_NETIF,
                                                  hostlist=self.hostlist,
                                                  sigint_trigger=self._sigint_trigger)
            except Exception:
                raise RuntimeError("Unable to acquire backend network configuration via workload manager")
        net_conf = self.net.get_network_config()
        log.debug(f"net_conf = {net_conf}")

        if self._sigint_trigger == 0:
            signal.raise_signal(signal.SIGINT)

        # Add the frontend config
        net_conf['f'] = NodeDescriptor.get_local_node_network_conf(network_prefix=self.network_prefix,
                                                                   port_range=dfacts.DEFAULT_FRONTEND_PORT)
        fe_host_id = str(net_conf['f'].host_id)
        fe_ip_addr = net_conf['f'].ip_addrs[0]  # it includes the port
        log.debug(f'node config: {net_conf}')

        # Create my memory pool
        conn_options = ConnectionOptions(min_block_size=2 ** 16)
        conn_policy = POLICY_INFRASTRUCTURE

        try:
            # Create my memory pool
            self.fe_mpool = MemoryPool(int(dfacts.DEFAULT_SINGLE_DEF_SEG_SZ),
                                       f'{os.getuid()}_{os.getpid()}_{fe_host_id}' + dfacts.DEFAULT_POOL_SUFFIX,
                                       dfacts.FE_OVERLAY_TRANSPORT_AGENT_MUID)
            puid, mpool_fname = MemoryPool.serialized_uid_fname(self.fe_mpool.serialize())
            log.debug(f'fe_mpool has uid {puid} and file {mpool_fname}')
            # Create my receiving channel
            fe_cuid = dfacts.FE_CUID
            local_in_cuid = dfacts.FE_LOCAL_IN_CUID
            local_out_cuid = dfacts.FE_LOCAL_OUT_CUID
            gw_cuid = dfacts.FE_GW_CUID

            # Channel for backend to come to
            self.fe_inbound = Channel(self.fe_mpool, fe_cuid)
            encoded_inbound = B64(self.fe_inbound.serialize())
            encoded_inbound_str = str(encoded_inbound)
            self.conn_in = Connection(inbound_initializer=self.fe_inbound,
                                      options=conn_options,
                                      policy=conn_policy)

            if self._sigint_trigger == 1:
                signal.raise_signal(signal.SIGINT)

            # Backdoor connection for breaking recv_msgs thread out of
            # its blocking receive at teardown
            self.conn_in_bd = Connection(outbound_initializer=self.fe_inbound)

            # Channel for tcp to tell me it's up
            self.local_ch_in = Channel(self.fe_mpool, local_in_cuid)
            self.local_ch_out = Channel(self.fe_mpool, local_out_cuid)
            self.local_inout = Connection(inbound_initializer=self.local_ch_in,
                                          outbound_initializer=self.local_ch_out)

            # Create a gateway and logging channel for my tcp agent
            self.gw_ch = Channel(self.fe_mpool, gw_cuid)
            self.dragon_logger = DragonLogger(self.fe_mpool)
        except (ChannelError, DragonPoolError, DragonLoggingError,
                DragonMemoryError) as init_err:
            log.fatal(f'could not create resources: {init_err}')
            raise RuntimeError('overlay transport resource creation failed') from init_err

        log.info('Memory pools and channels created')

        if self._sigint_trigger == 2:
            signal.raise_signal(signal.SIGINT)
        # Set gateway in my environment and register them
        encoded_ser_gw = B64(self.gw_ch.serialize())
        encoded_ser_gw_str = str(encoded_ser_gw)
        os.environ[dfacts.GW_ENV_PREFIX + str(dfacts.DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE)] = encoded_ser_gw_str
        register_gateways_from_env()

        # start my transport agent
        nnodes = len(net_conf) - 1  # Exclude the frontend node from this
        log.debug(f'requested {self.nnodes} and got {nnodes}')
        if self.nnodes > 0:
            if self.nnodes > nnodes:
                log.exception('too many nodes requested')
                raise ValueError('Not enough backend nodes allocated to match requested')
            nnodes = self.nnodes
        else:
            self.nnodes = nnodes

        log.debug(f'main has {nnodes} nodes')

        # Acquire the primary node and add
        # the frontend info to host_ids and ip_addrs, so the TA functions
        host_ids = [str(net_conf['0'].host_id)] + [fe_host_id]
        ip_addrs = [net_conf['0'].ip_addrs[0]] + [fe_ip_addr]
        hostnames = [net_conf['0'].name]

        # Add as many as needed to meet the requested node count.
        # We also find the minimum number of network interface cards
        # per node.
        min_nics_per_node = 99999
        for node_index in range(nnodes):
            sindex = str(node_index)
            min_nics_per_node = min(min_nics_per_node, len(net_conf[sindex].ip_addrs))
            if not net_conf[sindex].is_primary:
                host_ids.append(str(net_conf[sindex].host_id))
                ip_addrs.append(net_conf[sindex].ip_addrs[0])
                hostnames.append(net_conf[sindex].host_name)

        log.debug(f"ip_addrs={ip_addrs}, host_ids={host_ids}")
        log.debug(f'Found {min_nics_per_node} NICs per node.')

        log.debug(f'standing up tcp agent with gw: {encoded_ser_gw_str}')

        self._STATE = FrontendState.OVERLAY_STARTING

        try:
            self.over_proc = start_overlay_network(ch_in_sdesc=B64(self.local_ch_out.serialize()),
                                                   ch_out_sdesc=B64(self.local_ch_in.serialize()),
                                                   log_sdesc=B64(self.dragon_logger.serialize()),
                                                   host_ids=host_ids,
                                                   ip_addrs=ip_addrs,
                                                   frontend=True,
                                                   env=os.environ)
        except Exception as e:
            log.fatal('transport agent launch failed on FE')
            raise RuntimeError('Overlay transport agent launch failed on launcher frontend') from e

        # Wait on a started message
        self._wait_on_overlay_init()

        # Start a thread for monitoring messages into the logging channel
        # so we can see what comes out of the overlay network
        self.recv_logs_from_overlaynet_thread = threading.Thread(name='Logging Monitor',
                                                                 target=self.recv_log_msgs_from_overlaynet,
                                                                 daemon=False)
        self.recv_logs_from_overlaynet_thread.start()

        # Start the backend
        log.debug("standing up backend")
        # we need to send to backend only the ip_addr and host_id of the frontend
        log.debug(f"fe_ip_addr={fe_ip_addr}, fe_host_id={fe_host_id}")

        # Send/recv data with my OverlayNet FE server
        send_overlaynet_args = (self.la_fe_stdout,)
        self.send_overlaynet_thread = threading.Thread(name='OverlayNet Sender',
                                                       target=self._send_msgs_to_overlaynet,
                                                       args=send_overlaynet_args,
                                                       daemon=False)
        self.send_overlaynet_thread.start()
        log.info('started overlaynet sender thread')
        if self._sigint_trigger == 3:
            signal.raise_signal(signal.SIGINT)

        recv_overlaynet_args = (self.la_fe_stdin,)
        self.recv_overlaynet_thread = threading.Thread(name='OverlayNet Receiver',
                                                       target=self.recv_msgs_from_overlaynet,
                                                       args=recv_overlaynet_args,
                                                       daemon=False)
        self.recv_overlaynet_thread.start()

        if self._sigint_trigger == 4:
            signal.raise_signal(signal.SIGINT)

        # Start the backend
        log.debug("standing up backend")
        # we need to send to backend only the ip_addr and host_id of the frontend
        log.debug(f"fe_ip_addr={fe_ip_addr}, fe_host_id={fe_host_id}")

        self._STATE = FrontendState.STARTUP

        try:
            self.wlm_proc = self._launch_backend(nnodes=nnodes,
                                                 nodelist=hostnames,
                                                 fe_ip_addr=fe_ip_addr,
                                                 fe_host_id=fe_host_id,
                                                 frontend_sdesc=encoded_inbound_str,
                                                 network_prefix=self.network_prefix)
        except Exception as e:
            log.fatal('FE failed to stand up BE')
            log.debug(f'error: {e}')
            raise RuntimeError('Backend launch failed from launcher frontend') from e

        log.debug("LA BE started on each compute node")
        self._STATE = FrontendState.BACKEND_LAUNCHED

        if self._sigint_trigger == 5 or self._sigint_trigger == 8:
            signal.raise_signal(signal.SIGINT)

        # Receive BEIsUp msg - Try getting a backend channel descriptor
        be_ups = [dlutil.get_with_blocking(self.la_fe_stdin) for _ in range(nnodes)]
        assert len(be_ups) == self.nnodes

        self.conn_outs = {}  # key is the node_index and value is the Connection object
        for be_up in be_ups:
            assert isinstance(be_up, dmsg.BEIsUp), 'la_fe received invalid be up'
            log.debug(f'received descriptor: {be_up.be_ch_desc} and host_id: {be_up.host_id}')
            try:
                # TODO: VERIFY WE GET A UNIQUE CUID: keep a set of seen cuids.
                # After attaching, get the cuid and compare it against the set
                # of already seen cuids. Throw an exception if already seen.
                # Delete the set after this loop.
                be_sdesc = B64.from_str(be_up.be_ch_desc)
                be_ch = Channel.attach(be_sdesc.decode(), mem_pool=self.fe_mpool)
                conn_options = ConnectionOptions(default_pool=self.fe_mpool, min_block_size=2 ** 16)
                conn_out = Connection(outbound_initializer=be_ch,
                                      options=conn_options,
                                      policy=conn_policy)
                conn_out.ghost = True
            except ChannelError as ex:
                log.fatal(f'could not connect to BE channel with host_id {be_up.host_id}')
                raise RuntimeError('Connection with BE failed') from ex

            # Send node_index to backend - FENodeIdxBE msg
            for key, node_desc in net_conf.items():
                if str(be_up.host_id) == str(node_desc.host_id):
                    fe_node_idx = dmsg.FENodeIdxBE(tag=dlutil.next_tag(), node_index=key)
                    conn_out.send(fe_node_idx.serialize())
                    self.conn_outs[int(key)] = conn_out
                    log.debug(f"this is node_index: {key} with host_id: {be_up.host_id}")
                    break
        del be_ups
        log.info(f'received {nnodes} BEIsUp msgs and sent {nnodes} FENodeIdxBE msgs')

        chs_up = [dlutil.get_with_blocking(self.la_fe_stdin) for _ in range(nnodes)]
        for ch_up in chs_up:
            assert isinstance(ch_up, dmsg.SHChannelsUp), 'la_fe received invalid channel up'
        log.info(f'received {nnodes} SHChannelsUP msgs')

        gs_cds = [ch_up.gs_cd for ch_up in chs_up if ch_up.gs_cd is not None]
        if len(gs_cds) == 0:
            print('The Global Services CD was not returned by any of the SHChannelsUp messages. Launcher Exiting.')
            sys.exit(LAUNCHER_FAIL_EXIT)
        gs_cd = gs_cds[0]

        # Set the number of gateway channels per node. When
        # HSTA is used, and the pals library is present, then
        # we are on a Shasta (EX) machine and we can then use
        # multi-nic support under HSTA when multiple NICs per
        # node are available. In this circumstance, the number of
        # gateway channels will control the number of HSTA Agent
        # procs created and HSTA uses the default MPI binding to
        # bind each agent to a separate NIC. All other configurations
        # will have one NIC per node and one gateway per node used
        # (for now anyway - eventually the number of gateways may not
        # be tied to multiple NIC support).
        #FIXME-MULTI-NIC: force one nic per node for now
        if (min_nics_per_node > 1) and self.pals_lib_present and (self.transport is dfacts.TransportAgentOptions.HSTA) and False:
            num_gw_channels = min_nics_per_node
        else:
            num_gw_channels = 1

        # Send LAChannelsInfo in a test environment or to all
        la_ch_info = dmsg.LAChannelsInfo(tag=dlutil.next_tag(), nodes_desc=chs_up,
                                         gs_cd=gs_cd, num_gw_channels=num_gw_channels,
                                         port=self.port,
                                         transport=str(self.transport))
        log.debug(f'la_fe la_channels.nodes_desc: {la_ch_info.nodes_desc}')
        log.debug(f'la_fe la_channels.gs_cd: {la_ch_info.gs_cd}')
        log.debug(f'la_fe la_channels.transport: {la_ch_info.transport}')
        self.la_fe_stdout.send("A", la_ch_info.serialize())
        log.info('sent LACHannelsInfo to overlaynet fe')

        self.tas_up = [dlutil.get_with_blocking(self.la_fe_stdin) for _ in range(nnodes)]
        for ta_up in self.tas_up:
            assert isinstance(ta_up, dmsg.TAUp), 'la_fe received invalid channel up'
        log.info(f'received {nnodes} TAUp messages')

        if self._sigint_trigger == 6:
            signal.raise_signal(signal.SIGINT)

        if not self.transport_test_env:
            log.info('Now waiting on getting GSIsUp....')
            gs_up = dlutil.get_with_blocking(self.la_fe_stdin)
            assert isinstance(gs_up, dmsg.GSIsUp), 'la_fe expected GSIsUp msg'
            log.info('la_fe received GSIsUp. Prepping launch of user application')

        # Infrastructure is up
        self._STATE = FrontendState.STOOD_UP

    def run_app(self):
        """Start user app execution via GSProcessCreate or SHProcessCreate
        """
        self._STATE = FrontendState.APP_EXECUTION
        log = logging.getLogger(dls.LA_FE).getChild('run_app')

        make_inf_channels = True

        # Send message to start user application
        if self.transport_test_env:
            # start a process on each shepherd with an SHProcessCreate broadcast to all shepherds
            # build a dictionary mapping each node_index to the test channels created on that node.
            node_idx_to_channels_map = dict((msg.idx, msg.test_channels) for msg in self.tas_up)
            start_msg = dlutil.mk_shproc_start_msg(logbase=dls.LA_FE, stdin_str=str(node_idx_to_channels_map))
            self.la_fe_stdout.send('A', start_msg.serialize())
            log.info(f'broadcast the program to start to all shepherds with message={start_msg}')
        else:
            # Send GS user proc start msg
            log.info('Prepping launch of user application')
            start_msg = dlutil.mk_head_proc_start_msg(logbase=dls.LA_FE,
                                                      make_inf_channels=make_inf_channels,
                                                      args_map=self.args_map)
            self.la_fe_stdout.send("P", start_msg.serialize())
            if self._sigint_trigger == 7:
                signal.raise_signal(signal.SIGINT)
            log.info('transmitted GSProcessCreate')

    def run_msg_server(self):
        """Process messages from backend after user app starts
        """
        self.msg_log = logging.getLogger(dls.LA_FE).getChild('run_msg_server')
        running = True
        execute_teardown = True

        while running:

            msg = dlutil.get_with_blocking(self.la_fe_stdin)

            if hasattr(msg, 'r_c_uid'):
                msg.r_c_uid = dfacts.launcher_cuid_from_index(self.node_idx)
            self.msg_log.info(f"received {type(msg)}")

            try:
                if type(msg) in LauncherFrontEnd._DTBL:
                    self._DTBL[type(msg)][0](self, msg=msg)
                else:
                    # TODO: Make sure handling abnormal termniation works this way
                    self.msg_log.warning(f'unexpected msg type: {repr(msg)}')
            except LauncherImmediateExit:
                self.msg_log.debug(f'completing a sys.exit({LAUNCHER_FAIL_EXIT})')
                sys.exit(LAUNCHER_FAIL_EXIT)
            except SigIntImmediateExit:
                self.msg_log.debug('completing an exit due to SIGINT')
                execute_teardown = False
                running = False
                break

            # Break and wait on threads to exit
            if self._sh_halt_be == self.nnodes:
                running = False

        self.msg_log.debug('out of msg_server loop')
        if execute_teardown:
            self.msg_log.info('joining on message threads')
            self._close_threads(abnormal=self._abnormal_termination.is_set())
            self.msg_log.info('la_fe has shutdown overlaynet send/recv threads')

            # Waiting on teardown of wlm launched backend
            self.wlm_proc.wait()

            # TEARDOWN
            self.msg_log.info('WLM launched backend down. Tearing down comm infra')
            self._close_comm_overlaynet()
            self.msg_log.info("Leaving run_msg_server")

        if self._sigint_count > 0:
            self.msg_log.debug('raising keyboard interrupt')
            raise KeyboardInterrupt

        # Set state so our exit method correctly executes.
        self._STATE = FrontendState.LAUNCHER_DOWN

    def probe_teardown(self):
        """Check on whether to begin teardown based on received backend messages"""
        # Global services is up and we're using it
        if self._gs_head_exit_received and self._gs_process_create_resp_received:
            # m4.1 Send GSTeardown to Primary BE

            # Only update this state if we're not in an abnormal teardown. Too many
            # subsequent state checks depend on that differentiation
            if self._STATE != FrontendState.ABNORMAL_TEARDOWN:
                self._STATE = FrontendState.TEARDOWN

            gs_teardown = dmsg.GSTeardown(tag=dlutil.next_tag())
            self.la_fe_stdout.send("P", gs_teardown.serialize())
            self.msg_log.info('m4.1 la_fe transmitted teardown msg to BE')
        # No global services so it's being overlooked
        elif self.transport_test_env:
            if self._proc_create_resps == self.nnodes and self._proc_exits == self.nnodes:
                # m7.1 Send SHHaltTA to All BEs
                self._STATE = FrontendState.TEARDOWN
                sh_halt_ta = dmsg.SHHaltTA(tag=dlutil.next_tag())
                self.la_fe_stdout.send("A", sh_halt_ta.serialize())
                self.msg_log.info('m7.1 transmitted SHHaltTA msg to BE')
        else:
            self.msg_log.debug(f'gs_head_exit: {self._gs_head_exit_received} | gs_proc_create_resp: {self._gs_process_create_resp_received}')

    @route(dmsg.SHFwdOutput, _DTBL)
    def handle_sh_fwd_output(self, msg: dmsg.SHFwdOutput):
        msg_out = self.build_stdmsg(msg, self.args_map,
                                    msg.fd_num == dmsg.SHFwdOutput.FDNum.STDOUT.value)
        print(msg_out, end="")

    @route(dmsg.LAExit, _DTBL)
    def handle_la_exit(self, msg: dmsg.LAExit):
        self.msg_log.debug('Received LAExit message. Will break out of message loop')
        if msg.sigint:
            raise SigIntImmediateExit('Exiting due to receipt of SIGINT')
        else:
            raise LauncherImmediateExit(f'Exiting due to reception of {type(msg)}')

    @route(dmsg.ExceptionlessAbort, _DTBL)
    def handle_exceptionless_abort(self, msg: dmsg.ExceptionlessAbort):
        self.msg_log.debug('exceptionless abort routing')
        self._cleanup_abnormal_state()

    @route(dmsg.GSHeadExit, _DTBL)
    def handle_gs_head_exit(self, msg: dmsg.GSHeadExit):
        self.msg_log.info('The head process has exited')
        self._gs_head_exit_received = True
        self.probe_teardown()

    @route(dmsg.GSProcessCreateResponse, _DTBL)
    def handle_gs_proc_create_response(self, msg: dmsg.GSProcessCreateResponse):
        self.msg_log.info('The GSProcessCreateResponse was received in the launcher front end')
        self._gs_process_create_resp_received = True
        self.probe_teardown()

    @route(dmsg.SHProcessCreateResponse, _DTBL)
    def handle_sh_proc_create_response(self, msg: dmsg.SHProcessCreateResponse):
        self._proc_create_resps += 1
        self.msg_log.info(f'Got {self._proc_create_resps} SHProcessCreateResponse messages')
        self.probe_teardown()

    @route(dmsg.SHProcessExit, _DTBL)
    def handle_sh_proc_exit(self, msg: dmsg.SHProcessExit):
        self._proc_exits += 1
        self.msg_log.info(f'Got {self._proc_exits} SHProcessExit messages')
        self.probe_teardown()

    @route(dmsg.GSHalted, _DTBL)
    def handle_gs_halted(self, msg: dmsg.GSHalted):
        # m6.2 recv GSHalted from Primary BE
        self.msg_log.info('m6.2 la_fe received GSHalted')

        # m7.1 Send SHHaltTA to All BEs
        sh_halt_ta = dmsg.SHHaltTA(tag=dlutil.next_tag())

        if self._STATE != FrontendState.ABNORMAL_TEARDOWN:
            self.la_fe_stdout.send("A", sh_halt_ta.serialize())
            self.msg_log.info(f'm7.1 transmitted SHHaltTA msg to BE via threads: {self._STATE}')
        else:
            self._overlay_bcast(sh_halt_ta)
            self.msg_log.info('m7.1 transmitted SHHaltTA msg to BE directly')

    @route(dmsg.TAHalted, _DTBL)
    def handle_ta_halted(self, msg: dmsg.TAHalted):
        self._tas_halted += 1
        if self._tas_halted == self.nnodes:
            self.msg_log.info('m10.2 la_fe received All TAHalted messages')
            sh_teardown = dmsg.SHTeardown(tag=dlutil.next_tag())

            if self._STATE != FrontendState.ABNORMAL_TEARDOWN:
                # m11.1 Send SHTeardown to all BE's
                self.la_fe_stdout.send("A", sh_teardown.serialize())
                self.msg_log.info('m11.1 la_fe transmitted SHTeardown msg to BE via threads')
            else:
                self._overlay_bcast(sh_teardown)
                self.msg_log.info('m11.1 la_fe transmitted SHTeardown msg to BE directly')

            # Handle one of my sigint test cases
            try:
                if self._sigint_trigger == 9:
                    self.msg_log.debug("raising signal during teardown")
                    self._sigint_trigger = None
                    signal.raise_signal(signal.SIGINT)
                    self.msg_log.debug("SIGINT raised teardown")
            except TypeError:
                pass

    @route(dmsg.SHHaltBE, _DTBL)
    def handle_sh_halt_be(self, msg: dmsg.SHHaltBE):
        self._sh_halt_be += 1
        if self._sh_halt_be == self.nnodes:
            self.msg_log.info('Received all SHHaltBE messages')

            be_halted = dmsg.BEHalted(tag=dlutil.next_tag())

            if self._STATE != FrontendState.ABNORMAL_TEARDOWN:
                self.la_fe_stdout.send("A", be_halted.serialize())
                self.msg_log.info('m14 transmitted BEHalted msg to BE to be forwarded to local services on all nodes via threads')
                self._shutdown.set()

                # signal the end of the logging thread by sending its way
                # an halt message
                msg = dmsg.HaltLoggingInfra(tag=dlutil.next_tag())
                self.dragon_logger.put(msg.serialize(), logging.INFO)

                self.conn_in_bd.send(be_halted.serialize())
                self.msg_log.info('sent BEHalted to break overlay recv out of loop')

            else:
                self._overlay_bcast(be_halted)
                self.msg_log.info(f'm14 transmitted BEHalted msg to BE to be forwarded to local services on all nodes directly')

    def recv_log_msgs_from_overlaynet(self):
        """This service grabs messages from overlay logger and logs them in this

        Args:
            level (int): minimum log priority of messages to forward
        """
        _, level = _get_dragon_log_device_level(os.environ, LOGGING_OUTPUT_DEVICE_DRAGON_FILE)
        log = logging.getLogger(dls.LA_FE).getChild('recv_overlay_logs')
        try:
            # Set timeout to None which allows better interaction with the GIL
            while not self._shutdown.is_set():
                try:
                    msg = dmsg.parse(self.dragon_logger.get(level, timeout=None))
                    if isinstance(msg, dmsg.HaltLoggingInfra):
                        break
                    log = logging.getLogger(msg.name)
                    log.log(msg.level, msg.msg, extra=msg.get_logging_dict())
                except ChannelEmpty:
                    pass

        except Exception as ex:  # pylint: disable=broad-except
            log.warning(f'Caught exception {ex} in logging thread')
        log.info("exiting recv log msg thread")

    def send_msg_to_overlaynet(self, target, msg):
        """Send single message to backend

        Args:
            target (str): "A" for all backend nodes. "P" for primary node only
            msg (dmsg._MsgBase): non-serialized message to send
        """
        try:
            self.la_fe_stdout.send(target, msg.serialize())
        except Exception:
            raise RuntimeError("Unable to send message to overlaynet send thread")

    def _overlay_bcast(self, msg: dmsg._MsgBase):
        '''Send bcast of message to all backend nodes via overlay network

        :param msg: Message to send
        :type msg: dmsg._MsgBase
        '''

        for conn_out in self.conn_outs.values():
            conn_out.send(msg.serialize())

    def _send_msgs_to_overlaynet(self, la_fe_stdout: dlutil.LaOverlayNetFEQueue):
        """Thread that send Messages from frontend to backend service

        Args:
            la_fe_stdout (dlutil.LaOverlayNetFEQueue): SimpleQueue sending messages from launcher
                FE main thread to this OverlayNet server thread
        """
        log = logging.getLogger(dls.LA_FE).getChild('send_msgs_to_overlaynet')

        try:
            # recv msg from parent thread to send to backend
            while not self._shutdown.is_set():
                target, fe_msg = dlutil.get_with_blocking(la_fe_stdout)
                if isinstance(fe_msg, dmsg.HaltOverlay):
                    self._shutdown.set()
                    break

                if target == "P":
                    # node index 0 is the primary node
                    self.conn_outs[0].send(fe_msg.serialize())
                    log.info(f'target = {target} -- msg = {fe_msg.__class__}')
                else:
                    self._overlay_bcast(fe_msg)
                    log.info(f'sent {fe_msg.__class__} down OverlayNet tree')

                if isinstance(fe_msg, dmsg.BEHalted):
                    log.info('overlaynet send thread got halt msg')
                    self._shutdown.set()
                    break
        except Exception as err:
            log.exception(f'overlaynet sending thread failure: {err}')

        # Let the OverlayNet receiving thread know we're gone, and it can exit
        self._shutdown.set()

        # signal the end of the logging thread by sending its way
        # a HaltLoggingInfra message
        msg = dmsg.HaltLoggingInfra(tag=dlutil.next_tag())
        self.dragon_logger.put(msg.serialize(), logging.INFO)

        log.info('overlaynet sending thread exiting ...')

    def recv_msgs_from_overlaynet(self,
                                  la_fe_stdin: dlutil.OverlayNetLaFEQueue):
        """"Start and run OverlayNet FrontEnd service

        Args:
            la_fe_stdin (dlutil.OverlayNetLaFEQueue): SimpleQueue sending messages from
                OverlayNetServer thread to launcher FE main thread
        """

        # Get the logger I will use to write my own log messages
        log = logging.getLogger(dls.LA_FE).getChild('recv_msgs_from_overlaynet')

        try:
            # Just recv messages and send them to my parent thread
            # until my parent thread says stop
            while not self._shutdown.is_set():
                # send msg from callback to parent

                # the get_with_timeout is a decorated function (i.e. wrapped function
                # in a function) that filters out all the log messages coming up through
                # overlaynet. So we don't worry about them here. See dlutil for details.
                try:
                    be_msg = dlutil.get_with_blocking(self.conn_in)
                except AbnormalTerminationError:
                    # I need to get this to the main thread without throwing an exception
                    self._abnormal_termination.set()
                    if self._STATE < FrontendState.TEARDOWN:
                        be_msg = dmsg.ExceptionlessAbort(tag=dlutil.next_tag())
                    else:
                        log.debug("Caught AbnormalTermination during teardown. Skipping it in hopes of completing teardown")
                        continue

                if isinstance(be_msg, dmsg.HaltOverlay):
                    self._shutdown.set()
                    break

                # We could also use decorator here, but not yet...
                if isinstance(be_msg, dmsg.BEHalted):
                    log.info('overlaynet recv thread got halt msg')
                    break
                else:
                    log.info(f'Got {be_msg.__class__} in the front end monitor and forwarding to launcher front end.')
                    la_fe_stdin.send(be_msg.serialize())

            # Let the OverlayNet sending thread know we're gone
            self._shutdown.set()

            # signal the end of the logging thread by sending its way
            # a Halt message
            msg = dmsg.HaltLoggingInfra(tag=dlutil.next_tag())
            self.dragon_logger.put(msg.serialize(), logging.INFO)

            log.info('Overlaynet recv thread exiting.')
        except Exception as err:
            log.exception(f'Overlaynet recving thread failure. Exiting: {err}')
            la_fe_stdin.send(dmsg.AbnormalTermination(tag=dlutil.next_tag()).serialize())
            log.debug('send abnormal term from except block')
            raise

    def build_stdmsg(self, msg: dmsg.SHFwdOutput, arg_map, is_stdout=True):
        if arg_map['no_label']:
            return f'{msg.data}'

        msg_str = ""
        if is_stdout:
            msg_str += "[stdout: "
        else:
            msg_str += "[stderr: "

        if arg_map['verbose_label']:
            msg_str += f"PID {msg.pid} @ {msg.hostname}]"
        elif arg_map['basic_label']:
            msg_str += f"Dragon PID {msg.p_uid}]"

        msg_str += f" {msg.data}"
        return msg_str

src/dragon/launcher/frontend.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import logging
import socket

from ..localservices.local_svc import multinode
from ..infrastructure.facts import GS_LOG_BASE, TRANSPORT_TEST_ENV, PROCNAME_GS
from ..dlogging.util import DragonLoggingServices, setup_BE_logging
from ..cli import console_script_args


def main():

    _, fname = setup_BE_logging(service=DragonLoggingServices.LS)
    log = logging.getLogger(DragonLoggingServices.LS).getChild('launch_multi_ls')
    log.info('--------------------ls started-----------------------')
    log.debug(f'starting debug mode on pid={os.getpid()}')

    for key, val in os.environ.items():
        log.debug(f'{key}: {val}')


    env_update = {GS_LOG_BASE: f'gs-{socket.gethostname()}'}

    # Discover whether the transport service test mode is being requested
    # or not.
    transport_test_flag = os.environ.get(TRANSPORT_TEST_ENV) is not None

    gs_args = console_script_args(PROCNAME_GS)

    try:
        _ = multinode(gs_args=gs_args, gs_env=env_update,
                      fname=fname,
                      transport_test_env=transport_test_flag)

    except Exception as ex:
        log.exception('Final launch exception catch')
        raise ex


if __name__ == "__main__":
    main()

src/dragon/launcher/launch_multi_ls.py
DragonHPC/dragon
import os
import re
import shutil


def determine_environment(args=None):
    from dragon.launcher import launchargs

    # Check if the user requested a specific launcher mode
    arg_map = launchargs.get_args(args)

    single_arg = arg_map['single_node_override']
    multi_arg = arg_map['multi_node_override']
    if single_arg and multi_arg:
        msg = """Cannot request both single node and multi-node launcher simultaneously.
Please specify only '--single-node-override' or '--multi-node-override'"""
        raise ValueError(msg)
    elif single_arg:
        return False
    elif multi_arg:
        return True

    # If no override specific, do automatic detection

    # Try to figure out if we're in a support multinode case
    multi_mode = False

    # Try to determine if we're on a supported multinode system
    pbs = shutil.which('qstat')
    is_pbs = re.match('.*/pbs/.*', str(pbs))

    is_slurm = shutil.which('srun')

    if is_pbs:
        # Next see if we have a PALS environment loaded -- ALPS isn't supported
        mpiexec = shutil.which("mpiexec")
        if mpiexec and is_pbs:
            pals_mpiexec = re.match('.*/pals/.*', mpiexec)
            if pals_mpiexec:
                multi_mode = True
            else:
                msg = """PBS has been detected on the system. However, Dragon is only
    compatible with PALS when using PBS. Please load or install PALS into your path to continue"""
                raise RuntimeError(msg)

            if not os.environ.get("PBS_NODEFILE"):
                msg = """Using a supported PALS with PBS config. However, no active jobs allocation
    has been detected. Resubmit as part of a 'qsub' execution."""
                raise RuntimeError(msg)

    elif is_slurm:
        multi_mode = True

        # Check if we have nodes allocated. If not, raise an
        # exception
        if not os.environ.get("SLURM_JOB_ID"):
            msg = """Executing in a Slurm environment, but with no job allocation.
    Resubmit as part of an 'salloc' or 'sbatch' execution"""
            raise RuntimeError(msg)

    return multi_mode


def get_launcher():
    multi_mode = determine_environment()

    if multi_mode:
        from dragon.launcher import dragon_multi_fe as dl
    else:
        from dragon.launcher import dragon_single as dl
    return dl


def main():
    '''
    .. code-block::

        usage: dragon [-h] [-N NODE_COUNT] [--hostlist HOSTLIST | --hostfile HOSTFILE] [--network-prefix NETWORK_PREFIX] [--network-config NETWORK_CONFIG]
                      [--wlm WORKLOAD_MANAGER] [-p PORT] [--transport TRANSPORT_AGENT] [-s | -m] [-l LOG_LEVEL] [--no-label] [--basic-label] [--verbose-label]
                      [--version] [PROG] ...

        Dragon Launcher Arguments and Options

        positional arguments:
          PROG                  PROG specifies an executable program to be run on the primary compute node. In this case, the file may be either executable or not.
                                If PROG is not executable, then Python version 3 will be used to interpret it. All command-line arguments after PROG are passed to
                                the program as its arguments. The PROG and ARGS are optional.
          ARG                   Zero or more program arguments may be specified.

        optional arguments:
          -h, --help            show this help message and exit
          -N NODE_COUNT, --nodes NODE_COUNT
                                NODE_COUNT specifies the number of nodes to use. NODE_COUNT must be less or equal to the number of available nodes within the WLM
                                allocation. A value of zero (0) indicates that all available nodes should be used (the default).
          --hostlist HOSTLIST   Specify backend hostnames as a comma-separated list, eg: `--hostlist host_1,host_2,host_3`. `--hostfile` or `--hostlist` is a
                                required argument for WLM SSH and is only used for SSH
          --hostfile HOSTFILE   Specify a list of hostnames to connect to via SSH launch. The file should be a newline character separated list of hostnames.
                                `--hostfile` or `--hostlist` is a required argument for WLM SSH and is only used for SSH
          --network-prefix NETWORK_PREFIX
                                NETWORK_PREFIX specifies the network prefix the dragon runtime will use to determine which IP addresses it should use to build
                                multinode connections from. By default the regular expression r'^(hsn|ipogif|ib)\d+$' is used -- the prefix for known HPE-Cray XC
                                and EX high speed networks. If uncertain which networks are available, the following will return them in pretty formatting: `dragon-
                                network-ifaddrs --ip --no-loopback --up --running | jq`. Prepending with `srun` may be necessary to get networks available on
                                backend compute nodes
          --network-config NETWORK_CONFIG
                                NETWORK_CONFIG specifies a YAML or JSON file generated via a call to the launcher's network config tool that successfully generated
                                a corresponding YAML or JSON file (eg: `dragon-network-config --output-to-yaml`) describing the available backend compute nodes
                                specified either by a workload manager (this is what the tool provides). Alternatively, one can be generated manually as is needed
                                in the case of ssh-only launch. An example with keywords and formatting can be found in the documentation
          --wlm WORKLOAD_MANAGER, -w WORKLOAD_MANAGER
                                Specify what workload manager is used. Currently supported WLMs are: slurm, pbs+pals, ssh
          -p PORT, --port PORT  PORT specifies the port to be used for multinode communication. By default, 7575 is used.
          --transport TRANSPORT_AGENT, -t TRANSPORT_AGENT
                                TRANSPORT_AGENT selects which transport agent will be used for backend node-to-node communication. By default, the TCP
                                transport agent (tcp) is selected. Currently supported agents are: hsta, tcp
          -s, --single-node-override
                                Override automatic launcher selection to force use of the single node launcher
          -m, --multi-node-override
                                Override automatic launcher selection to force use of the multi-node launcher
          -l LOG_LEVEL, --log-level LOG_LEVEL
                                The Dragon runtime enables the output of diagnostic log messages to multiple different output devices. Diagnotic log messages can be
                                seen on the Dragon stderr console, via a combined 'dragon_*.log' file, or via individual log files created by each of the Dragon
                                'actors' (Global Services, Local Services, etc). By default, the Dragon runtime disables all diagnostic log messaging. Passing one
                                of NONE, DEBUG, INFO, WARNING, ERROR, or CRITICAL to this option, the Dragon runtime will enable the specified log verbosity. When
                                enabling DEBUG level logging, the Dragon runtime will limit the stderr and combined dragon log file to INFO level messages. Each
                                actor's log file will contain the complete log history, including DEBUG messages. This is done to help limit the number of messages
                                sent between the Dragon frontend and the Dragon backend at scale. To override the default logging behavior and enable specific
                                logging to one or more Dragon output devices, the LOG_LEVEL option can be formatted as a keyword=value pair, where the KEYWORD is
                                one of the Dragon log output devices (stderr, dragon_file or actor_file), and the VALUE is one of NONE, DEBUG, INFO, WARNING, ERROR
                                or CRITICAL (eg `-l dragon_file=INFO -l actor_file=DEBUG`). Multiple -l|--log-level options may be passed to enable the logging
                                desired.
          --no-label
          --basic-label
          --verbose-label
          --version             show program's version number and exit
    '''
    from dragon import _patch_multiprocessing
    _patch_multiprocessing()
    dl = get_launcher()
    dl.main()


if __name__ == '__main__':
    main()

src/dragon/launcher/launch_selector.py
DragonHPC/dragon
import argparse
import os
import re

from .wlm import WLM
from ..dlogging import util as dlogutil
from ..infrastructure.facts import TransportAgentOptions


# The argparse module automatically formats the help strings below so you don't have to
# worry about formatting them when editing.

VERSION = f'''
Dragon Version {os.environ.get('DRAGON_VERSION', '(unknown)')}
'''

PROGRAM_HELP = '''PROG specifies an executable program to be run on the primary
compute node. In this case, the file may be either executable or not.
If PROG is not executable, then Python version 3 will be used to
interpret it. All command-line arguments after PROG are passed
to the program as its arguments. The PROG and ARGS are optional.'''

ARG_HELP = '''Zero or more program arguments may be specified.'''

SINGLE_MODE_HELP = '''Override automatic launcher selection to force use of the
single node launcher'''

MULTI_MODE_HELP = '''Override automatic launcher selection to force use of the
multi-node launcher'''

NODES_HELP = '''NODE_COUNT specifies the number of nodes to use. NODE_COUNT must
be less or equal to the number of available nodes within the WLM allocation.
A value of zero (0) indicates that all available nodes should be used (the default).'''

NETWORK_HELP = '''NETWORK_PREFIX specifies the network prefix the dragon runtime will
use to determine which IP addresses it should use to build multinode connections from.
By default the regular expression r'^(hsn|ipogif|ib)\d+$' is used -- the prefix for known
HPE-Cray XC and EX high speed networks. If uncertain which networks are available,
the following will return them in pretty formatting:
`dragon-network-ifaddrs --ip --no-loopback --up --running | jq`.
Prepending with `srun` may be necessary to get networks available on backend
compute nodes'''

NETWORK_CONFIG_HELP = '''NETWORK_CONFIG specifies a YAML or JSON file generated via a call to the
launcher's network config tool that successfully generated a corresponding YAML
or JSON file (eg: `dragon-network-config --output-to-yaml`)
describing the available backend compute nodes specified either by a workload
manager (this is what the tool provides). Alternatively, one can be generated
manually as is needed in the case of ssh-only launch. An example with keywords
and formatting can be found in the documentation'''

HOSTFILE_HELP = \
    '''
    Specify a list of hostnames to connect to via SSH launch. The file
    should be a newline character separated list of hostnames.
    `--hostfile` or `--hostlist` is a required argument for WLM SSH
    and is only used for SSH
    '''

HOSTLIST_HELP = \
    '''
    Specify backend hostnames as a comma-separated list, eg:
    `--hostlist host_1,host_2,host_3`.
    `--hostfile` or `--hostlist` is a required argument for WLM SSH
    and is only used for SSH
    '''

WLM_HELP = f"""Specify what workload manager is used. Currently supported WLMs are: {', '.join([wlm.value for wlm in WLM])}"""

# NOTE: port 7575 is hardcoded rather than grabbed from facts.py because we
# NOTE: DO NOT want to import dragon when launchargs is imported.
PORT_HELP = '''PORT specifies the port to be used for multinode communication. By
default, 7575 is used.'''

LOGGING_HELP = '''The Dragon runtime enables the output of diagnostic log messages to multiple
different output devices. Diagnotic log messages can be seen on the Dragon stderr
console, via a combined 'dragon_*.log' file, or via individual log files created by
each of the Dragon 'actors' (Global Services, Local Services, etc).

By default, the Dragon runtime disables all diagnostic log messaging.

Passing one of NONE, DEBUG, INFO, WARNING, ERROR, or CRITICAL to this option,
the Dragon runtime will enable the specified log verbosity. When enabling DEBUG
level logging, the Dragon runtime will limit the stderr and combined dragon log file
to INFO level messages. Each actor's log file will contain the complete log history,
including DEBUG messages. This is done to help limit the number of messages sent
between the Dragon frontend and the Dragon backend at scale.

To override the default logging behavior and enable specific logging to one or
more Dragon output devices, the LOG_LEVEL option can be formatted as a keyword=value
pair, where the KEYWORD is one of the Dragon log output devices (stderr, dragon_file
or actor_file), and the VALUE is one of NONE, DEBUG, INFO, WARNING, ERROR or
CRITICAL (eg `-l dragon_file=INFO -l actor_file=DEBUG`). Multiple -l|--log-level
options may be passed to enable the logging desired.'''

TRANSPORT_HELP = f'''TRANSPORT_AGENT selects which transport agent will be used for
backend node-to-node communication. By default, the TCP transport agent
is selected. Currently supported agents are: {', '.join([ta.value for ta in TransportAgentOptions])}'''


class SplitArgsAtComma(argparse.Action):
    def __call__(self, parser, namespace, values, option_string=None):
        setattr(namespace, self.dest, [v.strip() for v in values.split(',')])


def non_negative_int(value):
    nvalue = int(value)
    if nvalue < 0:
        raise argparse.ArgumentTypeError(f"{value} must be zero or a positive int value")
    return nvalue


def valid_port_int(value):
    port = int(value)
    if port < 1024 or port > 65536:
        raise argparse.ArgumentTypeError(f"{value} must be in port range 1024-65535")
    return port


def is_hostname_valid(hostname):
    """Confirm hostname conforms to POSIX rules"""
    if len(hostname) > 255:
        return False
    if hostname[-1] == ".":
        hostname = hostname[:-1]  # strip exactly one dot from the right, if present
    allowed = re.compile(r"(?!-)[A-Z\d-]{1,63}(?<!-)$", re.IGNORECASE)
    return all(allowed.match(x) for x in hostname.split("."))


def parse_hosts(hostlist, hostfile):
    """Parse the hostfile if one is given and hostnames are valid"""

    if hostlist is None and hostfile is None:
        raise argparse.ArgumentError(hostlist, "When using WLM SSH, hostlist, hostfile, or existing network configuration is required.")
    else:
        # parse the hostfile if its set
        if hostfile is not None:
            try:
                with open(hostfile, 'r') as f:
                    hostlist = [line.strip() for line in f.readlines()]
            except Exception:
                raise argparse.ArgumentError(hostfile, f'Unable to parse {hostfile}')

    # Do a basic sanity check that the hostnames make sense
    try:
        for host in hostlist:
            if not is_hostname_valid(host):
                raise ValueError(f"Hostname is invalid: {host}")
    except Exception:
        raise

    return hostlist


def get_parser():

    from .network_config import WLM_HELP

    parser = argparse.ArgumentParser(prog='dragon',
                                     description='Dragon Launcher Arguments and Options')
    parser.add_argument('-N', '--nodes', metavar='NODE_COUNT', dest='node_count',
                        type=non_negative_int, help=NODES_HELP)

    host_group = parser.add_mutually_exclusive_group()
    host_group.add_argument('--hostlist', action=SplitArgsAtComma, metavar='HOSTLIST', type=str, help=HOSTLIST_HELP)
    host_group.add_argument('--hostfile', type=str, metavar='HOSTFILE', help=HOSTFILE_HELP)

    parser.add_argument('--network-prefix', metavar='NETWORK_PREFIX',
                        type=str, help=NETWORK_HELP)
    parser.add_argument('--network-config', metavar='NETWORK_CONFIG',
                        type=str, help=NETWORK_CONFIG_HELP)
    parser.add_argument('--wlm', '-w', metavar='WORKLOAD_MANAGER', type=WLM.from_str, choices=list(WLM), help=WLM_HELP)
    parser.add_argument('-p', '--port', metavar='PORT', type=valid_port_int, help=PORT_HELP)

    parser.add_argument('--transport', '-t', metavar='TRANSPORT_AGENT',
                        type=TransportAgentOptions.from_str, choices=list(TransportAgentOptions), help=TRANSPORT_HELP)

    group = parser.add_mutually_exclusive_group()
    group.add_argument('-s', '--single-node-override',
                       action='store_true', help=SINGLE_MODE_HELP)
    group.add_argument('-m', '--multi-node-override', action='store_true', help=MULTI_MODE_HELP)

    parser.add_argument('-l', '--log-level', nargs=1, default=dict(), action=dlogutil.LoggingValue,
                        dest='log_device_level_map', metavar='LOG_LEVEL', help=LOGGING_HELP)
    parser.add_argument('--no-label', action='store_true', default=True)
    parser.add_argument('--basic-label', action='store_true')
    parser.add_argument('--verbose-label', action='store_true')

    parser.add_argument('--version', action='version', version=VERSION)

    parser.add_argument('prog', type=str, nargs='?', metavar='PROG', help=PROGRAM_HELP)
    parser.add_argument('args', type=str, metavar='ARG', nargs=argparse.REMAINDER,
                        default=[], help=ARG_HELP)

    parser.set_defaults(
        transport=TransportAgentOptions.TCP,
    )

    return parser


def get_args(args_input=None):
    try:
        parser = get_parser()
        if args_input is None:
            args = parser.parse_args()
        else:
            args = parser.parse_args(args_input)

        # Override default no-label if higher verbosity selected
        if args.basic_label or args.verbose_label:
            args.no_label = False

        return {key: value for key, value in vars(args).items() if value is not None}

    except Exception:
        raise


def head_proc_args(args_map=None):
    if args_map is None:
        args_map = get_args()
    head_proc = args_map.get('prog')
    head_args = [head_proc]
    head_args.extend(args_map.get('args'))
    return head_proc, head_args

src/dragon/launcher/launchargs.py
DragonHPC/dragon
import os
import argparse
import logging
import json
import enum
import zlib
from typing import Optional
from base64 import b64encode, b64decode

from ..infrastructure.node_desc import NodeDescriptor
from ..infrastructure.facts import DEFAULT_TRANSPORT_NETIF, DEFAULT_OVERLAY_NETWORK_PORT
from ..dlogging import util as dlog


from .launchargs import HOSTFILE_HELP, HOSTLIST_HELP, WLM_HELP
from .launchargs import SplitArgsAtComma, parse_hosts
from .wlm import WLM
from .wlm import SlurmNetworkConfig, PBSPalsNetworkConfig, SSHNetworkConfig


@enum.unique
class ConfigOutputType(enum.Enum):
    """Enumerated list of supported network config I/O formats"""
    JSON = 'json'
    YAML = "yaml"

    def __str__(self):
        return self.value


class NetworkConfig():
    """Class for constructing network configurations for backend compute"""

    wlm_cls_dict = {
        WLM.SLURM: SlurmNetworkConfig,
        WLM.PBS_PALS: PBSPalsNetworkConfig,
        WLM.SSH: SSHNetworkConfig
    }

    def __init__(self):
        """Initializer for class. Assumes a Slurm workload manager presently"""

        self.network_config = {}

    def get_network_config(self):
        """Get the object's stored network configuration

        Returns:
            dict: network configuration with keys of node indices and values
                  dragon.infrastructure.NodeDescriptor
        """
        return self.network_config

    def output_to_file(self,
                       basename: str = "network-config",
                       basedir: str = None,
                       output_type: ConfigOutputType = ConfigOutputType.YAML):
        """Output network config to formatted file

        Args:
            basename (str, optional): Basename of output file. Defaults to "network-config".
            basedir (str, optional): Directory to output file. Defaults to None and uses
                                     current working directory
            output_type (ConfigOutputType, optional): Format to output to. Defaults to ConfigOutputType.YAML.
        """

        if basedir is None:
            basedir = os.getcwd()

        filename = os.path.join(basedir, ".".join([basename, str(output_type)]))

        if output_type == ConfigOutputType.JSON:
            with open(filename, 'w+') as f:
                from json import dumps
                f.write(dumps(self.serialize(), indent=4))

        elif output_type == ConfigOutputType.YAML:
            from yaml import dump
            with open(filename, 'w+') as f:
                f.write(dump(self.serialize(), indent=4))

    def serialize(self):
        """Serialize network config

        Returns:
            dict: network config with keys of node indices and values of serialized dicts
                  of dragon.infrastructure.NodeDescriptor
        """
        sdesc = {key: val.get_sdict() for key, val in self.network_config.items()}
        return sdesc

    def compress(self):
        """Serialize and compress the network configuration"""

        sdesc = self.serialize()
        compressed = b64encode(zlib.compress(json.dumps(sdesc).encode('utf-8'))).decode('ascii')
        return compressed

    def decompress(self, compressed_config):
        """Decompress a compressed serialized dictionary network config

        :param compressed_config: Compressed seralized NetworkConfig dictionary. Returned by
                                  NetworkConfig.compress() given an already existing object
        :type compressed_config: str
        """
        sdict = json.loads(zlib.decompress(b64decode(compressed_config)))
        return sdict

    def _initialize_from_sdict(self, serialized_dict):
        """Create a network config from seralized dictionary of one

        Args:
            serialized_dict (dict): network config with keys of node indices
                                    and values of serialized dicts
                                    of dragon.infrastructure.NodeDescriptor
        """
        self.network_config = {}
        for key, val in serialized_dict.items():
            self.network_config[key] = NodeDescriptor.from_sdict(val)

    def _set_primary_index_to_zero(self, ref_key):
        """Make sure the primary node is indexed to 0"""
        if int(ref_key) != 0:
            old_zero = self.network_config['0']
            self.network_config['0'] = self.network_config[ref_key]
            self.network_config['0'].is_primary = True
            self.network_config[ref_key] = old_zero
        else:
            self.network_config['0'].is_primary = True

    def _select_primary(self, primary_hostname: Optional[str] = None):
        """Select a backend node to be primary

        Args:
            primary_hostname (str, optional): Hostname of desired node to be primary.
                Defaults to None in which case node with smallest host ID is selected

        Raises:
            RuntimeError: If for any reason unable to select a node for primary host
        """
        if primary_hostname is None:
            # Find the index with minimum host ID
            try:
                primary_key = min(self.network_config, key=lambda k: self.network_config[k].host_id)
                self._set_primary_index_to_zero(primary_key)

                # If key is not 0, make it going forward
            except Exception:
                raise RuntimeError("Unable to select primary host via host ID")
        else:
            # Match given input to those in our network config
            try:
                key = [k for k, v in self.network_config.items() if v.name == primary_hostname]
                if len(key) == 0:
                    raise RuntimeError('Input hostname does not match any available')
                elif len(key) > 1:
                    raise RuntimeError("Invalid network configuration. hostname's are not all unique")
                else:
                    self._set_primary_index_to_zero(key[0])

            except Exception:
                raise

    @classmethod
    def from_wlm(cls,
                 workload_manager: WLM = WLM.SLURM,
                 network_prefix: str = DEFAULT_TRANSPORT_NETIF,
                 port: str = DEFAULT_OVERLAY_NETWORK_PORT,
                 primary_hostname: Optional[str] = None,
                 hostlist: Optional[list[str]] = None,
                 sigint_trigger=None):
        """Obtain a network configuration for a given worklaod manager

        Args:
            workload_manager (WLM, optional): Workload manager to be used to obtain info.
                Defaults to WLM.SLURM.
            network_prefix (str, optional): Network prefix to select for backend to frontend
                communication. Defaults to dragon.infrastructure.facts.DEFAULT_TRANSPORT_NETIF.
            port (str, optional): Port to use for listening. Defaults to
                dragon.infrastructure.facts.DEFAULT_OVERLAY_NETWORK_PORT.
            primary_hostname (str, optional): Hostname desired to be selected as primary. Defaults
                to None, in which case host with smallest host ID is selected.
            sigint_trigger (int, optional): Conditional location to raise a SIGINT for testing
                purposes

        Returns:
            NetworkConfig: object with complete network configuration
        """

        obj = cls()
        wlm_generator = obj.wlm_cls_dict.get(workload_manager)(network_prefix, port, hostlist)
        obj.network_config = wlm_generator.get_network_config(sigint_trigger=sigint_trigger)

        # If coming from the workload manager, we need to use the host ID to select the primary
        # node (unless overruled by user input) and assign a node index
        obj._select_primary(primary_hostname=primary_hostname)

        return obj

    @classmethod
    def from_file(cls,
                  filename: str,
                  format: ConfigOutputType = ConfigOutputType.YAML,
                  primary_hostname: Optional[str] = None):
        """Read input from a YAML or JSON network configuration file

        Args:
            filename (str): name of file containing configuration
            format (ConfigOutputType, optional): YAML or JSON structure of file. Defaults to ConfigOutputType.YAML.
            primary_hostname (str, optional): Node to be selected as primary if configuration
                doesn't have one already selected. Defaults to None.

        Raises:
            RuntimeError: if primary host is requested but already set in file or if
                file is not parsible

        Returns:
            NetworkConfig: object with complete network configuration
        """

        obj = cls()
        try:
            # Don't assume we've been told the correct format
            clean_read = False

            if not clean_read:
                with open(filename, 'r') as f:
                    try:
                        from json import load
                        sdict = load(f)
                        clean_read = True
                    except Exception:
                        pass

            if not clean_read:
                with open(filename, 'r') as f:
                    try:
                        from yaml import safe_load
                        sdict = safe_load(f)
                        clean_read = True
                    except Exception:
                        pass

            if not clean_read:
                raise RuntimeError("Unable to parse config file in JSON or YAML format")

            for key, value in sdict.items():
                obj.network_config[str(key)] = NodeDescriptor.from_sdict(value)

            # Confirm a primary host is selected. If not, set.
            has_primary = any(p.is_primary is True for p in obj.network_config.values())
            if not has_primary:
                obj._select_primary(primary_hostname=primary_hostname)
            elif has_primary and primary_hostname is not None:
                raise RuntimeError('Primary hostname input by user but is already set in network config file')

        except Exception:
            raise

        return obj

    @classmethod
    def from_sdict(cls, serialized_dict):
        """Initialize configuration from serialized dictionary

        Args:
            serialized_dict (dict): network config with keys of node indices
                                    and values of serialized dicts
                                    of dragon.infrastructure.NodeDescriptor

        Returns:
            NetworkConfig: object with complete network configuration
        """
        obj = cls.__new__(cls)
        obj._initialize_from_sdict(serialized_dict)
        return obj


def deliver_backend_node_descriptor(network_prefix=DEFAULT_TRANSPORT_NETIF,
                                    port=DEFAULT_OVERLAY_NETWORK_PORT):
    """Print to stdout the node descriptor of node being executed on

    Args:
        network_prefix (str, optional): Network prefix to use for IP address selection.
            Defaults to dragon.infrastructure.facts.DEFAULT_TRANSPORT_NETIF.
        port (int, optional): Listening port. Defaults to DEFAULT_OVERLAY_NETWORK_PORT.
    """

    from dragon.launcher.launchargs import valid_port_int

    parser = argparse.ArgumentParser(description='Starts the Dragon launch helper for generating network topology')

    parser.add_argument('--network-prefix', required=True, type=str, help='Network prefix to look for.')
    parser.add_argument('--port', required=True, type=valid_port_int, help='Port to connect to.')

    parser.set_defaults(
        port=DEFAULT_OVERLAY_NETWORK_PORT,
        network_prefix=DEFAULT_TRANSPORT_NETIF,
    )

    args = parser.parse_args()

    # Generate my node descriptor
    try:
        node_info = NodeDescriptor.get_local_node_network_conf(args.network_prefix, args.port)
    except RuntimeError:
        raise RuntimeError("Unable to acquire backend network configuration")

    # Dump to stdout
    print(json.dumps(node_info.get_sdict()))


def get_args(inputs=None):
    """Parse command line inputs for use in network config generator

    Returns:
        Namespace: namespace of expressing input arguments
    """
    from dragon.launcher.launchargs import NETWORK_HELP, valid_port_int

    parser = argparse.ArgumentParser(description='Runs Dragon internal tool for generating network topology')

    parser.add_argument('-p', '--port', type=valid_port_int, help="Infrastructure listening port (default: %(default)s)")
    parser.add_argument('--network-prefix', metavar='NETWORK_PREFIX', type=str, help=NETWORK_HELP)
    parser.add_argument('--wlm', '-w', metavar='WORKLOAD_MANAGER', type=WLM.from_str, choices=list(WLM), help=WLM_HELP)
    parser.add_argument('--log', '-l', help="Enable debug logging", action='store_true')
    parser.add_argument('--output-to-yaml', '-y', action='store_true', help='Output configuration to YAML file')
    parser.add_argument('--output-to-json', '-j', action='store_true', help='Output configuration to JSON file')
    parser.add_argument('--no-stdout', action='store_true', help='Do not print the configuration to stdout')
    parser.add_argument('--primary', type=str, help='Specify the hostname to be used for the primary compute node')

    host_group = parser.add_mutually_exclusive_group()
    host_group.add_argument('--hostlist', action=SplitArgsAtComma, metavar='HOSTLIST', type=str, help=HOSTLIST_HELP)
    host_group.add_argument('--hostfile', type=str, metavar='HOSTFILE', help=HOSTFILE_HELP)

    parser.set_defaults(
        port=DEFAULT_OVERLAY_NETWORK_PORT,
        network_prefix=DEFAULT_TRANSPORT_NETIF,
        log=False,
        output_to_yaml=False,
        output_to_json=False,
        no_stdout=False,
        primary=None,
        hostlist=None,
        hostfile=None
    )

    if inputs is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(inputs)

    if args.log:
        dlog.setup_logging(basename='wlm_config', level=logging.DEBUG)
        LOGGER = logging.getLogger('wlm_config')
        LOGGER.debug(f'Command line arguments: {args}')

    wlm = vars(args).get('wlm', None)
    if not wlm:
        try:
            for wlm, cls in NetworkConfig.wlm_cls_dict.items():
                if cls.check_for_wlm_support():
                    print(f'Detected a {str(wlm)} environment.')
                    break
        except Exception:
            raise RuntimeError('Could not detect a supported WLM environment.')

    # Handle the ssh case
    if wlm is WLM.SSH:
        args.hostlist = parse_hosts(args.hostlist, args.hostfile)
    else:
        if args.hostlist is not None or args.hostfile is not None:
            raise NotImplementedError('hostlist and hostfile arguments are only supported in the SSH case')

    return wlm, args


def main():
    '''
    .. code-block::
        :caption: **Dragon Network Config (dragon-network-config) tool's help and basic use**

        usage: dragon-network-config [-h] [-p PORT] [--network-prefix NETWORK_PREFIX] [--wlm WORKLOAD_MANAGER] [--log] [--output-to-yaml] [--output-to-json]
                                     [--no-stdout] [--primary PRIMARY] [--hostlist HOSTLIST | --hostfile HOSTFILE]

        Runs Dragon internal tool for generating network topology

        optional arguments:
          -h, --help            show this help message and exit
          -p PORT, --port PORT  Infrastructure listening port (default: 6565)
          --network-prefix NETWORK_PREFIX
                                NETWORK_PREFIX specifies the network prefix the dragon runtime will use to determine which IP addresses it should use to build
                                multinode connections from. By default the regular expression r'^(hsn|ipogif|ib)\d+$' is used -- the prefix for known HPE-Cray XC
                                and EX high speed networks. If uncertain which networks are available, the following will return them in pretty formatting: `dragon-
                                network-ifaddrs --ip --no-loopback --up --running | jq`. Prepending with `srun` may be necessary to get networks available on
                                backend compute nodes
          --wlm WORKLOAD_MANAGER, -w WORKLOAD_MANAGER
                                Specify what workload manager is used. Currently supported WLMs are: slurm, pbs+pals, ssh
          --log, -l             Enable debug logging
          --output-to-yaml, -y  Output configuration to YAML file
          --output-to-json, -j  Output configuration to JSON file
          --no-stdout           Do not print the configuration to stdout
          --primary PRIMARY     Specify the hostname to be used for the primary compute node
          --hostlist HOSTLIST   Specify backend hostnames as a comma-separated list, eg: `--hostlist host_1,host_2,host_3`. `--hostfile` or `--hostlist` is a
                                required argument for WLM SSH and is only used for SSH
          --hostfile HOSTFILE   Specify a list of hostnames to connect to via SSH launch. The file should be a newline character separated list of hostnames.
                                `--hostfile` or `--hostlist` is a required argument for WLM SSH and is only used for SSH

        # To create YAML and JSON files with a slurm WLM:
        $ dragon-network-config --wlm slurm --output-to-yaml --output-to-json
    '''

    try:
        wlm, args = get_args()

        net = NetworkConfig.from_wlm(workload_manager=wlm,
                                     port=args.port,
                                     network_prefix=args.network_prefix,
                                     primary_hostname=args.primary,
                                     hostlist=args.hostlist
                                     )
        if not args.no_stdout:
            config = net.get_network_config()
            local_sdict = {}
            for key, item in config.items():
                local_sdict[key] = item.get_sdict()
            print(json.dumps(local_sdict, indent=4))

        if args.output_to_yaml:
            net.output_to_file(basename=str(wlm), basedir=os.getcwd(), output_type=ConfigOutputType.YAML)

        if args.output_to_json:
            net.output_to_file(basename=str(wlm), basedir=os.getcwd(), output_type=ConfigOutputType.JSON)

    except RuntimeError:
        raise


if __name__ == '__main__':
    main()

src/dragon/launcher/network_config.py
DragonHPC/dragon
import asyncio
import traceback
import sys
import queue
import logging
import shutil
from functools import wraps, partial
from typing import Callable

from ..infrastructure.messages import AbnormalTerminationError
from ..infrastructure import messages as dmsg
from ..infrastructure import facts as dfacts
from ..infrastructure import process_desc as pdesc

from .wlm import wlm_launch_dict, WLM


# general amount of patience we have for an expected message
# in startup or teardown before we assume something has gone wrong
TIMEOUT_PATIENCE = 1  # seconds, 1 second.
QUEUE_PATIENCE = 0.1  # seconds of patience to wait for messages in MRNet
LA_TAG = 0
LOGBASE = 'launcher'


class SRQueue(queue.SimpleQueue):
    def send(self, msg):
        super().put(msg)

    def recv(self, timeout=None):
        try:
            return super().get(timeout=timeout)
        except queue.Empty:
            return None


def next_tag():
    global LA_TAG
    tmp = LA_TAG
    LA_TAG += 1
    return tmp


def detect_wlm():
    """Detect a supported WLM"""
    from .network_config import NetworkConfig

    wlm = None
    try:
        for wlm, cls in NetworkConfig.wlm_cls_dict.items():
            if cls.check_for_wlm_support():
                break

    except Exception:
        raise RuntimeError('Error searching for supported WLM')

    if wlm is None:
        raise RuntimeError('No supported WLM found')
    elif wlm is WLM.SSH:
        msg = \
"""
SSH was only supported launcher found. To use it, specify `--wlm ssh` as input to the dragon launcher.
It requires passwordless SSH to all backend compute nodes and a list of hosts. Please see documentation
and `dragon --help` for more information.
"""
        raise RuntimeError(msg)

    return wlm


def get_wlm_launch_args(args_map: dict = None,
                        nodes: tuple[int, list[str]] = None,
                        hostname: str = None,
                        wlm: WLM = None):
    """Get arguments for WLM to launch the backend"""

    try:
        if wlm is None:
            wlm = detect_wlm()
        wlm_args = wlm_launch_dict[wlm](args_map=args_map)

        if nodes is not None:
            wlm_args = wlm_args.format(nnodes=nodes[0], nodelist=','.join(nodes[1]))

        if hostname is not None:
            wlm_args = wlm_args.format(hostname=hostname)

    except Exception:
        raise RuntimeError("Unable to generate WLM backend launch args")

    return wlm_args.split()


def queue_monitor(func: Callable, *, log_test_queue=None):
    """Decorator to pull logging or abnormal messages out launch communications queues

    Take messages out of callback queue and puts into a separate logging queue to be dealt
    with later, thereby allowing us to immediately handle infrastructure

    Args:
        func (Callable): Callable that emits messages

        log_test_queue (queue, optional): Queue to drop log messages into. Used for internal
            unit testing. Defaults to None.

    Returns:
        function: The decorator wrapper function
    """

    if func is None:
        return partial(queue_monitor, log_test_queue=log_test_queue)

    @wraps(func)
    def wrapper(*args, **kwargs):
        while True:
            msg = func(*args, **kwargs)
            if isinstance(msg, dmsg.LoggingMsgList):
                if log_test_queue:
                    log_test_queue.put(msg)
                else:
                    for record in msg.records:
                        log = logging.getLogger(record.name)
                        log.log(record.level, record.msg, extra=record.get_logging_dict())
            elif isinstance(msg, dmsg.LoggingMsg):
                if log_test_queue:
                    log_test_queue.put(msg)
                else:
                    log = logging.getLogger(msg.name)
                    log.log(msg.level, msg.msg, extra=msg.get_logging_dict())
            elif isinstance(msg, dmsg.AbnormalTermination):
                raise AbnormalTerminationError("Abnormal exit detected")
            else:
                return msg
    return wrapper


@queue_monitor
def get_with_timeout(handle, timeout=TIMEOUT_PATIENCE):
    """Function for getting messages from a queue given a timeout"""

    # Handle queues we pass in either have a poll or
    # a timeout recv. Use the appropriate one
    if hasattr(handle, "poll"):
        if handle.poll(timeout=timeout):
            return dmsg.parse(handle.recv())
        else:
            raise TimeoutError
    else:
        msg = handle.recv(timeout=timeout)
        if isinstance(msg, tuple):
            return msg
        elif msg is not None:
            return dmsg.parse(msg)
        else:
            raise TimeoutError


@queue_monitor
def get_with_blocking(handle):
    """Function for getting messages from the queue while blocking"""
    try:
        msg = handle.recv()
        if isinstance(msg, tuple):
            return msg
        else:
            return dmsg.parse(msg)
    except Exception:
        raise


class LaOverlayNetFEQueue(queue.SimpleQueue):
    """Class for sending messages between launcher and OverlayNet threads"""

    def send(self, target, msg):
        super().put(target)
        super().put(msg)

    def recv(self, timeout=None):
        try:
            target = super().get(timeout=timeout)
            msg = dmsg.parse(super().get())
            return target, msg
        except queue.Empty:
            return None


class OverlayNetLaFEQueue(queue.SimpleQueue):
    """Class for sending messages between launcher and OverlayNet threads"""

    def send(self, msg):
        super().put(msg)

    def recv(self, timeout=None):
        try:
            msg = super().get(timeout=timeout)
            return msg
        except queue.Empty:
            return None


class AsyncQueueMonitor:
    def __init__(self, description, queue, dispatch, log=None, queue_wh=None):
        self.queue = queue

        # The write handle is needed to cleanly close the Queue Monitor when
        # it is waiting on a recv. We need to give it something to receive
        # because the run_in_executor is holding a lock that needs to be released.
        if queue_wh is None:
            self.queue_wh = self.queue
        else:
            self.queue_wh = queue_wh

        self.dispatch = dispatch
        self.description = description
        self.log = log
        self.running = True
        self.task = asyncio.create_task(self._receive_dispatch())

    def __str__(self):
        rv = f"""{self.description} AsyncQueueMonitor:
    task = {self.task}
    queue = {self.queue}"""
        return rv

    def get_task(self):
        # This creates a new task if one needs creating. Otherwise
        # it will continue to return the
        if self.task.done():
            self.task = asyncio.create_task(self._receive_dispatch())

        return self.task

    def close(self):
        self.running = False
        try:
            self.queue_wh.send("EXIT\n")
        except:
            pass

    async def _receive_dispatch(self):
        try:
            loop = asyncio.get_running_loop()
            while self.running:
                # Get the serialized message.
                json_string = await loop.run_in_executor(None, self.queue.recv)

                if self.running:
                    if json_string == '':
                        # Seems to happen at times during shutdown.
                        self.log.info("Got an empty string for a message in AsyncQueueMonitor for " + self.description)
                        raise EOFError("Got an empty string for a message in AsyncQueueMonitor for " + self.description)

                    # Parse it into a dictionary
                    msg = dmsg.parse(json_string)

                    # reset json_string in case a bad message comes in
                    # next time around. We don't want to get the previous
                    # json_string if an exception occurs.
                    del json_string

                    # Call dispatch on the server. The dispatch
                    # method will route this message to the
                    # appropriate message handler. By calling
                    # dispatch directly (and not through the object on which
                    # dispatch is defined) multiple dispatch functions
                    # are possible per server and potentially this queue
                    # could have its own custom dispatch function.
                    self.dispatch(msg)

            self.log.info('AsyncQueueMonitor exiting for queue ' + self.description)
            self.queue.close()

        except EOFError:
            # The underlying queue has closed.
            self.log.info('AsyncQueueMonitor exiting due to EOF of queue for ' + self.description)

        except Exception as ex:

            if not self.running:
                # If this is done, then the exception was likely a closed
                # queue. In that case, just return
                return

            if self.log is not None:
                exc_type, exc_value, exc_traceback = sys.exc_info()
                try:
                    try:
                        text = "Got an unexpected message in AsyncQueueMonitor _receive_dispatch for " + \
                               self.description + ". Message is '" + repr(
                            msg) + "'"
                    except Exception as ex:
                        text = "Got an invalid message in AsyncQueueMonitor _receive_dispatch for " + \
                               self.description + ". Message is '" + json_string + "'"
                except:
                    text = "Things are really bad in AysncQueueMonitor _receive_dispatch for " + self.description

                ex_message = text + '\n' + "".join(traceback.format_exception(exc_type, exc_value, exc_traceback))
                self.log.warning(ex_message)

            raise ex


def get_process_exe_args(args_map=None):
    """Parse sys.argv to get the user executable and arguments"""

    from .launchargs import head_proc_args as get_head_proc_args

    head_proc, head_proc_args = get_head_proc_args(args_map)

    if 1 == len(sys.argv):
        exe = sys.executable
        args = []
    else:
        target_fp = shutil.which(head_proc)
        if target_fp is not None:
            exe = target_fp
            args = head_proc_args[1:]  # Exclude the executable
        else:
            exe = sys.executable
            args = head_proc_args

    return exe, args


def mk_head_proc_start_msg(logbase='launcher',
                           make_inf_channels=True,
                           args_map=None):
    """Look at script arguments and return the head process launch message.

    If the first argument looks executable via 'which' then that will
    be what is launched as a managed process.  Otherwise sys.executable
    will be started with all the other arguments tacked on.

    :return: dmsg.GSProcessCreate message
    """
    log = logging.getLogger(logbase).getChild('mk_head_proc_start_msg')

    exe, args = get_process_exe_args(args_map=args_map)

    log.info(f'head process will be {exe} {args}')

    options = None
    if (make_inf_channels):
        options = pdesc.ProcessOptions(make_inf_channels=True)

    return dmsg.GSProcessCreate(tag=next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                                r_c_uid=dfacts.BASE_BE_CUID, exe=exe, args=args,
                                options=options)


def mk_shproc_echo_msg(logbase='launcher', stdin_str='', node_index=0):
    """Look at script arguments and return a SHProcessCreate message
    to be used when using the transport test environment mode.

    If the first argument looks executable via 'which' then that will
    be what is launched as a managed process.  Otherwise sys.executable
    will be started with all the other arguments tacked on.

    :return: dmsg.SHProcessCreate message
    """
    log = logging.getLogger(logbase).getChild('mk_shproc_echo_msg')

    exe, args = get_process_exe_args()
    args = ['-c', "import socket; print(f'hello from {socket.gethostname()}')"]

    log.info(f'The shprocesscreate message will include {exe} {args}')

    return dmsg.SHProcessCreate(tag=next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                                r_c_uid=dfacts.launcher_cuid_from_index(node_index),
                                t_p_uid=dfacts.FIRST_PUID,
                                exe=exe, args=args, initial_stdin=stdin_str)


def mk_shproc_start_msg(logbase='launcher', stdin_str=''):
    """Look at script arguments and return a SHProcessCreate message
    to be used when using the transport test environment mode.

    If the first argument looks executable via 'which' then that will
    be what is launched as a managed process.  Otherwise sys.executable
    will be started with all the other arguments tacked on.

    :return: dmsg.GSProcessCreate message
    """
    log = logging.getLogger(logbase).getChild('mk_head_shproc_create_msg')

    exe, args = get_process_exe_args()
    log.info(f'The shprocesscreate message will include {exe} {args}')

    return dmsg.SHProcessCreate(tag=next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                                r_c_uid=dfacts.BASE_BE_CUID, t_p_uid=dfacts.FIRST_PUID, exe=exe, args=args, initial_stdin=stdin_str)

src/dragon/launcher/util.py
DragonHPC/dragon
import enum

from .slurm import SlurmNetworkConfig, get_slurm_launch_be_args
from .pbs_pals import PBSPalsNetworkConfig, get_pbs_pals_launch_be_args
from .ssh import SSHNetworkConfig, get_ssh_launch_be_args


@enum.unique
class WLM(enum.Enum):
    """Enumerated list of supported workload manager"""
    SLURM = 'slurm'
    PBS_PALS = 'pbs+pals'
    SSH = 'ssh'

    def __str__(self):
        return self.value

    @staticmethod
    def from_str(s):
        """Obtain enum value from WLM string

        Args:
            s (str): name of WLM (eg: 'slurm', 'pbs_pals', etc.)

        Raises:
            ValueError: invalid string input

        Returns:
            WLM: enumerated value associated with string WLM
        """
        try:
            return WLM(s)
        except KeyError:
            raise ValueError()


wlm_launch_dict = {WLM.SLURM: get_slurm_launch_be_args,
                   WLM.PBS_PALS: get_pbs_pals_launch_be_args,
                   WLM.SSH: get_ssh_launch_be_args}

src/dragon/launcher/wlm/__init__.py
DragonHPC/dragon
from abc import ABC, abstractmethod
import json
import logging
import subprocess
import signal
from enum import Enum
from shlex import quote

from dragon.infrastructure.node_desc import NodeDescriptor
from dragon.infrastructure.util import NewlineStreamWrapper


class NetworkConfigState(Enum):
    """Enumerated states of Dragon FrontEnd
    """

    NONE = 0
    IN_PROGRESS = 1
    CONFIG_DONE = 2


class BaseNetworkConfig(ABC):
    def __init__(self, network_prefix, port, nnodes):
        self.NNODES = nnodes
        self.NETWORK_CFG_HELPER_LAUNCH_CMD = [
            "dragon-network-config-launch-helper",
            "--network-prefix", f"{network_prefix}",
            "--port", f"{port}",
        ]
        self.NETWORK_CFG_HELPER_LAUNCH_SHELL_CMD = [
            "dragon-network-config-launch-helper",
            "--network-prefix", f"{quote(network_prefix)}",
            "--port", f"{port}",


        ]
        self.LOGGER = logging.getLogger("NetworkConfig")

        self._sigint_trigger = None
        self._sigint_triggered = False
        self.config_helper = None
        self._hosts = None

        self._state = NetworkConfigState.NONE

    @classmethod
    def check_for_wlm_support(cls) -> bool:
        raise NotImplementedError

    @abstractmethod
    def _launch_network_config_helper(self) -> subprocess.Popen:
        raise NotImplementedError

    def _parse_network_configuration(self, stdout) -> None:
        self.node_descriptors = {}

        stdout_stream = NewlineStreamWrapper(stdout)
        while len(self.node_descriptors.keys()) != self.NNODES:
            if stdout_stream.poll():
                lines = []
                while stdout_stream.poll():
                    line = stdout_stream.recv()
                    # sattach returns an empty string if nothing to report. ignore
                    if line == "":
                        break
                    else:
                        lines.append(line)

                for line in lines:
                    self.LOGGER.debug(f'{line=}')
                    if line != "":
                        node_index, node_desc = line.split(": ", maxsplit=1)
                        if " " in node_index:
                            node_index = node_index.split(" ")[-1]
                        if str(node_index) not in self.node_descriptors.keys():
                            self.LOGGER.debug(json.loads(node_desc))
                            self.node_descriptors[
                                str(node_index)
                            ] = NodeDescriptor.from_sdict(json.loads(node_desc))

    def _sigint_teardown(self):
        """Safely teardown network config infrastructure"""

        try:
            if self.config_helper.is_alive():
                self.LOGGER.info('Transmitting SIGINT to config helper')
                self.config_helper.send_signal(signal.SIGINT)
                self.config_helper.wait()
        except AttributeError:
            pass

        raise KeyboardInterrupt

    def _sigint_handler(self, *args):
        """Handler for SIGINT signals for graceful teardown
        """
        self._sigint_triggered = True
        if self._state == NetworkConfigState.NONE or \
           self._state == NetworkConfigState.CONFIG_DONE:
            self._sigint_teardown()

    def get_network_config(self, sigint_trigger=None) -> map:

        try:
            self.orig_handler = signal.signal(signal.SIGINT, self._sigint_handler)

            self._state = NetworkConfigState.IN_PROGRESS

            self.LOGGER.debug("Launching config helper.")
            self.config_helper = self._launch_network_config_helper()

            if sigint_trigger == -2:
                signal.raise_signal(signal.SIGINT)

            self.LOGGER.debug("Waiting for config helper to exit.")
            self.config_helper.wait()
            if self.config_helper.returncode != 0:
                out, err = self.config_helper.communicate()
                raise RuntimeError(str(err))

            self.LOGGER.debug("Parsing configuration data.")
            self._parse_network_configuration(self.config_helper.stdout)

            self.LOGGER.debug("Closing config helper's stdout handle.")
            self.config_helper.stdout.close()
            self.config_helper.stderr.close()

            if sigint_trigger == -1:
                signal.raise_signal(signal.SIGINT)

            self._state = NetworkConfigState.CONFIG_DONE
            self.LOGGER.debug("All child procs exited. Returning output")

            if self._sigint_triggered:
                self._sigint_teardown()

            # Set the handler back before leaving
            signal.signal(signal.SIGINT, self.orig_handler)

            # return
            return self.node_descriptors

        except Exception as e:
            raise RuntimeError(e)

src/dragon/launcher/wlm/base.py
DragonHPC/dragon
import os
import re
import shutil
import subprocess

from .base import BaseNetworkConfig

PBS_PALS_LAUNCH_BE_ARGS = 'mpiexec --np {nnodes} --ppn 1 --cpu-bind none --hosts {nodelist} --line-buffer'


def get_pbs_pals_launch_be_args(args_map=None):
    return PBS_PALS_LAUNCH_BE_ARGS


def get_nodefile_node_count(filename) -> int:
    nnodes = 0
    with open(filename) as f:
        for nnodes, _ in enumerate(f, start=1):
            pass
    return nnodes


class PBSPalsNetworkConfig(BaseNetworkConfig):

    MPIEXEC_COMMAND_LINE = "mpiexec --np {nnodes} -ppn 1 -l"

    def __init__(self, network_prefix, port, hostlist):

        if not os.environ.get("PBS_NODEFILE"):
            msg = """Requesting a PBS network config outside of PBS job allocation.
Resubmit as part of a 'qsub' execution"""
            raise RuntimeError(msg)

        super().__init__(
            network_prefix,
            port,
            get_nodefile_node_count(os.environ.get("PBS_NODEFILE")),
        )

        self.MPIEXEC_ARGS = self.MPIEXEC_COMMAND_LINE.format(nnodes=self.NNODES).split()

    @classmethod
    def check_for_wlm_support(cls) -> bool:
        if (mpiexec := shutil.which("mpiexec")):
            return re.match('.*/pals/.*', mpiexec)
        return False

    def _launch_network_config_helper(self) -> subprocess.Popen:
        mpiexec_launch_args = self.MPIEXEC_ARGS[:]
        mpiexec_launch_args.append("--line-buffer")
        mpiexec_launch_args.extend(self.NETWORK_CFG_HELPER_LAUNCH_CMD)

        self.LOGGER.debug(f"Launching config with: {mpiexec_launch_args=}")

        return subprocess.Popen(
            args=mpiexec_launch_args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            bufsize=0,
            start_new_session=True
        )

src/dragon/launcher/wlm/pbs_pals.py
DragonHPC/dragon
import os
import shutil
import subprocess

from .base import BaseNetworkConfig

SLURM_LAUNCH_BE_ARGS = 'srun --nodes={nnodes} --ntasks={nnodes} --cpu_bind=none --nodelist={nodelist}'


def get_slurm_launch_be_args(args_map=None):
    return SLURM_LAUNCH_BE_ARGS


class SlurmNetworkConfig(BaseNetworkConfig):

    SRUN_COMMAND_LINE = "srun --nodes={nnodes} --ntasks={nnodes} --cpu_bind=none -u -l"

    def __init__(self, network_prefix, port, hostlist):

        if not os.environ.get("SLURM_JOB_ID"):
            msg = """Requesting a slurm network config outside of slurm job allocation.
Resubmit as part of a 'salloc' or 'sbatch' execution."""
            raise RuntimeError(msg)

        super().__init__(
            network_prefix, port, int(os.environ.get("SLURM_JOB_NUM_NODES"))
        )

        self.SRUN_ARGS = self.SRUN_COMMAND_LINE.format(nnodes=self.NNODES).split()

    @classmethod
    def check_for_wlm_support(cls) -> bool:
        return shutil.which("srun")

    def _launch_network_config_helper(self) -> subprocess.Popen:
        srun_launch_args = self.SRUN_ARGS[:]
        srun_launch_args.extend(self.NETWORK_CFG_HELPER_LAUNCH_CMD)

        self.LOGGER.debug(f"Launching config with: {srun_launch_args=}")

        return subprocess.Popen(
            args=srun_launch_args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            bufsize=0,
            start_new_session=True
        )

src/dragon/launcher/wlm/slurm.py
DragonHPC/dragon
import shutil
import subprocess
import selectors
import json
from itertools import groupby
from io import BufferedIOBase, TextIOBase
from os import environ, getcwd
from shlex import quote

from ...infrastructure.node_desc import NodeDescriptor
from ...infrastructure.parameters import this_process
from .base import BaseNetworkConfig

# TODO: we're propagating environment variables to make sure
#       the SSH envionment has the same envrionment as the
#       frontend. This works if there's a shared filesystem.
#       It's garbage otherwise. Longterm, we have to figure out
#       how to propagate Dragon out absent a shared filesystem.
ENV_VARS = None


def get_ssh_env_vars(args_map=None) -> list[str]:

    global ENV_VARS
    ENV_VARS = f"""DRAGON_BASE_DIR={environ.get('DRAGON_BASE_DIR','')} \
                   DRAGON_VERSION={environ.get('DRAGON_VERSION', '')} \
                   DRAGON_HSTA_DEBUG={environ.get('DRAGON_HSTA_DEBUG',"")} \
                   PATH={environ.get('PATH', '')} \
                   PYTHONPATH={environ.get('PYTHONPATH', '')} \
                   LD_LIBRARY_PATH={environ.get('LD_LIBRARY_PATH', '')} \
                   PYTHONSTARTUP={environ.get('PYTHONSTARTUP', '')} \
                   VIRTUAL_ENV={environ.get('VIRTUAL_ENV', '')} \
                """
    local_env = ENV_VARS
    # Make sure we propogate the device logging
    try:
        for log_device, log_level in args_map['log_device_level_map'].items():

            local_env = " ".join([local_env, f"{log_device}={log_level}"])
    except (KeyError, TypeError):
        pass

    # Construct a list of all the environment variables
    all_envs = [f'{var}={quote(val)}' for var, val in this_process.env().items()]
    all_envs.append(local_env)

    ENV_VARS = all_envs


def get_ssh_launch_be_args(hostname=None, args_map=None) -> str:

    global ENV_VARS
    if ENV_VARS is None or args_map is not None:
        get_ssh_env_vars(args_map=args_map)

    be_args = " ".join(["ssh -oBatchMode=yes {hostname}", f"cd {getcwd()} &&"] + ENV_VARS)
    return be_args


class SSHIOStream(BufferedIOBase):

    def __init__(self, streams: list[(int, TextIOBase)]):

        self._inputs = streams
        self.selectors = {}

    def _iterative_poll(self, idx, stream, timeout=0):
        '''Test if there's a read event in the input stream'''

        try:
            sel = self.selectors[idx]
        except (AttributeError, KeyError):
            sel = self.selectors[idx] = selectors.DefaultSelector()
            sel.register(stream, selectors.EVENT_READ)
        events = sel.select(timeout)
        for key, mask in events:
            if key.fileobj is not stream:
                continue
            if mask & selectors.EVENT_READ:
                return True
        return False

    def poll(self, timeout=0):
        '''Loop through our streams testing for a read event

        Return as soon as an event is found with True. False otherwise
        '''
        for idx, stream in self._inputs:
            if self._iterative_poll(idx, stream, timeout=timeout):
                return True
        return False

    def readline(self, **kwargs):
        '''Find a line to read and return it, prefixed with an index'''
        for idx, stream in self._inputs:
            if self._iterative_poll(idx, stream):
                msg = stream.readline()

                # See if this gets parsed into a Node Descriptor. If so,
                # remove it from the stream from our monitoring list:
                try:
                    NodeDescriptor.from_sdict(json.loads(msg))
                    self._inputs.remove((idx, stream))
                except (TypeError, json.decoder.JSONDecodeError):
                    pass

                if isinstance(msg, str):
                    msg = f"{idx}: {msg}"
                else:
                    msg = f"{idx}: {msg.decode()}"

                return msg

    def close(self):
        '''Close the streams'''
        for _, stream in self._inputs:
            try:
                stream.close()
            except Exception:
                pass

    def write(self, b):
        for _, stream in self._inputs:
            stream.write(b)

    def flush(self):
        for _, stream in self._inputs:
            stream.flush()


class SSHSubprocessPopen():
    """A convenience class providing Popen methods and attributes on a list of Popen objects"""

    def __init__(self, popen_dict: dict[str, subprocess.Popen]):

        self.procs = popen_dict
        self.nprocs = len(self.procs)
        self.args = [proc.args for proc in self.procs.values()]
        self.pid = [proc.pid for proc in self.procs.values()]

        self._returncode = None
        self._ret_codes = []

        # Set up newlinestreamwrappers for each of the subprocess stdout/stderr
        self._stdouts = [(i, proc.stdout) for i, proc in enumerate(self.procs.values())
                         if proc.stdout is not None]
        self._stderrs = [(i, proc.stderr) for i, proc in enumerate(self.procs.values())
                         if proc.stderr is not None]
        self._stdins = [(i, proc.stdin) for i, proc in enumerate(self.procs.values())
                        if proc.stdin is not None]

        self.stdout = None
        self.stderr = None
        self.stdin = None
        if len(self._stdouts) > 0:
            self.stdout = SSHIOStream(self._stdouts)

        if len(self._stderrs) > 0:
            self.stderr = SSHIOStream(self._stderrs)

        if len(self._stdins) > 0:
            self.stdin = SSHIOStream(self._stdins)

    def poll(self):
        '''Check exit status on all Popen objects

        If not all procs have exited, return None. Otherwise returncode
        '''
        for proc in self.procs.values():
            if proc.poll() is not None:
                self._ret_codes.append(proc.returncode)

        if len(self._ret_codes) == self.nprocs:
            return self.returncode
        else:
            return None

    def wait(self, timeout=None):
        '''Wait on exit for all Popen objects'''

        for proc in self.procs.values():
            # TODO: We don't want to necessarily wait the full timeout
            #       for each Popen object. Ideally, we'd aggregate this
            #       timeout over multiple calls to wait
            if proc.wait(timeout=timeout) is not None:
                self._ret_codes.append(proc.returncode)

        if len(self._ret_codes) == self.nprocs:
            return self.returncode
        else:
            return None

    def terminate(self):
        '''Send SIGTERM to all Popen objects'''
        for proc in self.procs.values():
            proc.terminate()

    def kill(self):
        '''Send SIGKILL to all Popen objects'''
        for proc in self.procs.values():
            proc.kill()

    def send_signal(self, signal):
        '''Send signal to all Popen objects'''
        for proc in self.procs.values():
            proc.send_signal(signal)

    def communicate(self, input=None, timeout=None):
        '''Send data to stdin or recv data from stdout,stderr of all Popen objects'''
        if input is not None:
            self.stdin.write(input)

        out_msg = None
        err_msg = None
        if self.stderr is not None:
            err_msg = self.stderr.readline()

        if self.stdout is not None:
            out_msg = self.stdout.readline()

        return out_msg, err_msg

    @property
    def returncode(self):
        '''Return return codes for Popen objects'''

        if isinstance(self._ret_codes, list):
            code_group = groupby(self._ret_codes)
            # If they're all the same, return any one of them
            if next(code_group, True) and not next(code_group, False):
                self._returncode = self._ret_codes[0]
            # if not, return the first non-zero:
            else:
                self._returncode = next((code for code in self._ret_codes if code != 0), 0)

        return self._returncode


class SSHNetworkConfig(BaseNetworkConfig):

    def __init__(self, network_prefix, port, hostlist):

        super().__init__(
            network_prefix, port, len(hostlist)
        )
        self.hostlist = hostlist

    @classmethod
    def check_for_wlm_support(cls) -> bool:
        return shutil.which("ssh")

    def _launch_network_config_helper(self) -> subprocess.Popen:
        popen_dict = {}

        for host in self.hostlist:

            ssh_args = get_ssh_launch_be_args(hostname=host)
            ssh_args = ssh_args.format(hostname=host).split()
            ssh_args.extend(self.NETWORK_CFG_HELPER_LAUNCH_SHELL_CMD)
            self.LOGGER.debug(f"Launching config with: {ssh_args}")

            popen_dict[host] = subprocess.Popen(
                                   args=ssh_args,
                                   stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE,
                                   bufsize=0,
                                   start_new_session=True
                                   )

        return SSHSubprocessPopen(popen_dict)

src/dragon/launcher/wlm/ssh.py
DragonHPC/dragon
"""Dragon local services"""
import json
import logging
import os
import subprocess
import sys

from typing import Optional

import dragon.channels as dch
import dragon.managed_memory as dmm
import dragon.utils as dutils

from ..infrastructure.node_desc import NodeDescriptor
from ..transport import start_transport_agent
from ..launcher.launchargs import TransportAgentOptions

from ..infrastructure import connection as dconn
from ..infrastructure import facts as dfacts
from ..infrastructure import messages as dmsg
from ..infrastructure import parameters as dparms
from ..infrastructure import util as dutil
from ..dlogging.util import setup_BE_logging, DragonLoggingServices as dls
from ..utils import B64, set_procname

from .server import LocalServer, PopenProps, get_new_tag, ProcessProps, OutputConnector


def maybe_start_gs(gs_args: Optional[list], gs_env: Optional[dict], hostname: str, be_in: object, logger_sdesc: B64 = None):
    """Starts gs if there are any args to start it with.

    Args:
        gs_args (list): args to start global services

        gs_env (dict): environment variable to set prior to gs launcher

        logger_sdesc (B64, optional): logging channel descriptor

    Raises:
        RuntimeError: Launch of global services was not successful.

    Returns:
        gs_proc (PopenProps) of gs launch, or None if gs_args is None
    """
    log = logging.getLogger(dls.LS).getChild('maybe_start_gs')
    if gs_args is not None:
        if gs_env is None:
            the_env = os.environ
        else:
            the_env = dict(os.environ)
            the_env.update(gs_env)

        if logger_sdesc is not None:
            the_env[dfacts.DRAGON_LOGGER_SDESC] = str(logger_sdesc)

        try:
            log.debug(f'launcher gs: {gs_args}')
            stdout_connector = OutputConnector(be_in=be_in, puid=dfacts.GS_PUID,
                    hostname=hostname, out_err=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                    conn=None, root_proc=True, critical_proc=True)
            stderr_connector = OutputConnector(be_in=be_in, puid=dfacts.GS_PUID,
                    hostname=hostname, out_err=dmsg.SHFwdOutput.FDNum.STDERR.value,
                    conn=None, root_proc=True, critical_proc=True)

            gs_proc = PopenProps(
                ProcessProps(p_uid=dfacts.GS_PUID, critical=True, r_c_uid=None,
                             stdin_req=None, stdout_req=None, stderr_req=None,
                             stdin_connector=None, stdout_connector=stdout_connector,
                             stderr_connector=stderr_connector),
                gs_args,
                bufsize=0,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                env=the_env
            )

            gs_proc.props.stdout_connector.add_proc_info(gs_proc)
            gs_proc.props.stderr_connector.add_proc_info(gs_proc)

            log.info('started gs')
        except (OSError, ValueError) as gs_launch_err:
            log.fatal('gs launch failed')
            raise RuntimeError('gs launch failed') from gs_launch_err
    else:
        gs_proc = None

    return gs_proc


def mk_inf_resources(node_index):
    log = logging.getLogger(dls.LS).getChild('inf resource maker')
    log.info('creating  pools and channels')

    inf_pool = None
    def_pool = None

    _user = os.environ.get('USER', str(os.getuid()))

    try:
        inf_muid = dfacts.infrastructure_pool_muid_from_index(node_index)
        ips = dparms.this_process.inf_seg_sz
        ipn = f'{_user}_{os.getpid()}_' + dfacts.INFRASTRUCTURE_POOL_SUFFIX
        log.info(f'inf pool: {ipn} size {ips}')
        inf_pool = dmm.MemoryPool(ips, ipn, inf_muid)

        def_muid = dfacts.default_pool_muid_from_index(node_index)
        dps = dparms.this_process.default_seg_sz
        dpn = f'{_user}_{os.getpid()}_' + dfacts.DEFAULT_POOL_SUFFIX
        log.info(f'def pool: {dpn} size {dps}')

        # PJM TODO: how pre-allocated blocks are done here is temporary.  We also need a better way
        # to manage Pool attributes as there are many more we will expose.
        def_pool = dmm.MemoryPool(dps, dpn, def_muid, pre_alloc_blocks=[8,0,0,0,0,8])
        start_pools = {inf_muid: inf_pool, def_muid: def_pool}
        dparms.this_process.inf_pd = dutils.B64.bytes_to_str(inf_pool.serialize())
        dparms.this_process.default_pd = dutils.B64.bytes_to_str(def_pool.serialize())

        shep_input_cuid = dfacts.shepherd_cuid_from_index(node_index)
        shep_ch = dch.Channel(inf_pool, shep_input_cuid, None)
        dparms.this_process.local_shep_cd = dutils.B64.bytes_to_str(shep_ch.serialize())
        shep_input = dconn.Connection(inbound_initializer=shep_ch, policy=dparms.POLICY_INFRASTRUCTURE)

        la_input_cuid = dfacts.launcher_cuid_from_index(node_index)
        la_input_ch = dch.Channel(inf_pool, la_input_cuid, None)
        dparms.this_process.local_be_cd = dutils.B64.bytes_to_str(la_input_ch.serialize())
        dparms.this_process.be_cuid = dfacts.launcher_cuid_from_index(node_index)
        la_input = dconn.Connection(outbound_initializer=la_input_ch, policy=dparms.POLICY_INFRASTRUCTURE)

        ta_input_cuid = dfacts.transport_cuid_from_index(node_index)
        ta_input_ch = dch.Channel(inf_pool, ta_input_cuid, None)
        ta_input_descr = ta_input_ch.serialize()
        dparms.this_process.local_ta_cd = dutils.B64.bytes_to_str(ta_input_descr)
        ta_input = dconn.Connection(outbound_initializer=ta_input_ch, policy=dparms.POLICY_INFRASTRUCTURE)

        start_channels = {shep_input_cuid: shep_ch, la_input_cuid: la_input_ch, ta_input_cuid: ta_input_ch}

        if 0 == node_index:
            gs_chan = dch.Channel(inf_pool, dfacts.GS_INPUT_CUID, None)
            start_channels[dfacts.GS_INPUT_CUID] = gs_chan
            dparms.this_process.gs_cd = dutils.B64.bytes_to_str(gs_chan.serialize())
            gs_input = dconn.Connection(outbound_initializer=gs_chan, policy=dparms.POLICY_INFRASTRUCTURE)
        else:
            gs_input = None

    except (dch.ChannelError, dmm.DragonPoolError,
            dmm.DragonMemoryError) as init_err:
        log.fatal('could not make infrastructure resources')
        # last ditch attempt to clean up pools.
        if inf_pool is not None:
            inf_pool.destroy()
        if def_pool is not None:
            def_pool.destroy()
        raise RuntimeError('infrastructure resource creation failed') from init_err

    log.info('infrastructure resources constructed')
    return start_pools, start_channels, shep_input, la_input, ta_input_descr, ta_input, gs_input


def get_shepherd_msg_queue(stdin=None, stdout=None):

    if stdin is None:
        bin_stdin = os.fdopen(sys.stdin.fileno(), 'rb', buffering=0)
        shep_stdin_msg = dutil.NewlineStreamWrapper(bin_stdin, write_intent=False)
    else:
        shep_stdin_msg = stdin

    if stdout is None:
        bin_stdout = os.fdopen(sys.stdout.fileno(), 'wb', buffering=0)
        shep_stdout_msg = dutil.NewlineStreamWrapper(bin_stdout, write_intent=True)
    else:
        shep_stdout_msg = stdout

    return shep_stdin_msg, shep_stdout_msg


def single(make_infrastructure_resources: bool = True,
           gs_input=None, la_input=None, shep_input=None,
           gs_args=None, gs_env=None, ls_stdin=None,
           ls_stdout=None, ta_input=None):

    log = logging.getLogger('local_svc single node main')
    log.info(f'starting on pid={os.getpid()}')

    start_pools = {}
    start_channels = {}

    try:  # attempt to construct infrastructure objects and establish communications
        os.setpgid(0, 0)

        shep_stdin_msg, shep_stdout_msg = get_shepherd_msg_queue(ls_stdin,
                                                                 ls_stdout)

        msg = dmsg.parse(shep_stdin_msg.recv())
        assert isinstance(msg, dmsg.BENodeIdxSH), 'startup msg expected on stdin'
        log.info('got BENodeIdxSH')
        assert msg.node_idx == 0, 'single node'
        node_index = msg.node_idx

        node_sdesc = NodeDescriptor.make_for_current_node(is_primary=True).sdesc
        log.info(f"Node: {str(node_sdesc)}")

        if make_infrastructure_resources:
            assert not any((gs_input, la_input, shep_input, ta_input))
            start_pools, start_channels, shep_input, la_input, ta_input_descr, ta_input, gs_input = mk_inf_resources(node_index)

        dparms.this_process.index = node_index
        os.environ.update(dparms.this_process.env())
        be_ping = dmsg.SHPingBE(tag=get_new_tag(), shep_cd=dparms.this_process.local_shep_cd,
                                be_cd=dparms.this_process.local_be_cd,
                                gs_cd=dparms.this_process.gs_cd,
                                default_pd=dparms.this_process.default_pd,
                                inf_pd=dparms.this_process.inf_pd)
        shep_stdout_msg.send(be_ping.serialize())
        log.info('wrote SHPingBE')

        msg = dmsg.parse(shep_input.recv())
        assert isinstance(msg, dmsg.BEPingSH), 'startup expectation on shep input'
        log.info('got BEPingSH')

        ch_up_msg = dmsg.SHChannelsUp(tag=get_new_tag(), host_name='localhost',
                                      host_id=dutils.host_id(), ip_addrs=['127.0.0.1'],
                                      shep_cd=dparms.this_process.local_shep_cd,
                                      gs_cd=dparms.this_process.gs_cd)
        la_input.send(ch_up_msg.serialize())
        log.info('sent SHChannelsUp')
        gs_proc = maybe_start_gs(gs_args, gs_env, hostname='localhost', be_in=la_input)
        msg = dmsg.parse(shep_input.recv())
        assert isinstance(msg, dmsg.GSPingSH), 'startup expectation shep input'
        log.info('got GSPingSH')
    except (OSError, EOFError, json.JSONDecodeError, AssertionError, RuntimeError) as rte:
        log.fatal('startup failed')
        LocalServer.clean_pools(start_pools, log)
        raise RuntimeError('startup fatal error') from rte

    gs_input.send(dmsg.SHPingGS(tag=get_new_tag(), node_sdesc=node_sdesc).serialize())

    server = LocalServer(channels=start_channels, pools=start_pools, hostname='localhost')

    if gs_proc is not None:
        server.add_proc(gs_proc)
        server.gs_proc = gs_proc

    try:
        server.run(shep_in=shep_input, gs_in=gs_input, be_in=la_input, is_primary=True)
    except Exception:
        log.fatal('There was an exception and LS is going to clean up.')
        server.cleanup()
        raise

    try:
        msg = dmsg.parse(shep_stdin_msg.recv())
        assert isinstance(msg, dmsg.BEHalted)
        shep_stdout_msg.send(dmsg.SHHalted(tag=get_new_tag()).serialize())
    except (OSError, AssertionError, EOFError, json.JSONDecodeError) as tde:
        log.fatal('teardown sequence error')
        raise RuntimeError('teardown') from tde
    finally:
        server.cleanup()

    log.info('shutdown complete')


def multinode(make_infrastructure_resources: bool = True,
              gs_input=None, la_input=None, ls_input=None,
              gs_args=None, gs_env=None,
              ta_args=None, ta_env=None,
              ls_stdin=None, ls_stdout=None,
              fname=None, ta_input=None,
              transport_test_env: bool = False):
    '''Local services for multi-node dragon'''

    setup_BE_logging(service=dls.LS, fname=fname)
    log = logging.getLogger(dls.LS).getChild('multinode')
    log.debug(f'starting debug modes on pid={os.getpid()}')

    set_procname(dfacts.PROCNAME_LS)

    start_pools = {}
    start_channels = {}

    try:  # attempt to construct infrastructure objects and establish communications
        os.setpgid(0, 0)

        ls_stdin_queue, ls_stdout_queue = get_shepherd_msg_queue(ls_stdin,
                                                                 ls_stdout)

        msg = dmsg.parse(ls_stdin_queue.recv())
        assert isinstance(msg, dmsg.BENodeIdxSH), 'startup msg expected on stdin'
        node_index = msg.node_idx
        hostname = msg.host_name
        ip_addrs = msg.ip_addrs
        is_primary = msg.primary
        logger_sdesc = msg.logger_sdesc
        log.info(f'got BENodeIdxSH (id={node_index}, ips={ip_addrs}, host={hostname}, primary={is_primary}) - m2.1')

        # Add the dragon logging handler to our already existing log
        setup_BE_logging(service=dls.LS, logger_sdesc=logger_sdesc, fname=fname)
        log.debug(f'dragon logging initiated on pid={os.getpid()}')

        node_sdesc = NodeDescriptor.make_for_current_node(ip_addrs=ip_addrs, name=hostname, is_primary=is_primary).sdesc

        if make_infrastructure_resources:
            assert not any((gs_input, la_input, ls_input, ta_input))
            start_pools, start_channels, ls_input, la_input, ta_input_descr, ta_input, gs_input = mk_inf_resources(node_index)

        dparms.this_process.index = node_index
        os.environ.update(dparms.this_process.env())

        if is_primary:
            gs_cd = dparms.this_process.gs_cd
        else:
            gs_cd = None

        be_ping = dmsg.SHPingBE(tag=get_new_tag(), shep_cd=dparms.this_process.local_shep_cd,
                                be_cd=dparms.this_process.local_be_cd,
                                gs_cd=gs_cd,
                                default_pd=dparms.this_process.default_pd,
                                inf_pd=dparms.this_process.inf_pd)
        ls_stdout_queue.send(be_ping.serialize())
        log.info('wrote SHPingBE')

        msg = dmsg.parse(ls_input.recv())
        assert isinstance(msg, dmsg.BEPingSH), 'startup expectation on shep input'
        log.info('got BEPingSH')

        ch_up_msg = dmsg.SHChannelsUp(tag=get_new_tag(), host_name=hostname, host_id=dutils.host_id(),
                                      ip_addrs=ip_addrs, shep_cd=dparms.this_process.local_shep_cd,
                                      gs_cd=gs_cd, idx=node_index)

        la_input.send(ch_up_msg.serialize())
        log.info('sent SHChannelsUp')

        # Recv LAChannelsInfo Broadcast
        la_channels_info = dmsg.parse(ls_input.recv())
        assert isinstance(la_channels_info, dmsg.LAChannelsInfo), 'expected LAChannelsInfo'
        log.info(f'node index {node_index} received all channels info')
        log.debug(f'la_channels.nodes_desc: {la_channels_info.nodes_desc}')
        log.debug(f'la_channels.gs_cd: {la_channels_info.gs_cd}')

        # Work out transport agent launch arguments:
        ta_args = [f'{la_channels_info.transport}']


        # Create gateway channels
        gw_channels = []
        num_ls_gw_channels = la_channels_info.num_gw_channels

        # We set the number of gateways in the environment for other processes
        # that will be created by Local Services and that need to call
        # register_gateways_from_env. C level programs that want to use multi-node
        # channels need to call this function.
        dparms.this_process.set_num_gateways_per_node(num_ls_gw_channels)
        node_count = len(la_channels_info.nodes_desc)
        # get the cuid of the first gateway channel on node index
        gw_cuid = dfacts.gw_cuid_from_index(node_index, num_ls_gw_channels)
        def_muid = dfacts.default_pool_muid_from_index(node_index)
        for id in range(num_ls_gw_channels):
            gw_ch = dch.Channel(start_pools[def_muid], gw_cuid + id, capacity=dparms.this_process.gw_capacity)
            encoded_ser_gw = B64(gw_ch.serialize())
            # add the serialized descriptor of the gw channel to the environment
            # for the transport agent to get it
            encoded_ser_gw_str = str(encoded_ser_gw)
            os.environ[dfacts.GW_ENV_PREFIX + str(id + 1)] = encoded_ser_gw_str
            gs_env[dfacts.GW_ENV_PREFIX + str(id + 1)] = encoded_ser_gw_str
            gw_channels.append(gw_ch)
        log.info(f'ls created {num_ls_gw_channels} gateway channels')

        # Here we register the gateway channels
        # created by this process as environment
        # variables. This is to enable the shepherd
        # (local services) once the transport service
        # is started, which happens next. Order is
        # not important here. The gateways can be
        # registered as soon as they are created and part
        # of the environment.
        dch.register_gateways_from_env()
        log.info('ls has registered the gateways in its environment.')

        # Start TA (telling it its node ID by appending to args) and send LAChannelsInfo
        log.info('standing up ta')
        try:
            ta = start_transport_agent(
                node_index,
                B64(ta_input_descr),
                logger_sdesc,
                args=ta_args,
                env=ta_env,
            )
            # Cast the Popen instance returned by start_transport_agent() to
            # PopenProps mainly for consistency, though it doesn't appear to
            # matter since the TA process isn't used elsewhere.
            ta.__class__ = PopenProps
            ta.props = ProcessProps(
                p_uid=dfacts.transport_puid_from_index(node_index),
                critical=True,
                r_c_uid=None,
                stdin_connector=None,
                stdout_connector=None,
                stderr_connector=None,
                stdin_req=None,
                stdout_req=None,
                stderr_req=None
            )
        except Exception as e:
            logging.getLogger(dls.LS).getChild('start_ta').fatal(f'transport agent launch failed on {node_index}')
            raise RuntimeError(f'transport agent launch failed on node {node_index}') from e

        # Send LAChannelsInfo to TA
        ta_input.send(la_channels_info.serialize())

        # Confirmation TA is up
        ta_ping = dmsg.parse(ls_input.recv())
        assert isinstance(ta_ping, dmsg.TAPingSH), 'ls did not receive ping from TA)'
        log.info('ls received TAPingSH - m7')
        ch_list = []
        if transport_test_env:
            first_cuid = dfacts.FIRST_CUID + node_index * 2
            second_cuid = dfacts.FIRST_CUID + node_index * 2 + 1
            barrier_cuid = dfacts.FIRST_CUID + node_count * 2

            ch1 = dch.Channel(start_pools[def_muid], first_cuid)
            ch2 = dch.Channel(start_pools[def_muid], second_cuid)
            ch1_ser = B64(ch1.serialize())
            ch2_ser = B64(ch2.serialize())
            ch_list.append(str(ch1_ser))
            ch_list.append(str(ch2_ser))
            if node_index == 0:
                # Create one Barrier channel for use in testing.
                ch3 = dch.Channel(start_pools[def_muid], barrier_cuid, capacity=node_count)
                ch3_ser = B64(ch3.serialize())
                ch_list.append(str(ch3_ser))

        la_input.send(dmsg.TAUp(tag=get_new_tag(), idx=node_index, test_channels=ch_list).serialize())
        log.info("ls send TAUp to la_be - m8.1")

        # Init these here so if in transport test mode, they have a value even though GS will not
        # be started or used in transport test mode.
        gs = None
        gs_in_wh = None

        if not transport_test_env:
            # Start global services on primary node.
            if is_primary:
                log.info('Starting global services on primary')
                gs = maybe_start_gs(gs_args, gs_env, hostname=hostname, be_in=la_input, logger_sdesc=logger_sdesc)
                if gs is None:
                    log.info('did not start gs')
                else:
                    log.info('gs up')

                gs_stdin = os.fdopen(gs.stdin.fileno(), 'wb')
                gs_stdin_send = dutil.NewlineStreamWrapper(gs_stdin, read_intent=False)
                gs_stdin_send.send(la_channels_info.serialize())
                log.info('transmitted la_channels_info to gs')

            gs_ch = la_channels_info.gs_cd
            gs_in_ch = dch.Channel.attach(dutils.B64.str_to_bytes(gs_ch))
            log.info('ls attached to gs channel')
            gs_in_wh = dconn.Connection(outbound_initializer=gs_in_ch, policy=dparms.POLICY_INFRASTRUCTURE)

            gs_ping_ls = dmsg.parse(ls_input.recv())
            assert isinstance(gs_ping_ls, dmsg.GSPingSH), 'ls expected GSPingSH'
            log.info('ls received GSPingSH from gs - m10')

            # Send response to GS
            gs_in_wh.send(dmsg.SHPingGS(tag=get_new_tag(), idx=node_index, node_sdesc=node_sdesc).serialize())
            log.info('ls sent SHPingGS - m11')
    except (OSError, EOFError, json.JSONDecodeError, AssertionError, RuntimeError) as rte:
        log.fatal('startup failed')
        LocalServer.clean_pools(start_pools, log)
        raise RuntimeError('startup fatal error') from rte
    server = LocalServer(channels=start_channels, pools=start_pools, transport_test_mode=transport_test_env,
                         hostname=hostname)

    if gs is not None:
        server.add_proc(gs)
        server.gs_proc = gs

    if ta is not None:
        server.add_proc(ta)
        server.ta_proc = ta

    try:
        server.run(shep_in=ls_input, gs_in=gs_in_wh, be_in=la_input,
                   is_primary=is_primary, ta_in=ta_input, gw_channels=gw_channels)
    except Exception:
        log.fatal('There was an exception in LS and clean up is called.')
        server.cleanup()
        raise

    try:
        # m14 Recv BEHalted from BE
        # Wait till I'm told the backend is detached from me -- listening on the original pmsgqueue
        be_halted = dmsg.parse(ls_stdin_queue.recv())
        assert isinstance(be_halted, dmsg.BEHalted), f'm14 BEHalted msg expected. Received {type(be_halted)}'
        log.debug('m14 Received final BEHalted. Cleaning up and exiting.')
        ls_stdin_queue.close()
        ls_stdout_queue.close()

    except (OSError, AssertionError, EOFError, json.JSONDecodeError) as tde:
        log.fatal('teardown sequence error')
        raise RuntimeError('teardown') from tde
    finally:
        server.cleanup()

    log.info('shutdown complete')

src/dragon/localservices/local_svc.py
DragonHPC/dragon
import asyncio
import enum
import io
import sys
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.facts as dfacts

INVALID_COMMAND = 42


class Process:
    """
    This class provides facilities for managing information
    about a managed process.
    """

    class State(enum.Enum):
        INIT = enum.auto()
        RUN = enum.auto()
        COMPLETE = enum.auto()
        FAIL = enum.auto()

    class StdInState(enum.Enum):
        # is_open below means it is open to writing, but may not yet
        # be connected to its pipe. When ProcessState is run and
        # there is data to write, then the pipe is connected.
        IS_OPEN = enum.auto()
        CLOSE_CALLED = enum.auto()
        CLOSING = enum.auto()
        CLOSED = enum.auto()

    def __init__(self, p_uid: int, tag_ref: int, cmd: str, args: tuple, env: dict = None,
                 cwd: str = None, resp_queue=None, stdout_queue=None,
                 stderr_queue=None, term_queue=None):

        # helps for debugging to have id in process. This is user-assigned id.
        self.p_uid = p_uid
        # the tag/ref is used by GS to correlate between create and create response
        # messages. It comes in on the create message as a tag field. It is sent
        # back as the ref field in the create response.
        self.tag_ref = tag_ref
        self.cmd = cmd
        self.args = args
        self.env = env
        self.cwd = cwd
        self.output_queue = {ProcessManager.TaskType.STDOUT: stdout_queue,
                             ProcessManager.TaskType.STDERR: stderr_queue}
        self.term_queue = term_queue
        self._kill_timer = None
        self._signal_tag_ref = 0

        # If this is not None then it is the response queue to send a response message
        # to once creation is complete.
        self.resp_queue = resp_queue

        # The keys stdout and stderr are used to select the correct stream.
        self.stream = {}

        self.stdin_buf = bytearray()
        self.state = Process.State.INIT
        self.pid = None
        self._stat = None

        self.stdin_state = self.StdInState.IS_OPEN
        self.close_pending = False

        # keep track of the stream tasks responsible for reading
        # output so we can clean up at the end.
        self.stream_task = {}

    def __str__(self):
        rv = f"""
   Process p_uid = {self.p_uid}
     state = {self.state}
     pid = {self.pid}
     tag_ref = {self.tag_ref}
     cmd = {self.cmd}
     args = {self.args}
     env = {self.env}
     cwd = {self.cwd}
     output_queues = {self.output_queue}
     term_queue = {self.term_queue}
     kill_timer = {self._kill_timer}
     signal_tag_ref = {self._signal_tag_ref}
     resp_queue = {self.resp_queue}
     stdin_buffer = {self.stdin_buf}
     stdin_state = {self.stdin_state}
     Posix stat = {self._stat}
     close_pending = {self.close_pending}
"""
        return rv

    @property
    def signal_tag_ref(self):
        return self._signal_tag_ref

    def write_stdin(self, data: str) -> None:
        """
        Given a bytearray, write it to standard input. Returns None.
        """
        if not isinstance(data, str):
            raise ValueError(f'Only string objects can be written to stdin: {data}')

        if self.stdin_state != self.StdInState.IS_OPEN:
            raise io.UnsupportedOperation(f'stdin of subprocess is not open. State: {self.stdin_state!s}')

        byte_data = data.encode()

        if self.stdin_buf is not None:
            self.stdin_buf.extend(byte_data)
        else:
            self.stdin_buf = byte_data

    def close_stdin(self) -> None:
        """
        Close the stdin stream.
        """
        self.stdin_state = self.StdInState.CLOSE_CALLED

    def stop(self):
        """
        Send the process SIGTERM.
        """
        return self._stat.terminate()

    def signal(self, sig, tag_ref):
        """
        Send the process a signal.
        """
        self._signal_tag_ref = tag_ref
        return self._stat.send_signal(sig)

    def set_kill_timer(self, timer_task):
        self._kill_timer = timer_task

    def clear_kill_timer(self):
        self._kill_timer = None

    def cancel_kill_timer(self, logger):
        if self._kill_timer is not None:
            logger.info(f'Attempting to cancel kill timer for process {self.p_uid}')
            self._kill_timer.cancel()
            logger.info(f'Cancel Success for process {self.p_uid}')
            self._kill_timer = None


class ProcessManager:
    """
    ProcessManager

    This class provides asynchronous methods for starting up
    processes on the local node.  It forwards stdout/stderr
    and redirects stdin to processes. It is also the driver
    of asynchronous execution of system tasks provided by the
    shepherd. The ProcessManager and the Shepherd work together
    to provide the services of the Shepherd component in Dragon.
    """

    class TaskType(enum.Enum):
        """
        The TaskType is a identifier of the type of task being executed asynchronously
        in the ProcessManager and also identifies the type of stream.
        """
        PROCESS = -2  # This indicates a user process task.
        SYS = -1  # This indicates a system task.
        # The stdin, stdout, and stderr values should remain
        # 0,1, and 2 respectively. They coincide with Posix
        # file descriptors but also with the FileDescriptor
        # enum in messages.
        STDIN = 0
        STDOUT = 1
        STDERR = 2

    def __init__(self, shepherd):

        self.processes = {}  # This is a dictionary of running/active processes.
        self.tasks = []
        self.new_processes = []  # This is for new processes that are to be created.
        self.task_tbl = {}
        self.shepherd = shepherd

        self.ect = asyncio.get_event_loop().create_task

    def __str__(self):

        tasks = '\n'.join([str(t) for t in self.tasks])

        procs = '\n'.join([str(p) for p in self.processes])
        new_procs = '\n'.join([str(p) for p in self.new_processes])

        rv = f"""
    Process Manager
Tasks:
{tasks}
       
Task Names Reverse Map:
{self.task_tbl.values()!s}

Processes:
{procs}

New Processes:
{new_procs}"""
        return rv

    def create_task(self, fun, task_type, id_num):
        task_name = f'{task_type!s}:{id_num=!s}'
        self.task_tbl[(task := self.ect(fun, name=task_name))] = (task_type, id_num)
        return task

    def _handle_started_procs(self, finished):

        local_new_tasks = []

        for task in set(finished):
            task_type, p_uid = self.task_tbl[task]
            if task_type == self.TaskType.PROCESS:
                # This check is necessary because it is possible that a task that finishes might
                # refer to a process that has already finished and if so, we don't want to
                # get an exception trying to find it in the process table.
                if p_uid in self.processes:
                    process = self.processes[p_uid]

                    if process.state == Process.State.INIT:
                        process.state = Process.State.RUN

                        process.stream[self.TaskType.STDOUT] = process._stat.stdout
                        process.stream[self.TaskType.STDERR] = process._stat.stderr
                        process.stream[self.TaskType.STDIN] = process._stat.stdin

                        proc_task = self.create_task(self.async_app_done(p_uid),
                                                     self.TaskType.PROCESS, p_uid)

                        stdout_task = self.create_task(self.async_read_output(p_uid, self.TaskType.STDOUT),
                                                       self.TaskType.STDOUT, p_uid)

                        stderr_task = self.create_task(self.async_read_output(p_uid, self.TaskType.STDERR),
                                                       self.TaskType.STDERR, p_uid)

                        process.stream_task[self.TaskType.STDOUT] = stdout_task
                        process.stream_task[self.TaskType.STDERR] = stderr_task
                        local_new_tasks.extend([proc_task, stdout_task, stderr_task])
                    else:
                        if process.state == Process.State.FAIL:
                            # remove it from the process table.
                            self.app_cleanup(process.p_uid)

                # remove it from the set of finished tasks
                del self.task_tbl[task]  # ??
                finished.remove(task)

        return local_new_tasks

    def _handle_done_procs(self, finished):
        local_new_tasks = []
        theTasks = set(finished)

        for task in theTasks:
            task_type, p_uid = self.task_tbl[task]
            if task_type == self.TaskType.PROCESS:
                if p_uid in self.processes and self.processes[p_uid].state != Process.State.COMPLETE:
                    process = self.processes[p_uid]
                    self.shepherd.log.warning(f'Process {process.p_uid} finished but was not marked complete.')

                del self.task_tbl[task]
                finished.remove(task)

        return local_new_tasks

    def _handle_output(self, finished):
        local_new_tasks = []
        theTasks = set(finished)

        for task in theTasks:
            stream_type, p_uid = self.task_tbl[task]
            if stream_type in frozenset([ProcessManager.TaskType.STDOUT, ProcessManager.TaskType.STDERR]):
                if p_uid in self.processes:
                    # If p_uid is not in self._processes it is because the process already ended and
                    # this finished task had yet to be processed. There is no process so the
                    # output is ignored but likely there is no output.
                    process = self.processes[p_uid]
                    # register another read output task
                    read_task = self.create_task(self.async_read_output(p_uid, stream_type), stream_type, p_uid)
                    process.stream_task[stream_type] = read_task
                    local_new_tasks.append(read_task)

                # remove it from the set of finished tasks
                del self.task_tbl[task]
                finished.remove(task)

        return local_new_tasks

    def _handle_new_stdin_tasks(self):
        local_new_tasks = []
        for p_uid, process in self.processes.items():
            # For each process look to see if something is waiting to be written
            # to stdin. If so and there is not a current outstanding task that
            # is waiting to write to stdin already (and the stream is not closing)
            # then create a new task to write to stdin. Don't empty the buffer here.
            # The buffer will be emptied when the task runs that will actually write
            # it to stdin.
            if process.state != Process.State.INIT:
                buf = process.stdin_buf
                if buf is not None and len(buf) > 0:
                    # We don't look at state here because it will be at least open and
                    # possibly closing. But, if there is data to write, then we write it
                    # before closing.
                    if self.TaskType.STDIN not in process.stream_task:
                        stdin_task = self.create_task(self.async_write_input(p_uid), self.TaskType.STDIN, p_uid)
                        process.stream_task[self.TaskType.STDIN] = stdin_task
                        local_new_tasks.append(stdin_task)
                elif process.stdin_state == Process.StdInState.CLOSE_CALLED:
                    process.stdin_state = Process.StdInState.CLOSING
                    stdin_task = self.create_task(self.async_close_stdin(p_uid), self.TaskType.STDIN, p_uid)
                    process.stream_task[self.TaskType.STDIN] = stdin_task
                    local_new_tasks.append(stdin_task)

        return local_new_tasks

    def _handle_finished_stdin_tasks(self, finished):

        theTasks = set(finished)

        for task in theTasks:
            stream_type, p_uid = self.task_tbl[task]
            if stream_type == self.TaskType.STDIN and p_uid in self.processes:
                process = self.processes[p_uid]
                del process.stream_task[self.TaskType.STDIN]
                del self.task_tbl[task]
                finished.remove(task)

    def _handle_new_process_creation(self):
        # We create the new process tasks here outside of any other
        # running task. Otherwise, if we were to create a new task from
        # within a running task, the task loop is already running and the
        # new task then gets scheduled and run immediately. By creating
        # the new task outside of the task loop, the task will run
        # when we want it to run.
        local_new_tasks = []

        for p in self.new_processes:
            task = self.create_task(self.async_spawn_app(p), self.TaskType.PROCESS, p.p_uid)
            local_new_tasks.append(task)

        # Now reset the new process list.
        self.new_processes = []

        return local_new_tasks

    async def async_spawn_app(self, p: Process) -> bool:
        """
        Async event to start cmd with arguments args and identifier p_uid.
        """
        try:
            # check to make sure there really is a process with this p_uid.
            if p.p_uid not in self.processes:
                self.shepherd.log.warning(f'No managed process with p_uid={p.p_uid} in async_spawn_app.')
                raise ValueError(f'No managed process with p_uid={p.p_uid} in async_spawn_app.')

            self.shepherd.log.info(f'Starting Process: {p.cmd!s} {p.args!s}')
            p._stat = await asyncio.create_subprocess_exec(p.cmd, *p.args,
                                                           stdout=asyncio.subprocess.PIPE,
                                                           stderr=asyncio.subprocess.PIPE,
                                                           stdin=asyncio.subprocess.PIPE,
                                                           env=p.env, cwd=p.cwd)

            self.shepherd.log.info('Process Started....')

            p.pid = p._stat.pid

            # If notification needs to be sent with the PID then send it.
            if p.resp_queue is not None:
                self.shepherd.log.info('SENDING SHProcessCreateResponse')
                p.resp_queue.send(
                    dmsg.SHProcessCreateResponse(tag=0, ref=p.tag_ref,
                                                 err=dmsg.SHProcessCreateResponse.Errors.SUCCESS).serialize())

        except Exception as ex:
            try:
                p.state = Process.State.FAIL
                p.resp_queue.send(
                    dmsg.SHProcessCreateResponse(tag=0, ref=p.tag_ref, err=dmsg.SHProcessCreateResponse.Errors.FAIL,
                                                 err_info=str(ex)).serialize())
                self.shepherd.log.info('Sent Failed ProcessCreateResponse')
            except:
                try:
                    self.shepherd.log.warning('Unable to send Failed ProcessCreateResponse')
                except:
                    pass

                self.shepherd.abend(sys.exc_info())
                raise ex

        return True

    async def async_app_done(self, p_uid: int) -> None:
        """
        Async event that completes if the managed application with identifier id has completed.
        """
        # check if we even have this identifier
        if p_uid not in self.processes:
            raise ValueError(f'no managed process with p_uid={p_uid}')

        process = self.processes[p_uid]
        await process._stat.wait()
        self.shepherd.log.info('In async_app_done: Process has exited.')
        self.shepherd.log.info(f'return code {process._stat.returncode!s}')

        # Call to cancel any timer set to go off for terminating the process. If no
        # timer exists, it won't hurt anything to call this.
        process.cancel_kill_timer(self.shepherd.log)

        if process.state == Process.State.RUN:
            process.state = Process.State.COMPLETE
            if process.term_queue is not None:
                self.shepherd.log.info('SENDING SHProcessExit message')
                process.term_queue.send(
                    dmsg.SHProcessExit(tag=0, p_uid=process.p_uid,
                                       exit_code=process._stat.returncode).serialize())

    async def async_read_output(self, p_uid: int, stream_type: TaskType) -> None:
        """
        Async event to read data from the managed process stream.
        """
        try:
            process = self.processes[p_uid]
            stream = process.stream[stream_type]

            if stream is not None:
                while True:
                    if stream.at_eof():
                        break

                    output = await stream.read(dfacts.MANAGED_PROCESS_MAX_OUTPUT_SIZE_PER_MESSAGE)

                    if output is None:
                        # Must be because the stream closed.
                        break

                    if len(output) <= 0:
                        # Possibly because the stream closed.
                        break

                    output_queue = process.output_queue[stream_type]

                    if output_queue is not None:
                        self.shepherd.log.info(f'SENDING SHFwdOutput: {output!s}')
                        fwd_output_msg = dmsg.SHFwdOutput(tag=0, idx=self.shepherd.node_index,
                                                          p_uid=process.p_uid, fd_num=stream_type.value,
                                                          data=output.decode())
                        output_queue.send(fwd_output_msg.serialize())

        except Exception as ex:
            self.shepherd.abend(sys.exc_info())
            raise ex

    async def async_write_input(self, p_uid: int) -> None:
        """
        Async event to write data to the managed process stdin.
        """
        try:
            process = self.processes[p_uid]
            stream = process.stream[self.TaskType.STDIN]
            buf = process.stdin_buf
            # reset the buffer so if while we await below, more
            # data is written, we get that the next time.
            process.stdin_buf = bytearray()
            self.shepherd.log.info('In async_write_input')
            if not stream.is_closing():
                self.shepherd.log.info(f'Now writing input {buf!s}')
                stream.write(buf)
                await stream.drain()
                self.shepherd.log.info('Done writing')
        except Exception as ex:
            self.shepherd.abend(sys.exc_info())
            raise ex

    async def async_close_stdin(self, p_uid: int) -> None:
        """
        Async event to close stdin.
        """
        process = self.processes[p_uid]
        stream = process.stream[self.TaskType.STDIN]
        if not stream.is_closing():
            stream.close()
            await stream.wait_closed()

        process.stdin_state = Process.StdInState.CLOSED

    def new_spawn_app(self, p: Process) -> None:
        """
        Create a task to start cmd with arguments args and identifier p_uid.
        """
        # check we do not already have a process
        # with this identifier

        # TODO should not report this as a ValueError probably
        if p.p_uid in self.processes:
            raise ValueError(f'existing managed process with p_uid={p.p_uid}')

        # add the process to our registry of processes
        self.processes[p.p_uid] = p

        # add the new process to our registry of new processes to be created.
        self.new_processes.append(p)

    def run_tasks(self, systemTasks) -> (int, int):
        """
        Run the task loop until one task returns complete. Returns a tuple
        of the set of finished system tasks and the number of remaining
        tasks to execute.
        """
        local_new_tasks = systemTasks

        # new process creation tasks are created and added
        # to the local_new_tasks list.
        local_new_tasks += self._handle_new_process_creation()

        # stdin handling
        local_new_tasks += self._handle_new_stdin_tasks()

        self.tasks.extend(local_new_tasks)

        finished = []
        unfinished = []

        # This deserves a little explanation. Asyncio in some cases
        # runs a task and completes it before the run_until_complete
        # executes below. This could happen as a result of a task
        # being created inside another task - but currently this is
        # not the case. However, it still happens under some race conditions
        # that a task completes outside of the run_until_complete below.
        # We don't want that to happen. So, we'll detect if it has happened here and flag
        # it if it happens. We'll also recover from it here by processing
        # tasks that have completed before running additional tasks.

        for task in self.tasks:
            if task.done():
                self.shepherd.log.info(f'already completed task: {task!s}')
                finished.append(task)
            elif task.cancelled():
                self.shepherd.log.info(f'cancelled task: {task!s}')
                finished.append(task)
            else:
                unfinished.append(task)

        if len(finished) == 0:
            # If the length of finished tasks is not zero then that means
            # we had a task that was already finished (see above) and we'll
            # skip running tasks right now until we have processed this
            # task that needs processing.

            if len(unfinished) == 0:
                # If there are no tasks to run and there are no finished tasks,
                # then return
                return [], 0

            finished_tasks, unfinished = asyncio.get_event_loop().run_until_complete(
                asyncio.wait(unfinished, timeout=1,
                             return_when=asyncio.FIRST_COMPLETED))

            finished.extend(finished_tasks)

        # Now accumulate the new tasks to add to the task list.
        local_new_tasks = []

        # first loop for started processes
        local_new_tasks += self._handle_started_procs(finished)

        # then handle any completed processes
        local_new_tasks += self._handle_done_procs(finished)

        # process any completed stdout and stderr reads
        local_new_tasks += self._handle_output(finished)

        self._handle_finished_stdin_tasks(finished)

        self.tasks = list(unfinished) + local_new_tasks

        # This removes the finished tasks being returned from the internal
        # reverse lookup map in the process manager.
        for task in finished:
            del self.task_tbl[task]

        return finished, len(self.tasks)

    def app_cleanup(self, p_uid: int) -> int:
        """
        Clean up all data structures associated with the process.  Returns
        the exit status of the process.
        """
        if p_uid not in self.processes:
            raise RuntimeError(f'cannot clean up non-existent process with p_uid={p_uid}')

        process = self.processes[p_uid]

        # make sure the app is not still running
        if process.state == Process.State.RUN:
            raise RuntimeError(f'cannot cleanup a running process with ID={p_uid}')

        # cancel any outstanding tasks on this app
        # There may be stdout and stderr tasks that may need clean up
        for task_id in process.stream_task:
            task = process.stream_task[task_id]
            if task is not None and not task.done():
                task.cancel()

        if process._stat is not None:
            rv = process._stat.returncode
        else:
            # None for the stat means that it never sucessfully ran.
            rv = INVALID_COMMAND

        del self.processes[p_uid]

        return rv

    def shutdown(self):
        """
        Called when the process manager is shutting down to do any clean up that is required.
        """
        proc_ids = list(self.processes.keys())
        for p_uid in proc_ids:
            self.app_cleanup(p_uid)

src/dragon/localservices/manager.py
DragonHPC/dragon
"""
Options objects for local services object creation messages.


All should have dict-like init/serialization.

TODO: lock type parameter in here?  Better handshake interface-wise
maybe between this and cython desirable.
"""


class ChannelOptions:
    def __init__(self, sattr='', capacity=None, block_size=None):
        self.capacity = capacity
        self.sattr = sattr
        self.block_size = block_size

    # wonder what __dict__(self) does.
    def get_sdict(self):
        return {'sattr': self.sattr,
                'capacity': self.capacity,
                'block_size': self.block_size}

    # this method seems kinda derpy in hindsight.
    @staticmethod
    def from_sdict(sdict):
        return ChannelOptions(**sdict)


class ProcessOptions:
    """Options for process creation at local services level.

    Attributes
     (e.g.)
     mon_std - whether to monitor stdout/stderr, write to stdin
               bool default True
    """
    def __init__(self, mon_std=True):
        self.mon_std = bool(mon_std)

    def get_sdict(self):
        return {'mon_std': self.mon_std}

    @staticmethod
    def from_sdict(sdict):
        return ProcessOptions(**sdict)

src/dragon/localservices/options.py
DragonHPC/dragon
class ProcessOptions:
    """Open type, placeholder, for Shepherd options on how to launch a process."""

    def __init__(self, dummy=None):
        self.dummy = dummy  # until there are things to put in here

    def get_sdict(self):
        return {}

    @staticmethod
    def from_sdict(sdict):
        return ProcessOptions(**sdict)

src/dragon/localservices/process_desc.py
DragonHPC/dragon
import os
import traceback
import time
import logging
import queue
import threading
import io
import subprocess
import json
import collections
import selectors
import signal
import socket

from .. import channels as dch
from .. import managed_memory as dmm

from .. import pmod
from .. import utils as dutils
from ..infrastructure import messages as dmsg
from ..infrastructure import util as dutil
from ..infrastructure import facts as dfacts
from ..infrastructure import parameters as parms
from ..infrastructure import connection as dconn
from ..infrastructure import parameters as dp

from ..dlogging import util as dlog
from ..dlogging.util import DragonLoggingServices as dls
from ..utils import B64

_TAG = 0
_TAG_LOCK = threading.Lock()

def get_new_tag():
    global _TAG
    with _TAG_LOCK:
        tmp = _TAG
        _TAG += 1
    return tmp

ProcessProps = collections.namedtuple('ProcessProps', ['p_uid', 'critical', 'r_c_uid',
                                      'stdin_req', 'stdout_req', 'stderr_req',
                                      'stdin_connector', 'stdout_connector', 'stderr_connector'])

class PopenProps(subprocess.Popen):

    def __init__(self, props: ProcessProps, *args, **kwds):
        assert isinstance(props, ProcessProps)
        super().__init__(*args, **kwds)
        self.props = props
        # TODO Add affinity control to the process options. To be on the safe
        # TODO side for now, open affinity to all cores.
        os.sched_setaffinity(self.pid, range(os.cpu_count()))
        # XXX Affinity settings are only inherited by grandchild processes
        # XXX created after this point in time. Any grandchild processes
        # XXX started when the child process starts up most certainly will
        # XXX not have the desired affinity. To guarantee CPU affinity settings
        # XXX processes should be launched with e.g. taskset(1) or otherwise
        # XXX be configured to set their own affinity as appropriate.

class TerminationException(Exception):
    pass

class InputConnector:
    def __init__(self, conn:dconn.Connection):
        self._proc = None
        self._conn = conn
        self._log = logging.getLogger('input connector')
        self._closed = False

    def add_proc_info(self, proc:PopenProps):
        self._proc = proc

    def poll(self):
        return self._conn.poll()

    def forward(self):
        try:
            data = None
            while self._conn.poll(timeout=0):
                data = self._conn.recv()
                self._proc.stdin.write(data.encode('utf-8'))
                self._proc.stdin.flush()
                self._log.info(f'Stdin data that was written:{data}')
        except EOFError as ex:
            return True
        except TimeoutError:
            return False
        except:
            self._log.info(f'The input could not be forwarded from cuid={self._conn.inbound_chan.cuid} to process pid={self._proc.pid}')

        return False

    def __hash__(self):
        return hash(self._conn.inbound_channel.cuid)

    def __eq__(self, other):
        # since all cuid's are unique we can compare cuids. In addition,
        # if we want we can ask if a cuid is in a set of connectors in O(1) time.
        return hash(self) == hash(other)

    def close(self):
        if self._closed:
            return

        try:
            if self._conn is not None:
                if self.poll():
                    self.forward()
                self._conn.close()
        except:
            pass

        try:
            self._proc.stdin.flush()
            self._proc.stdin.close()
        except:
            pass

        self._closed = True

    @property
    def proc_is_alive(self):
        return self._proc.returncode is None

    @property
    def conn(self):
        return self._conn

    @property
    def cuid(self):
        if self._conn is None:
            return None

        return self._conn.inbound_channel.cuid

class OutputConnector:
    def __init__(self, be_in, puid, hostname, out_err, conn:dconn.Connection = None, root_proc=None, critical_proc=False):
        self._be_in = be_in
        self._puid = puid
        self._hostname = hostname
        self._out_err = out_err
        self._conn = conn
        self._root_proc = root_proc
        self._log = logging.getLogger('output connector')
        self._writtenTo = False
        self._critical_proc = critical_proc
        self._closed = False
        self._proc = None

    def __hash__(self):
        return hash(self._conn.outbound_channel.cuid)

    def __eq__(self, other):
        # since all cuid's are unique we can compare cuids. In addition,
        # if we want we can ask if a cuid is in a set of connectors in O(1) time.
        return hash(self) == hash(other)

    def _sendit(self, block):
        if len(block) > 0:
            self._writtenTo = True

        if self._conn is None:
            self._be_in.send(dmsg.SHFwdOutput(tag=get_new_tag(), idx=parms.this_process.index,
                    p_uid=self._puid, data=block,
                    fd_num=self._out_err,
                    pid=self._proc.pid, hostname=self._hostname).serialize())
            return

        try:
            # A process has requested that output be forwarded from this process to it.
            # Because a connection could be owned by a parent and if the parent
            # exits the connection could be destroyed, we'll use exception handling
            # here and as a backup, we'll forward lost output to the launcher.
            self._conn.send(block)
        except:
            self._be_in.send(dmsg.SHFwdOutput(tag=get_new_tag(), idx=parms.this_process.index,
                        p_uid=self._puid, data='[orphaned output]: '+block,
                        fd_num=self._out_err,
                        pid=self._proc.pid, hostname=self._hostname).serialize())

    def add_proc_info(self, proc):
        self._proc = proc

    def forward(self, data):
        if self._conn is not None:
            while len(data) > 300:
                chunk = data[:300]
                self._sendit(chunk)
                data = data[300:]

            if len(data) > 0:
                self._sendit(data)
        else:
            self._sendit(data)

    def flush(self):

        is_stderr = self._out_err == dmsg.SHFwdOutput.FDNum.STDERR.value

        try:
            file_obj = self.file_obj

            if file_obj is not None:
                io_data = file_obj.read(dmsg.SHFwdOutput.MAX)
            else:
                io_data = None

        except EOFError:
            io_data = None
        except ValueError:  # race, file could be closed
            io_data = None
        except OSError: # race, file is closed
            io_data = None

        if not io_data:  # at EOF because we just selected
            return True # To indicate EOF

        str_data = io_data.decode()

        if not is_stderr and self._puid == dfacts.GS_PUID:
            raise TerminationException(str_data)
        else:
            self.forward(str_data)

        if self._critical_proc and is_stderr:
            raise TerminationException(str_data)

        return False

    def close(self):
        if self._closed:
            return

        self.flush()

        try:
            self.file_obj.close()
        except:
            pass

        if not self._root_proc:
            # Don't call close on the connection unless a root proc.
            # Children share the same connection so we don't want to
            # close it from a child.
            self._closed = True
            return

        if self._conn is not None and not self._writtenTo:
            # If it is not written to yet, the connection must be
            # written to before it is closed so that EOF gets signaled
            # for the receiving process.
            try:
                self._conn.send('')
            except:
                pass
            self._writtenTo = True

        if self._conn is not None:
            try:
                self._conn.close()
            except:
                pass

        self._closed = True

    @property
    def file_obj(self):
        is_stderr = self._out_err == dmsg.SHFwdOutput.FDNum.STDERR.value

        if is_stderr:
            return self._proc.stderr

        return self._proc.stdout

    @property
    def proc_is_alive(self):
        return self._proc.returncode is None

    @property
    def puid(self):
        return self._puid

    @property
    def conn(self):
        return self._conn

    @property
    def cuid(self):
        if self._conn is None:
            return None
        return self._conn.outbound_channel.cuid

def mk_response_pairs(resp_cl, ref):
    err_cl = resp_cl.Errors

    def success_resp(desc=None, **kwargs):
        if desc is not None:
            return resp_cl(tag=get_new_tag(), ref=ref, err=err_cl.SUCCESS, desc=desc, **kwargs)
        else:
            return resp_cl(tag=get_new_tag(), ref=ref, err=err_cl.SUCCESS, **kwargs)


    def fail_resp(msg):
        return resp_cl(tag=get_new_tag(), ref=ref, err=err_cl.FAIL, err_info=msg)

    return success_resp, fail_resp


def mk_output_connection_over_channel(ch_desc):

    channel_descriptor = B64.str_to_bytes(ch_desc)
    the_channel = dch.Channel.attach(channel_descriptor)
    return dconn.Connection(outbound_initializer=the_channel,
                            options=dconn.ConnectionOptions(min_block_size=512), policy=dp.POLICY_INFRASTRUCTURE)

def mk_input_connection_over_channel(ch_desc):

    channel_descriptor = B64.str_to_bytes(ch_desc)
    the_channel = dch.Channel.attach(channel_descriptor)
    return dconn.Connection(inbound_initializer=the_channel,
                            options=dconn.ConnectionOptions(min_block_size=512), policy=dp.POLICY_INFRASTRUCTURE)


class LocalServer:
    """Handles shepherd messages in normal processing.

    This object does not handle startup/teardown - instead
     it expects to be given whatever channels/pools have been made for it
     by startup, handles to what it needs to talk to in normal processing,
     and offers a 'run' and 'cleanup' method.
    """

    _DTBL = {}  # dispatch router, keyed by type of shepherd message

    SHUTDOWN_RESP_TIMEOUT = 0.010  # seconds, 10 ms
    QUIESCE_TIME = 1  # seconds, 1 second, join timeout for thread shutdown.

    def __init__(self, channels=None, pools=None,
                 transport_test_mode=False,
                 hostname='NONE'):

        self.transport_test_mode = transport_test_mode
        self.new_procs = queue.SimpleQueue()  # outbound PopenProps of newly started processes
        self.new_channel_input_monitors = queue.SimpleQueue()
        self.exited_channel_output_monitors = queue.SimpleQueue()
        self.hostname = hostname
        self.cuid_to_input_connector = {}

        if channels is None:
            self.channels = {}  # key c_uid value channel
        else:
            self.channels = channels
        if pools is None:
            self.pools = {}  # key m_uid value memory pool
        else:
            self.pools = pools

        self.apt = {}  # active process table. key: pid, value PopenProps obj
        self.puid2pid = {}  # key: p_uid, value pid
        self.apt_lock = threading.Lock()

        self.shutdown_sig = threading.Event()
        self.gs_shutdown_sig = threading.Event()
        self.ta_shutdown_sig = threading.Event()
        self.stashed_threading_excepthook = None
        self.exit_reason = None

    def _logging_ex_handler(self, args):
        log = logging.getLogger('fatal exception')
        ex_type, ex_value, ex_tb, thread = args
        if ex_type is SystemExit:
            return

        buf = io.BytesIO()
        wrap = io.TextIOWrapper(buf, write_through=True)
        traceback.print_exception(ex_type, ex_value, ex_tb, file=wrap)
        log.error(f'from {thread.name}:\n{buf.getvalue().decode()}')
        self._abnormal_termination(f'from {thread.name}:\n{buf.getvalue().decode()}')

    def set_shutdown(self, msg):
        log = logging.getLogger('shutdown event')
        self.shutdown_sig.set()
        log.info(f'shutdown called after receiving {repr(msg)}')

    def check_shutdown(self):
        return self.shutdown_sig.is_set()

    def set_gs_shutdown(self):
        log = logging.getLogger('gs shutdown event')
        self.gs_shutdown_sig.set()
        log.info('set GS shutdown')

    def check_gs_shutdown(self):
        return self.gs_shutdown_sig.is_set()

    def set_ta_shutdown(self):
        log = logging.getLogger('ta shutdown event')
        self.ta_shutdown_sig.set()
        log.info('set TA shutdown')

    def check_ta_shutdown(self):
        return self.ta_shutdown_sig.is_set()

    def __str__(self):
        with self.apt_lock:
            plist = [f'\t{p_uid}:{pid} {self.apt[pid]!s}'
                     for p_uid, pid in self.puid2pid.items()]

        procs = '\n'.join(plist)
        chans = ' '.join([f'{k!s}' for k in self.channels.keys()])
        pools = ' '.join([f'{k!s}' for k in self.pools.keys()])

        return f'Procs:\n{procs}\nChans:\n{chans}\nPools:\n{pools}'

    def add_proc(self, proc):
        with self.apt_lock:
            self.apt[proc.pid] = proc
            self.puid2pid[proc.props.p_uid] = proc.pid

        self.new_procs.put(proc)

    @staticmethod
    def clean_pools(pools, log):
        log.info(f'{len(pools)} pools outstanding')
        for m_uid, pool in pools.items():
            try:
                pool.destroy()
            except (dmm.DragonMemoryError, dmm.DragonPoolError) as dpe:
                log.warning(f'm_uid={m_uid} failed: {dpe!s}')

    def _clean_procs(self):
        log = logging.getLogger('kill procs')

        with self.apt_lock:
            log.info(f'{len(self.apt)} processes outstanding')
            for p_uid, pid in self.puid2pid.items():
                proc = self.apt[pid]
                try:
                    proc.kill()
                    log.info(f'kill sent to p_uid={p_uid}:proc.pid={proc.pid}')
                except (subprocess.SubprocessError, OSError) as ose:
                    log.warning(f'kill on p_uid={p_uid}: prod.pid={proc.pid} failed: {ose}')

            self.puid2pid = {}

            for proc in self.apt.values():
                try:
                    proc.wait(10)

                except subprocess.SubprocessError as spe:
                    log.warning(f'wait on puid={proc.props.p_uid} failed: {spe}')

            self.apt = {}

    def cleanup(self):
        """Tries to destroy channels and pools and kill outstanding processes.

        None of the other threads should be running at this point.
        """
        log = logging.getLogger('cleanup')

        log.info('start')

        # clean outstanding processes
        self._clean_procs()

        log.info(f'{len(self.channels)} channels outstanding')
        for c_uid, chan in self.channels.items():
            try:
                chan.destroy()
            except dch.ChannelError as cse:
                log.warning(f'c_uid={c_uid} failed: {cse!s}')
        self.channels = {}

        self.clean_pools(self.pools, log)
        self.pools = {}

        log.info('end')

    def _abnormal_termination(self, error_str):
        """Triggers LS abnormal termination.
        Sends AbnormalTermination message to Launcher BE and logs.

        :param error_str: error message with the cause of abnormal termination
        :type error_str: string
        """
        log = logging.getLogger('Abnormal termination')
        try:
            self.be_in.send(dmsg.AbnormalTermination(tag=get_new_tag(), err_info=error_str).serialize())
            log.critical(f"Abnormal termination sent to launcher be: {error_str}")
        except Exception as ex:
            log.exception(f'Abnormal termination exception: {ex}')

    def run(self, shep_in, gs_in, be_in, is_primary, ta_in=None, gw_channels=None):
        """Local services main function.

        :param shep_in: ls channel
        :type shep_in: Connection object
        :param gs_in: global services channel
        :type gs_in: Connection object
        :param be_in: launcher be channel
        :type be_in: Connection object
        :param is_primary: indicates if this is the primary LS or not
        :type is_primary: bool
        :param ta_in: transport agent channel, defaults to None
        :type ta_in: Connection object, optional
        :param gw_channels: list of gateway channels for multinode only, defaults to None
        :type gw_channels: list, optional
        """
        log = logging.getLogger('ls run')
        log.info('start')

        if gw_channels is None:
            gw_channels = []

        self.ta_in = ta_in
        self.be_in = be_in
        self.gs_in = gs_in
        self.is_primary = is_primary

        th = threading.Thread
        threads = [th(name='output mon', target=self.watch_output, daemon=True),
                   th(name='watch death', target=self.watch_death, daemon=True),
                   th(name='input mon', target=self.watch_input, daemon=True)]

        self.stashed_threading_excepthook = threading.excepthook
        threading.excepthook = self._logging_ex_handler

        try:
            log.info('starting runtime service threads')
            for th in threads:
                th.start()
            log.info('runtime service threads started')
        except Exception as ex:
            self._abnormal_termination(f'ls run starting threads exception: {ex}')

        try:
            self.main_loop(shep_in)
        except Exception as ex:
            tb = traceback.format_exc()
            self._abnormal_termination(f'ls main loop exception: {ex}\n{tb}')

        threading.excepthook = self.stashed_threading_excepthook

        try:
            for th in threads:
                th.join(self.QUIESCE_TIME)
            for th in threads:
                if th.is_alive():
                    log.error(f'thread {th.name} seems to have hung!')
        except Exception as ex:
            self._abnormal_termination(f'ls run joining threads exception: {ex}')

        if self.transport_test_mode and self.exit_reason is None:
            # In the transport test mode, normal set_shutdown was called.
            log.info('beginning normal shutdown during transport test mode run')

        try:
            # Destroy gateway channels
            gw_count = 0
            for id, gw_ch in enumerate(gw_channels):
                gw_ch.destroy()
                try:
                    del os.environ[dfacts.GW_ENV_PREFIX+str(id+1)]
                except KeyError:
                    pass
                gw_count += 1
            log.info(f'ls isPrimary={self.is_primary} destroyed {gw_count} gateway channels')
        except Exception as ex:
            self._abnormal_termination(f'ls run destroying gateway channels exception: {ex}')

        try:
            # m12 Send SHHaltBE to BE
            # tell launcher be to shut mrnet down and detach from logging
            log.info('m12 transmitting final messsage from ls SHHaltBE')
            dlog.detach_from_dragon_handler(dls.LS)
            self.be_in.send(dmsg.SHHaltBE(tag=get_new_tag()).serialize())
        except Exception as ex:
            self._abnormal_termination(f'ls run sending SHHaltBE exception: {ex}')

        log.info('exit')

    def main_loop(self, shep_rh):
        """Monitors the main LS input channel and receives messages.
        If the received message is not one of the expected that are
        handled by corresponding route decorators, then it signals
        abnormal termination of LS.

        :param shep_rh: ls input channel
        :type shep_rh: Connection object
        """
        log = logging.getLogger('main loop')
        log.info('start')

        while not self.check_shutdown():
            msg_pre = shep_rh.recv()

            if msg_pre is None:
                continue

            if isinstance(msg_pre, str):
                try:
                    msg = dmsg.parse(msg_pre)
                except (json.JSONDecodeError, KeyError, NotImplementedError, ValueError) as err:
                    self._abnormal_termination(f'msg\n{msg_pre}\nfailed parse!\n{err!s}')
                    continue
            else:
                msg = msg_pre

            if type(msg) in LocalServer._DTBL:
                resp_msg = self._DTBL[type(msg)][0](self, msg=msg)
                if resp_msg is not None:
                    self._send_response(target_uid=msg.r_c_uid, msg=resp_msg)
            else:
                self._abnormal_termination(f'unexpected msg type: {repr(msg)}')

        log.info('exit')

    def _send_response(self, target_uid, msg):
        """Sends response to either Global Services or Launcher BE
        depending on the target/return channel uid. Signals
        abnormal termination in the case of an unknown r_c_uid.

        :param target_uid: return channel uid, r_c_uid
        :type target_uid: int
        :param msg: message for the response
        :type msg: string
        """
        if target_uid == dfacts.GS_INPUT_CUID:
            self.gs_in.send(msg.serialize())
        elif target_uid == dfacts.launcher_cuid_from_index(parms.this_process.index):
            self.be_in.send(msg.serialize())
        else:
            self._abnormal_termination(f'unknown r_c_uid: {repr(msg)}')

    @dutil.route(dmsg.SHPoolCreate, _DTBL)
    def create_pool(self, msg: dmsg.SHPoolCreate) -> None:
        log = logging.getLogger('create pool')
        success, fail = mk_response_pairs(dmsg.SHPoolCreateResponse, msg.tag)

        error = ''
        if msg.m_uid in self.pools:
            error = f'msg.m_uid={msg.m_uid!s} already in use'

        if not error:
            try:
                mpool = dmm.MemoryPool(msg.size, msg.name, msg.m_uid)
            except (dmm.DragonPoolError, dmm.DragonMemoryError) as dme:
                error = f'{msg!r} failed: {dme!s}'

        if error:
            log.warning(error)
            resp_msg = fail(error)
        else:
            self.pools[msg.m_uid] = mpool
            encoded_desc = B64.bytes_to_str(mpool.serialize())
            resp_msg = success(encoded_desc)

        return resp_msg

    @dutil.route(dmsg.SHPoolDestroy, _DTBL)
    def destroy_pool(self, msg: dmsg.SHPoolDestroy) -> None:
        log = logging.getLogger('destroy pool')
        success, fail = mk_response_pairs(dmsg.SHPoolDestroyResponse, msg.tag)

        error = ''
        if msg.m_uid not in self.pools:
            error = f'msg.m_uid={msg.m_uid!s} does not exist'

        if not error:
            mpool = self.pools.pop(msg.m_uid)
            try:
                mpool.destroy()
            except (dmm.DragonPoolError, dmm.DragonMemoryError) as dme:
                error = f'{msg!r} failed: {dme!s}'

        if error:
            log.warning(error)
            resp_msg = fail(error)
        else:
            resp_msg = success()

        return resp_msg

    @dutil.route(dmsg.SHChannelCreate, _DTBL)
    def create_channel(self, msg: dmsg.SHChannelCreate) -> None:
        log = logging.getLogger('create channel')
        log.info("Received an SHChannelCreate")
        success, fail = mk_response_pairs(dmsg.SHChannelCreateResponse, msg.tag)

        error = ''
        if msg.c_uid in self.channels:
            error = f'msg.c_uid={msg.c_uid!s} already in use'

        if msg.m_uid not in self.pools:
            error = f'msg.m_uid={msg.m_uid!s} does not exist'

        if not error:
            try:
                ch = dch.Channel(mem_pool=self.pools[msg.m_uid], c_uid=msg.c_uid,
                                 block_size=msg.options.block_size,
                                 capacity=msg.options.capacity)
            except dch.ChannelError as cex:
                error = f'{msg!r} failed: {cex!s}'

        if error:
            log.warning(error)
            resp_msg=fail(error)
        else:
            self.channels[msg.c_uid] = ch
            encoded_desc = B64.bytes_to_str(ch.serialize())
            resp_msg = success(encoded_desc)
            log.info("Received and Created a channel via SHChannelCreate")

        return resp_msg

    @dutil.route(dmsg.SHChannelDestroy, _DTBL)
    def destroy_channel(self, msg: dmsg.SHChannelDestroy) -> None:
        log = logging.getLogger('destroy channel')
        success, fail = mk_response_pairs(dmsg.SHChannelDestroyResponse, msg.tag)

        error = ''
        if msg.c_uid not in self.channels:
            error = f'{msg.c_uid} does not exist'

        if not error:
            ch = self.channels.pop(msg.c_uid)
            try:
                ch.destroy()
            except dch.ChannelError as cex:
                error = f'{msg!r} failed: {cex!s}'

        if error:
            log.warning(error)
            resp_msg = fail(error)
        else:
            resp_msg = success()

        return resp_msg

    @dutil.route(dmsg.SHProcessCreate, _DTBL)
    def create_process(self, msg: dmsg.SHProcessCreate) -> None:
        log = logging.getLogger('create process')
        success, fail = mk_response_pairs(dmsg.SHProcessCreateResponse, msg.tag)

        if msg.t_p_uid in self.puid2pid:
            error = f'msg.t_p_uid={msg.t_p_uid} already exists'
            log.warning(error)
            self._send_response(target_uid=msg.r_c_uid, msg=fail(error))
            return

        if not msg.rundir:
            working_dir = None
        else:
            working_dir = msg.rundir

        # TODO for multinode: whose job is it to update which parameters?
        log.debug(f'The number of gateways per node is configured to {parms.this_process.num_gw_channels_per_node}')
        log.debug(f'Removing these from environment: {parms.LaunchParameters.NODE_LOCAL_PARAMS}')
        req_env = dict(msg.env)
        parms.LaunchParameters.remove_node_local_evars(req_env)
        the_env = dict(os.environ)
        the_env.update(req_env)

        stdin_conn = None
        stdin_resp = None
        stdout_conn = None
        stdout_resp = None
        stderr_conn = None
        stderr_resp = None
        stdout_root = False
        stderr_root = False

        if msg.stdin_msg is not None:
            stdin_resp = self.create_channel(msg.stdin_msg)
            if stdin_resp.err != dmsg.SHChannelCreateResponse.Errors.SUCCESS:
                resp_msg = fail(f'Failed creating the stdin channel for new process: {stdin_resp.err_info}')
                return resp_msg
            stdin_conn = mk_input_connection_over_channel(stdin_resp.desc)

        if msg.stdout_msg is not None:
            stdout_resp = self.create_channel(msg.stdout_msg)
            if stdout_resp.err != dmsg.SHChannelCreateResponse.Errors.SUCCESS:
                # TBD: We need to destroy the stdin channel if it exists
                resp_msg = fail(f'Failed creating the stdout channel for new process: {stdout_resp.err_info}')
                return resp_msg
            stdout_conn = mk_output_connection_over_channel(stdout_resp.desc)
            stdout_root = True
            the_env[dfacts.STDOUT_DESC] = stdout_resp.desc
        elif dfacts.STDOUT_DESC in the_env:
            # Finding the STDOUT descriptor in the environment
            # means a parent requested PIPE and so all children
            # inherit this as well. If a child of the parent requested
            # PIPE itself, then it would have been caught above.
            stdout_conn = mk_output_connection_over_channel(the_env[dfacts.STDOUT_DESC])

        if msg.stderr_msg is not None:
            stderr_resp = self.create_channel(msg.stderr_msg)
            if stderr_resp.err != dmsg.SHChannelCreateResponse.Errors.SUCCESS:
                # TBD: We need to destroy the stdin and stdout channels if they exist
                resp_msg = fail(f'Failed creating the stderr channel for new process: {stderr_resp.err_info}')
                return resp_msg
            stderr_conn = mk_output_connection_over_channel(stderr_resp.desc)
            stderr_root = True
            the_env[dfacts.STDERR_DESC] = stderr_resp.desc
        elif msg.stderr == dmsg.STDOUT:
            # We put stderr connection in environment so any subprocesses
            # will also write to the stdout connection.
            stderr_conn = stdout_conn
            the_env[dfacts.STDERR_DESC] = stdout_resp.desc
        elif dfacts.STDERR_DESC in the_env:
            # Finding the STDERR descriptor in the environment
            # means inherit it. See above for STDOUT explanation.
            stderr_conn = mk_output_connection_over_channel(the_env[dfacts.STDERR_DESC])

        real_args = [msg.exe] + msg.args

        try:
            stdout = subprocess.PIPE
            if msg.stdout == subprocess.DEVNULL:
                # if user is requesting devnull, then we'll start
                # process that way. Otherwise, LS gets output
                # via PIPE forwards it where requested.
                stdout = subprocess.DEVNULL

            stderr = subprocess.PIPE
            if msg.stderr == subprocess.DEVNULL:
                # Same handling as stdout explanation above.
                stderr = subprocess.DEVNULL

            if msg.stderr == subprocess.STDOUT:
                stderr = subprocess.STDOUT

            if msg.pmi_info:
                log.debug(f'{msg.pmi_info}')
                log.info(f'p_uid {msg.t_p_uid} looking up pmod launch cuid')
                pmod_launch_cuid = dfacts.pmod_launch_cuid_from_jobinfo(
                    dutils.host_id(),
                    msg.pmi_info.job_id,
                    msg.pmi_info.lrank
                )

                log.info(f'p_uid {msg.t_p_uid} Creating pmod launch channel using {pmod_launch_cuid=}')
                node_index = parms.this_process.index
                inf_muid = dfacts.infrastructure_pool_muid_from_index(node_index)
                pmod_launch_ch = dch.Channel(self.pools[inf_muid], pmod_launch_cuid)
                the_env['DRAGON_PMOD_CHILD_CHANNEL'] = str(dutils.B64(pmod_launch_ch.serialize()))

                log.info(f'p_uid {msg.t_p_uid} Setting required PMI environment variables')
                the_env['PMI_CONTROL_PORT'] = str(msg.pmi_info.control_port)
                the_env['MPICH_OFI_CXI_PID_BASE'] = str(msg.pmi_info.pid_base)
                the_env['DL_PLUGIN_RESILIENCY'] = "1"
                the_env['LD_PRELOAD'] = 'libdragon.so'
                the_env['_DRAGON_PALS_ENABLED'] = '1'
                the_env['FI_CXI_RX_MATCH_MODE'] = 'hybrid'
                # the_env['DRAGON_DEBUG'] = '1'
                # the_env['PMI_DEBUG'] = '1'

            stdin_connector = InputConnector(stdin_conn)

            stdout_connector = OutputConnector(be_in = self.be_in, puid=msg.t_p_uid,
                    hostname=self.hostname, out_err=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                    conn=stdout_conn, root_proc=stdout_root, critical_proc=False)

            stderr_connector = OutputConnector(be_in = self.be_in, puid=msg.t_p_uid,
                    hostname=self.hostname, out_err=dmsg.SHFwdOutput.FDNum.STDERR.value,
                    conn=stderr_conn, root_proc=stderr_root, critical_proc=False)

            with self.apt_lock:  # race with death watcher; hold lock to get process in table.
                # The stdout_conn and stderr_conn will be filled in just below.
                the_proc = PopenProps(
                    ProcessProps(p_uid=msg.t_p_uid, critical=False, r_c_uid=msg.r_c_uid,
                        stdin_req=msg.stdin, stdout_req=msg.stdout, stderr_req=msg.stderr,
                        stdin_connector=stdin_connector, stdout_connector=stdout_connector,
                        stderr_connector=stderr_connector),
                    real_args,
                    bufsize=0,
                    stdin=subprocess.PIPE,
                    stdout=stdout,
                    stderr=stderr,
                    cwd=working_dir,
                    env=the_env
                )

                stdout_connector.add_proc_info(the_proc)
                stderr_connector.add_proc_info(the_proc)
                stdin_connector.add_proc_info(the_proc)

                if msg.stdin == dmsg.PIPE:
                    self.cuid_to_input_connector[msg.stdin_msg.c_uid] = stdin_connector
                    self.new_channel_input_monitors.put(stdin_connector)

                self.puid2pid[msg.t_p_uid] = the_proc.pid
                self.apt[the_proc.pid] = the_proc
                log.info(f'Now created process with args {real_args} and pid={the_proc.pid}')

            if msg.pmi_info:
                log.info(f'p_uid {msg.t_p_uid} sending mpi data for {msg.pmi_info.lrank}')
                pmod.PMOD(
                    msg.pmi_info.ppn,
                    msg.pmi_info.nid,
                    msg.pmi_info.nnodes,
                    msg.pmi_info.nranks,
                    msg.pmi_info.nidlist,
                    msg.pmi_info.hostlist,
                    msg.pmi_info.job_id
                ).send_mpi_data(msg.pmi_info.lrank, pmod_launch_ch)
                log.info(f'p_uid {msg.t_p_uid} DONE: sending mpi data for {msg.pmi_info.lrank}')

            log.info(f'p_uid {msg.t_p_uid} created as {the_proc.pid}')
            self.new_procs.put(the_proc)
            if msg.initial_stdin is not None and msg.initial_stdin != '':
                # we are asked to provide a string to the started process.
                log.info(f'Writing {msg.initial_stdin} to newly created process')
                proc_stdin = os.fdopen(the_proc.stdin.fileno(), 'wb')
                proc_stdin_send = dutil.NewlineStreamWrapper(proc_stdin, read_intent=False)
                proc_stdin_send.send(msg.initial_stdin)
                log.info('The provided string was written to stdin of the process by local services.')

            resp_msg = success(stdin_resp=stdin_resp, stdout_resp=stdout_resp, stderr_resp=stderr_resp)
        except (OSError, ValueError) as popen_err:
            error = f'{msg!r} encountered {popen_err}'
            log.warning(error)
            resp_msg = fail(error)

        return resp_msg

    @dutil.route(dmsg.SHProcessKill, _DTBL)
    def kill_process(self, msg: dmsg.SHProcessKill) -> None:
        log = logging.getLogger('kill process')
        success, fail = mk_response_pairs(dmsg.SHProcessKillResponse, msg.tag)

        try:
            target = self.puid2pid[msg.t_p_uid]
        except KeyError:
            error = f'{msg.t_p_uid} not present'
            log.warning(error)
            self._send_response(target_uid=msg.r_c_uid, msg=fail(error))
            return

        try:
            os.kill(target, msg.sig)
            log.info(f'{msg!r} delivered to pid {target}')
            resp_msg = success()
        except OSError as ose:
            error = f'delivering {msg!r} to pid {target} encountered {ose}'
            log.warning(error)
            resp_msg = fail(error)

        return resp_msg

    @dutil.route(dmsg.SHFwdInput, _DTBL)
    def fwd_input(self, msg: dmsg.SHFwdInput) -> None:
        log = logging.getLogger('fwd input handler')

        target = msg.t_p_uid
        error = ''
        with self.apt_lock:
            if target in self.puid2pid:
                targ_proc = self.apt[self.puid2pid[target]]
                if targ_proc.stdin is None:
                    error = f'p_uid {target} has no stdin'
            else:
                targ_proc = None
                error = f'p_uid {target} does not exist here and now'

        if not error:
            input_sel = selectors.DefaultSelector()
            input_sel.register(targ_proc.stdin, selectors.EVENT_WRITE)
            sel = input_sel.select(timeout=self.SHUTDOWN_RESP_TIMEOUT)
            if sel:
                try:
                    output_data = msg.input.encode()
                    if len(output_data) > dmsg.SHFwdInput.MAX:
                        log.warning(f'truncating request of {len(output_data)} to {dmsg.SHFwdInput.MAX}')

                    fh = sel[0][0].fileobj
                    fh.write(output_data[:dmsg.SHFwdInput.MAX])
                except (OSError, BlockingIOError) as err:
                    error = f'{err!s}'
            else:
                error = f'input of target={target} not ready for writing'
            input_sel.close()

        if error:
            log.warning(f'error={error} from {msg!s}')
            if targ_proc is not None and targ_proc.stdin is not None:
                targ_proc.stdin.close()
                targ_proc.stdin = None

        if msg.confirm:
            success, fail = mk_response_pairs(dmsg.SHFwdInputErr, msg.tag)
            if error:
                resp_msg = fail(error)
            else:
                resp_msg = success()

            return resp_msg

    @dutil.route(dmsg.AbnormalTermination, _DTBL)
    def handle_abnormal_term(self, msg: dmsg.AbnormalTermination) -> None:
        log = logging.getLogger('abnormal termination')
        log.info("received abnormal termination signal. starting shutdown process.")
        self._abnormal_termination(msg.err_info)

    @dutil.route(dmsg.GSHalted, _DTBL)
    def handle_gs_halted(self, msg: dmsg.GSHalted) -> None:
        self.set_gs_shutdown()
        log = logging.getLogger('forward GSHalted msg')
        log.info('is_primary=True and GSHalted recvd')
        self.be_in.send(msg.serialize())

    @dutil.route(dmsg.SHTeardown, _DTBL)
    def teardown_ls(self, msg: dmsg.SHTeardown) -> None:
        log = logging.getLogger('teardown LS')
        log.info(f'isPrimary={self.is_primary} handling SHTeardown')
        self.set_shutdown(msg)

    @dutil.route(dmsg.SHHaltTA, _DTBL)
    def handle_halting_ta(self, msg):
        log = logging.getLogger('forward SHHaltTA msg')
        log.info(f'handling {msg!s}')
        # m8 Forward SHHaltTA to TA
        self.ta_in.send(msg.serialize())

    @dutil.route(dmsg.TAHalted, _DTBL)
    def handle_ta_halted(self, msg):
        self.set_ta_shutdown()
        log = logging.getLogger('forward TAHalted msg')
        log.info(f'handling {msg!s}')
        self.be_in.send(msg.serialize())

    @dutil.route(dmsg.SHDumpState, _DTBL)
    def dump_state(self, msg: dmsg.SHDumpState) -> None:
        log = logging.getLogger('dump state')
        the_dump = f'{self!s}'
        if msg.filename is None:
            log.info('\n' + the_dump)
        else:
            try:
                with open(msg.filename, 'w') as dump_fh:
                    dump_fh.write(the_dump)

                log.info(f'to {msg.filename}')
            except (IOError, OSError) as e:
                log.warning(f'failed: {e!s}')

    def watch_death(self):
        """Thread monitors the demise of child processes of this process.

        Not all children do we care about; only the ones in our process group.

        :return: None, but exits on self.check_shutdown()
        """

        log = logging.getLogger('watch death')
        log.info('starting')

        while not self.check_shutdown():
            try:
                died_pid, exit_status = os.waitpid(0, os.WNOHANG)
            except ChildProcessError:  # no child processes at the moment
                # There is no error here. There just isn't a child process.
                died_pid, exit_status = (0, 0)

            if died_pid:
                with self.apt_lock:
                    try:
                        proc = self.apt.pop(died_pid)
                        self.puid2pid.pop(proc.props.p_uid)
                    except KeyError:
                        log.warning(f'unknown child pid {died_pid} exited!')
                        proc = None

                if proc is None:
                    continue

                ecode = os.waitstatus_to_exitcode(exit_status)
                log.info(f'p_uid: {proc.props.p_uid} pid: {died_pid} ecode={ecode}')
                resp = dmsg.SHProcessExit(tag=get_new_tag(), exit_code=ecode, p_uid=proc.props.p_uid)

                if proc.props.p_uid != dfacts.GS_PUID:
                    if proc.props.r_c_uid is None:
                        self.gs_in.send(resp.serialize())
                        log.info(f'transmit {repr(resp)} via gs_in')
                    else:
                        r_c_uid = proc.props.r_c_uid
                        self._send_response(target_uid=r_c_uid, msg=resp)
                        log.info(f'transmit {repr(resp)} via _send_response')

                # If we haven't received SHTeardown yet
                if proc.props.critical and not self.check_shutdown():
                    if proc.props.p_uid == dfacts.GS_PUID:
                        # if this is GS and we haven't received GSHalted yet and SHTeardown
                        # has not been received then this is an abnormal termination condition.
                        if self.is_primary and (not self.check_gs_shutdown()) and not self.check_shutdown():
                            # Signal abnormal termination and notify Launcher BE
                            err_msg = f'LS watch death - GS exited - puid {proc.props.p_uid}'
                            self._abnormal_termination(err_msg)
                    elif dfacts.is_transport_puid(proc.props.p_uid):
                        if (not self.check_ta_shutdown()) and (not self.check_shutdown()):
                            err_msg = f'LS watch death - TA exited - puid {proc.props.p_uid}'
                            self._abnormal_termination(err_msg)
                    else:
                        # Signal abnormal termination and notify Launcher BE
                        err_msg = f'LS watch death - critical process exited - puid {proc.props.p_uid}'
                        self._abnormal_termination(err_msg)

                try:  # keep subprocess from giving spurious ResourceWarning
                    proc.wait(0)
                    # Remember to close any open connections for stdout and stderr.
                    # If they weren't opened, the close methods will handle that. The
                    # underlying channel will be decref'ed when the SHProcessExit is
                    # received by global services (see server.py in GS).
                    if proc.props.stdout_connector is not None:
                        self.exited_channel_output_monitors.put(proc.props.stdout_connector)
                    if proc.props.stderr_connector is not None:
                        self.exited_channel_output_monitors.put(proc.props.stderr_connector)
                except OSError:
                    pass
            else:
                time.sleep(self.SHUTDOWN_RESP_TIMEOUT)

        log.info('exit')

    def update_watch_set(self, connectors, dead_connector):
        changed = False
        while not self.new_channel_input_monitors.empty():
            changed = True
            connector = self.new_channel_input_monitors.get()
            connectors.add(connector)

        if dead_connector is not None:
            changed = True
            connectors.discard(dead_connector)

        return changed

    def watch_input(self):
        """Thread monitors inbound traffic directed to stdin of a process."""

        log = logging.getLogger('watch input')

        log.info('starting')

        connectors = set()
        channel_set = None
        dead_connector = None

        # Gets us the default pool.
        node_index = parms.this_process.index
        def_muid = dfacts.default_pool_muid_from_index(node_index)
        def_pool = self.pools[def_muid]

        while not self.check_shutdown():
            EOF = False

            new_channel_set_needed = self.update_watch_set(connectors, dead_connector)

            if len(connectors) == 0:
                time.sleep(self.SHUTDOWN_RESP_TIMEOUT)
                if channel_set is not None:
                    del channel_set
                    channel_set = None
            else:
                try:
                    if new_channel_set_needed:
                        del channel_set
                        channel_list = [connector.conn.inbound_channel for connector in connectors]
                        channel_set = dch.ChannelSet(def_pool, channel_list)

                    dead_connector = None

                    connector = None
                    channel, event = channel_set.poll(self.SHUTDOWN_RESP_TIMEOUT)
                    connector = self.cuid_to_input_connector[channel.cuid]
                    if event == dch.POLLIN:
                        EOF = connector.forward()

                    if EOF or event == dch.POLLNOTHING or not connector.proc_is_alive:
                        dead_connector = connector
                        connector.close()


                except dch.ChannelSetTimeout:
                    pass
                except Exception as ex:
                    # Any error is likely due to the child exiting
                    log.info(f'InputConnector Error:{repr(ex)}')
                    try:
                        if connector is not None:
                            dead_connector = connector
                            connector.close()
                    except:
                        pass

        log.info('exiting')


    def watch_output(self):
        """Thread monitors outbound std* activity from processes we started.

            Any stderr activity on a 'critical' (e.g. infrastructure) process
            running under this shepherd will cause an error shutdown.

        :return: None, exits on self.check_shutdown()
        """
        log = logging.getLogger('watch output')

        log.info('starting')
        p_hostname = self.hostname

        class WatchingSelector(selectors.DefaultSelector):
            """Enhanced DefaultSelector to register stdout/stderr of PopenProps

            Automates registering the file handles with the selector base class
            and maps it to an OutputConnector object to handle the logic for
            forwarding data where it needs to go.
            """

            def add_proc_streams(self, server, proc: PopenProps):
                # carried data is (ProcessProps, closure to make SHFwdOutput, whether stderr or not)
                try:
                    self.register(proc.stdout, selectors.EVENT_READ,
                                  data=proc.props.stdout_connector)
                except ValueError:  # file handle could be closed or None: a race, so must catch
                    pass
                except KeyError as ke:  # already registered
                    self.unregister(proc.stdout)
                    #log.warning(f'ke={ke}, proc.stdout={proc.stdout}, new props={proc.props}')
                    self.register(proc.stdout, selectors.EVENT_READ,
                                  data=proc.props.stdout_connector)

                try:
                    self.register(proc.stderr, selectors.EVENT_READ,
                                  data=proc.props.stderr_connector)
                except ValueError:
                    pass
                except KeyError as ke:
                    self.unregister(proc.stderr)
                    #log.warning(f'ke={ke}, proc.stderr={proc.stderr}, new props={proc.props}')
                    self.register(proc.stderr, selectors.EVENT_READ,
                                  data=proc.props.stderr_connector)

        stream_sel = WatchingSelector()

        while not self.check_shutdown():
            work = stream_sel.select(timeout=self.SHUTDOWN_RESP_TIMEOUT)

            for sel_k, _ in work:
                output_connector = sel_k.data

                try:
                    EOF = output_connector.flush()
                except TerminationException as ex:
                    EOF = False
                    handled = False
                    str_data = str(ex)
                    try:  # did we get a GSHalted?
                        msg = dmsg.parse(str_data)
                        if isinstance(msg, dmsg.GSHalted):
                            handled = True
                            self._DTBL[dmsg.GSHalted][0](self, msg=msg)
                            EOF = True
                    except json.JSONDecodeError:
                        pass

                    if not handled:
                        err_msg = f'output from critical puid {output_connector.puid}'
                        log.error(err_msg)
                        log.error(f'output is:\n{str_data}\n')
                        # Signal abnormal termination and notify Launcher BE
                        self._abnormal_termination(err_msg)

                if EOF:
                    stream_sel.unregister(output_connector.file_obj)
                    output_connector.close()

            if self.check_shutdown():
                break

            try:
                while True:
                    new_proc = self.new_procs.get_nowait()
                    stream_sel.add_proc_streams(self, new_proc)
            except queue.Empty:
                pass
            except ValueError: # The file was closed, possible race condition.
                pass

            try:
                while True:
                    exited_proc_connector = self.exited_channel_output_monitors.get_nowait()
                    stream_sel.unregister(exited_proc_connector.file_obj)
                    exited_proc_connector.close()
            except queue.Empty:
                pass
            except ValueError: # The file was closed, possible race condition.
                pass

        stream_sel.close()
        log.info('exit')

src/dragon/localservices/server.py
DragonHPC/dragon
"""The Dragon MPBridge component maps the Python Multiprocessing API onto Dragon
Native components.
"""

import io
import os
import pickle
import sys

import multiprocessing
import multiprocessing.spawn
from multiprocessing.util import debug
from multiprocessing.spawn import prepare
from multiprocessing.process import current_process

from dragon.infrastructure.parameters import this_process

from .connection import ConnGlue
from .process import DragonProcess, PUID, DragonPopen, DragonPopenStub
from .context import DragonContext
from .queues import DragonQueue, DragonSimpleQueue, DragonJoinableQueue
from .pool import DragonPool
from .synchronize import DragonSemaphore, DragonBoundedSemaphore, DragonLock, DragonRLock, DragonBarrier
from .reduction import dragon_DupFd, dragon_recvfds, dragon_sendfds
from .sharedctypes import DragonArray, DragonRawArray, DragonValue, DragonRawValue


def dragon_spawn_main():
    """Main entrypoint for running a Multiprocessing program backed by Dragon.
    I.e. this is the function first executed when a new Dragon process is
    started through a Multiprocessing call to `process.start()`.
    See `dragon.mpbridge.process.Popen._launch`.
    """
    assert multiprocessing.spawn.is_forking(sys.argv), "Not forking"
    if sys.platform == "win32":
        raise NotImplementedError("win32 unsupported by dragon")
    exitcode = _dragon_main()
    sys.exit(exitcode)


def _dragon_main():
    """This function gets the payload for the new process from the parent,
    constructs the new process object and runs it.
    """

    debug(f"_dragon_main starting os.getpid()={os.getpid()}")

    current_process()._inheriting = True

    try:
        import dragon.globalservices.api_setup as dga

        dga.connect_to_infrastructure()

        # Get argdata
        from_parent = pickle.loads(dga._ARG_PAYLOAD)
        assert isinstance(from_parent, io.BytesIO)
        from_parent.seek(0)

        preparation_data = multiprocessing.context.reduction.pickle.load(from_parent)
        prepare(preparation_data)

        # get process objects
        process_obj = multiprocessing.context.reduction.pickle.load(from_parent)

        # force it to be the right class & install fake _popen stub
        # object to fool some methods inside...
        process_obj._popen = DragonPopenStub(PUID(this_process.my_puid))

    except Exception as e:
        print(
            f"Starting Dragon process {this_process.my_puid} with pid {os.getpid()} failed: \n{e}", flush=True
        )
        raise e
    finally:
        del multiprocessing.process.current_process()._inheriting

    return process_obj._bootstrap(process_obj._p_p_uid)  # run process

src/dragon/mpbridge/__init__.py
DragonHPC/dragon
"""Dragon's replacement for the Multiprocessing Connection and Pipe objects
"""
import os
import pickle

import multiprocessing.connection
import multiprocessing.context
from multiprocessing.util import debug, sub_debug, info, sub_warning

import inspect  # Only needed for demo, not critical functionality
from pprint import pprint  # Only needed for demo, not critical functionality

_original_os_write = os.write
_original_os_read = os.read
_original_os_close = os.close
_original_os_pipe = os.pipe

# This file mostly contains dead code. The real implementation is in infrastructure.connection

class ConnGlueBase:
    def __init__(self, ucid, note, *args, **kwargs):
        self.ucid = ucid
        self.note = note

    def __repr__(self):
        return f'{self.__class__.__name__}({", ".join(f"{k}={v!r}" for k, v in self.__dict__.items())})'

    def get_fd(self):
        # TODO: Eliminate once not using file descriptors anymore and modify
        #       connection.wait() accordingly.
        return int(self.ucid)

    @staticmethod
    def extract_fd(conn_fileno):
        # TODO: Eliminate once not using file descriptors anymore.
        # Nice thing is, to track down all the places where a file descriptor
        # is still really being used, just search for this func's invocation.
        return conn_fileno.get_fd() if hasattr(conn_fileno, "get_fd") else conn_fileno

    def write(self, *args, **kwargs):
        sub_debug(f"{self.__class__.__name__}.write called len(args[0])={len(args[0])} args={args}")
        fd = self.__class__.extract_fd(self)
        return _original_os_write(fd, *args, **kwargs)

    def read(self, *args, **kwargs):
        sub_debug(f"{self.__class__.__name__}.read called args={args}")
        try:
            fd = self.__class__.extract_fd(self)
            retval = _original_os_read(fd, *args, **kwargs)
        except Exception as e:
            sub_warning(f"{self.__class__.__name__}.read exception: {e}")
            raise e
        return retval

    def close(self, *args, **kwargs):
        sub_debug(f"{self.__class__.__name__}.close called with ({args}, {kwargs}) by pid={os.getpid()}")
        # if inspections:
        #    frame = inspect.currentframe()
        #    pprint(f"{inspect.getouterframes(frame)}")
        fd = self.__class__.extract_fd(self)
        return _original_os_close(fd, *args, **kwargs)

    @classmethod
    def pipe(cls, _log_note=""):
        a, b = _original_os_pipe()
        retval = (
            cls(str(a), f"r from {cls.__name__}.pipe {_log_note}"),
            cls(str(b), f"w from {cls.__name__}.pipe {_log_note}"),
        )
        debug(f"{cls.__name__}.pipe called, returning: {retval}")
        return retval

    def fileno(self):
        # TODO: Fix the dependency on this in pool._wait_for_updates() which
        # depends upon connection.wait() (very similar to need for fix to
        # DragonConnection.fileno()) by either monkeypatching connection.wait()
        # (and making sure pool picks it up) or by providing a viable object
        # for the selectors module to use to monitor and wait on.
        info(f"{self.__class__.__name__}.fileno() still being called; mind the TODO here")
        return self.get_fd()


try:
    from dragon.infrastructure.connection import CUID, Connection, Pipe as DragonPipe
    from .. import managed_memory
    from .. import channels
    from .process import PUID
except ImportError as e:
    raise ImportError(f"import dragon failed: {e}")
else:
    ##################################################
    # Use dragon.infrastructure.connection.Connection
    ##################################################

    class ConnGlue(ConnGlueBase):
        "Glue around dragon.infrastructure.connection.Connection."
        _mem_pool = None
        _last_channel_uid = None

        def __init__(self, ucid, note, lowerlevel_conn, *args, **kwargs):
            # TODO: Obtain ucid from lowerlevel_conn instead of requiring separate input?
            # TODO: Enforce uniqueness or trust this in lowerlevel_conn objects and leave well enough alone?
            self.ucid = ucid
            self.note = note
            if isinstance(lowerlevel_conn, bytes):
                self.lowerlevel_conn = pickle.loads(lowerlevel_conn)
            else:
                self.lowerlevel_conn = lowerlevel_conn  # MinConnection
            self.read_buf = b""  # TODO: Verify either this is needed or eliminate it

        def __repr__(self):
            try:
                pickled_lowerlevel_conn = pickle.dumps(self.lowerlevel_conn)
            except (ImportError, AttributeError):
                pickled_lowerlevel_conn = "unavailable during shutdown"
            return (
                f"{self.__class__.__name__}(ucid={self.ucid!r}, note={self.note!r}, lowerlevel_c"
                f"onn={pickled_lowerlevel_conn!r})"
            )

        def write(self, data, *args, **kwargs):
            sub_debug(
                f"{self.__class__.__name__}.write called len(data)={len(data)} args={args} kwargs={kwargs}"
            )
            try:
                self.lowerlevel_conn.send_bytes(data)
            except Exception as e:
                sub_warning(f"{self.__class__.__name__}.write exception: {e}")
                raise
            return len(data)

        def read(self, n_bytes=None, *args, **kwargs):
            # TODO: Should this be "read up to n_bytes" instead of "read until n_bytes avail"?
            # Buffer receiving from the Connection object.
            sub_debug(f"{self.__class__.__name__}.read called n_bytes={n_bytes} args={args} kwargs={kwargs}")
            _read_buf = self.read_buf

            limit = n_bytes if (n_bytes is not None) else (len(_read_buf) + 1)
            while len(_read_buf) < limit:
                try:
                    temp = self.lowerlevel_conn.recv_bytes()
                    _read_buf += temp
                except EOFError:
                    return b""  # Indicate EOF with an empty.
                except Exception as e:
                    sub_warning(f"{self.__class__.__name__}.read exception: e={e}")
                    raise

            cutpoint = n_bytes if (n_bytes is not None) else len(_read_buf)
            retval, self.read_buf = _read_buf[:cutpoint], _read_buf[cutpoint:]

            return retval

        def close(self, *args, **kwargs):
            sub_debug(
                f"{self.__class__.__name__}.close called with ({args}, {kwargs}) by os.getpid()={os.getpid()}"
            )
            self.lowerlevel_conn.close()
            super().close(*args, **kwargs)  # TODO: Remove once no longer employing file descriptors at all.


inspections = False


class DragonConnection(multiprocessing.connection.Connection):
    _write = ConnGlue.write
    _read = ConnGlue.read

    def __init__(self, handle: ConnGlue, readable=True, writable=True):
        debug(
            f"{self.__class__.__name__}.__init__(handle={handle}, readable={readable}, writable={writable}) was "
            f"called!"
        )
        if not readable and not writable:
            raise ValueError("at least one of `readable` and `writable` must be True")
        # TODO: Extra processing of handle as appropriate for dragon.
        # self._handle = handle.__index__()
        # assert self._handle >= 0
        self._handle = handle
        self._readable = readable
        self._writable = writable

    def _close(self, _close=ConnGlue.close):
        debug(f"{self.__class__.__name__}.close() was called on {self._handle}!")
        try:
            _close(self._handle)
        except OSError:
            # TEMPORARY to see if we can make things work for now with Popen.DupFd
            pass

    def _send(self, buf, write=_write):
        "TODO: Wrapper to pick up dragon-specific _write."
        return super()._send(buf, write)

    def _recv(self, size, read=_read):
        "TODO: Wrapper to pick up dragon-specific _read."
        return super()._recv(size, read)

    def __repr__(self):
        return (
            f"{self.__class__.__name__}({self._handle}, readable={self._readable}, writable={self._writable})"
        )

    def fileno(self):
        # TODO: Eliminate once file descriptors not used and fix connection.wait()
        # Leveraged by multiprocessing.connection.wait() which in turn
        # indirectly calls selectors._fileobj_to_fd().
        info("DragonConnection.fileno() still being called; mind the TODO here")
        retval = super().fileno()
        fd = ConnGlue.extract_fd(retval)
        return fd


def Pipe(duplex=True, *, original=None, use_base_impl=True):
    """Generate two Connection objects to be used by only one process each.

    :param duplex: If True, both ends of the Pipe may send and receive. Optional, defaults to True.
    :type duplex: bool
    :return: tuple of connection objects
    """
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(duplex)
    else:
        return DragonPipe(duplex)


class DragonListener(multiprocessing.connection.Listener):
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonListener is not implemented, yet")
        super().__init__(*args, **kwargs)


class BaseImplListener(multiprocessing.connection.Listener):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


def DragonClient(address, family=None, authkey=None):
    raise NotImplementedError(f"DragonClient is not implemented, yet")


def Listener(address=None, family=None, backlog=1, authkey=None, *, ctx=None, use_base_impl=True):
    if use_base_impl:
        return BaseImplListener(address=address, family=family, backlog=backlog, authkey=authkey)
    else:
        return DragonListener(address=address, family=family, backlog=backlog, authkey=authkey)


def Client(address, family=None, authkey=None, *, original=None, use_base_impl=True):
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(address, family=family, authkey=authkey)
    else:
        return DragonClient(address, family=family, authkey=authkey)


def reduce_connection(conn):
    debug(f"reduce_connection(conn={conn}) invoked by os.getpid()={os.getpid()}")
    # Switching from reduction.DupFd to get_spawning_popen().DupFd has the
    # consequence that the file descriptor will not be appended to the list in
    # the Popen instance's _fds attribute.  Only file descriptors in that _fds
    # list will be passed to util.spawnv_passfds() for marking as inheritable
    # and closing prior to vfork (spawning the child process).
    df = multiprocessing.context.reduction.DupFd(conn.fileno())
    # TODO: Replace the above line with the following one when
    #       DragonConnection objects no longer employ true file descriptors.
    # df = multiprocessing.context.get_spawning_popen().DupFd(conn.fileno())
    return rebuild_connection, (df, conn.readable, conn.writable, conn._handle)


def rebuild_connection(df, readable, writable, handle):
    fd = df.detach()  # TODO: adjust for DragonConnection having no file descriptor
    debug(
        f'rebuild_connection({df}->{getattr(fd, "__index__", int)()}, readable={readable}, writable={writable}) invo'
        f"ked by os.getpid()={os.getpid()}"
    )
    if not isinstance(multiprocessing.connection.os, PseudoOS):
        # Ensure child processes have this activated to aid tracking/debugging.
        activate_monkeypatch_connection_os()
    # Original returned Connection(), necessitating monkeypatching of
    # Connection until this function was reimplemented for DragonConnection.
    return DragonConnection(handle, readable, writable)


multiprocessing.connection.reduction.register(DragonConnection, reduce_connection)




class PseudoOS:
    __os = os

    def __getattr__(self, name):
        info(f"{self.__class__.__name__}.{name} requested")
        return getattr(self.__os, name)

    def write(self, *args, **kwargs):
        sub_warning(f"{self.__class__.__name__}.write called")
        return self.__os.write(*args, **kwargs)

    def read(self, *args, **kwargs):
        sub_warning(f"{self.__class__.__name__}.read called")
        return self.__os.read(*args, **kwargs)

    def close(self, *args, **kwargs):
        sub_warning(f"{self.__class__.__name__}.close called with ({args}, {kwargs})")
        return self.__os.close(*args, **kwargs)

    def pipe(self):
        retval = self.__os.pipe()
        sub_warning(f"{self.__class__.__name__}.pipe called, returning: {retval}")
        if inspections:
            frame = inspect.currentframe()
            pprint(f"{inspect.getouterframes(frame)}")
        return retval


def activate_monkeypatch_connection_os():
    multiprocessing.connection.os = PseudoOS()

src/dragon/mpbridge/connection.py
DragonHPC/dragon
"""Dragon's replacement for the Multiprocessing context object.

Also contains a high performance implementation of `wait()` on Multiprocessing and Dragon
Native objects. A blocking wait on one or more objects with a timeout is a
common task when handling communication. Objects can be of very different nature
(Queue, Process, Pipe), below them is usually a Dragon Channel.  Our
implementation here handles multiple waiters on lists that contain some of the
same objects gracefully and in a high perfomant way. Objects are categorized by
type and threads are spawned for every type. Objects of the same type are
grouped into a multi-join call, so the number of threads spawned is minimized.
The life-cycle of the process handling the wait is independent of the waiter
process, i.e. multiple repeated calls on the same objects by the same or
different processes will create only minimal overhead.

:raises NotImplementedError: If used with win32
"""


import multiprocessing
import multiprocessing.pool
import multiprocessing.queues
import multiprocessing.connection
import multiprocessing.shared_memory
import multiprocessing.synchronize
import multiprocessing.sharedctypes

import threading
import time
import os

import dragon
from .monkeypatching import original_multiprocessing
from .queues import DragonJoinableQueue, DragonQueue
from .shared_memory import DragonSharedMemory, DragonShareableList
from .process import DragonProcess, DragonPopen, PUID
from .connection import DragonListener, DragonClient

from dragon.globalservices.process import multi_join

import dragon.infrastructure.connection


def _sort_and_check_for_ready(object_list):
    """Sort objects by waiting type and check, if any are ready already.
    Addresses Multiprocessing unit test TestWait.test_wait_integer

    :param object_list: object to wait on
    :type object_list: list of PUID, Connection, etc.
    :return: ready objects, puids, connections and others
    :rtype: tuple of lists
    """

    # sort by type
    puids = set()
    conns = set()
    queues = set()
    others = list()

    for obj in object_list:
        if isinstance(obj, PUID):
            puids.add(obj)
        elif isinstance(obj, dragon.infrastructure.connection.Connection):
            conns.add(obj)
        elif isinstance(obj, dragon.infrastructure.connection.CUID):
            conns.add(obj)
        elif isinstance(obj, dragon.mpbridge.queues.DragonQueue):
            queues.add(obj)
        elif isinstance(obj, dragon.mpbridge.queues.DragonJoinableQueue):
            queues.add(obj)
        elif isinstance(obj, dragon.mpbridge.queues.DragonSimpleQueue):
            queues.add(obj)
        elif isinstance(obj, dragon.mpbridge.queues.BaseImplQueue):
            queues.add(obj)
        elif isinstance(obj, dragon.mpbridge.queues.BaseImplJoinableQueue):
            queues.add(obj)
        elif isinstance(obj, dragon.mpbridge.queues.BaseImplSimpleQueue):
            queues.add(obj)
        elif isinstance(obj, dragon.mpbridge.queues.FakeConnectionHandle):
            conns.add(obj)
        else:
            others.append(obj)

    # Check if any are done already
    # I have decided not to use the methods of _WaitService here,
    # but to call into Dragon directly. It reduces complexity, we might not have to spawn the
    # _WaitService class at all and it separates the concern of this function with _WaitService.

    ready_puids = []
    ready_conns = []
    ready_queues = []
    ready_others = []

    if puids:
        object_tuples, proc_status = multi_join(list(puids), 0)
        if object_tuples:
            ready_puids = [PUID(i[0]) for i in object_tuples]

    if conns:
        ready_conns = [conn for conn in list(conns) if conn.poll(0)]

    if queues:
        ready_queues = [q for q in list(queues) if q._poll(0)]

    if others:
        ready_others = original_multiprocessing.connection_wait(others, 0)

    ready_objects = ready_puids + ready_conns + ready_queues + ready_others

    return ready_objects, list(puids), list(conns), list(queues), list(others)


class _WaitService:
    """Provide a thread safe way of waiting on multiple types of objects at
    once.  This is a performance oriented implementation with shared data
    structures and one lock.  Service threads are waiting on the same event as
    their parent, that joins on them.  We never spawn more than one service
    thread per object. We handle Dragon PUIDs, connection, queues and the
    standard multiprocessing objects. If only a single new object is waited on,
    we do not spawn a service thread, but use the incoming thread.

    NOTE: Due to timing/scheduling of threads in the OS, this implementation
    cannot behave exactly like a single threaded version. In particular, when
    waiting on several object types at once, one type will always return first.
    A single threaded version might return a puid and 2 connections together.
    This implementation will return either a puid or 2 connections, but not both
    of them, even if they get ready at virtually the same time.
    `_sort_and_check_for_ready` alleviates the problem for objects that are
    already 'ready', when waited upon.
    """

    THREAD_WAIT = 0.1  # [sec] time interval for a thread to check if it needs to stop waiting, used when timeout is None
    GC_INTERVAL = 10  # [sec] how often to look for abandoned waiters/ready_objects from vanished waiters
    # exact value not important as it's not related to timeout

    def __init__(self, timeout):

        self.wait_timeout = timeout # [sec] time interval for a thread to check if it needs to stop waiting

        self.insert_lock = threading.Lock()  # protect inserting new objects and waiters

        self.events = {}  # event of waiting tid

        self.objects = {}  # list of objects by service thread tid
        self.waiters = {}  # list of waiting tids by service thread tid

        self.ready_objects = {}  # list of ready objects by waiter tid

        self.garbage_truck = threading.Thread(target=self._collect_garbage, daemon=True)
        self.garbage_truck.start()

    def __del__(self):
        """Shut down service threads"""

        for ev in self.events.values():
            ev.set()  # stop every waiter and the garbage truck

    def _collect_garbage(self):
        """Cleanup leftover tids that have never been picked up by a parent."""

        my_pid = os.getpid()
        my_tid = threading.get_native_id()

        with self.insert_lock:
            self.events[my_tid] = threading.Event()

        while True:

            # Wait on the event for 10sec.
            # If the event is set, this breaks immediately
            res = self.events[my_tid].wait(timeout=self.GC_INTERVAL)
            if res:
                break

            self.insert_lock.acquire()
            active_tids = os.listdir(f"/proc/{my_pid}/task")  # Linux only
            waiting_tids = {item for sublist in self.waiters.values() for item in sublist}
            ready_tids = self.ready_objects.keys()

            # find waiter tids that are no longer active, i.e. have disappeared
            bad_tids = []
            for tid in waiting_tids:
                if str(tid) not in active_tids:
                    bad_tids.append(tid)

            # remove bad waiter tids from the data structures. This is a bit tedious.
            for tid in bad_tids:
                for serv_tid in self.waiters.keys():
                    if tid in self.waiters[serv_tid]:
                        self.waiters[serv_tid].remove(tid)
                if tid in ready_tids:
                    del self.ready_objects[tid]
            self.insert_lock.release()

    def _wait_on_connection(self, conn):
        """Paranoid wait for a single connection object. We wake up from time to time
        to check if we are actually still needed.

        :param conns: connection objects to wait on
        :type conns: list of
                     dragon.infrastructure.connection.Connection,
                     dragon.infrastructure.connection.CUID,
                     dragon.mpbridge.queues.DragonQueue,
                     dragon.mpbridge.queues.Faker
        """

        my_tid = threading.get_native_id()

        joined = False
        lonely = False

        while not lonely and not joined:
            joined = conn.poll(self.wait_timeout)

            self.insert_lock.acquire()
            lonely = self.waiters[my_tid] == []
            self.insert_lock.release()

        if joined and not lonely:
            self._handle_wait_return([conn])

        self.insert_lock.acquire()
        del self.objects[my_tid]
        del self.waiters[my_tid]
        self.insert_lock.release()

    def _wait_on_queue(self, queue):
        """Paranoid wait for a single queue object. We wake up from time to time
        to check if we are actually still needed.

        :param queue: queue object to wait on
        :type queue: one of
                     dragon.mpbridge.queues.DragonQueue,
                     dragon.mpbridge.queues.DragonSimpleQueue,
                     dragon.mpbridge.queues.DragonJoinableQueue,
        """

        my_tid = threading.get_native_id()

        joined = False
        lonely = False

        while not lonely and not joined:
            joined = queue._poll(self.wait_timeout)

            self.insert_lock.acquire()
            lonely = self.waiters[my_tid] == []
            self.insert_lock.release()

        if joined and not lonely:
            self._handle_wait_return([queue])

        self.insert_lock.acquire()
        del self.objects[my_tid]
        del self.waiters[my_tid]
        self.insert_lock.release()

    def _wait_on_puids(self, puids):
        """Paranoid wait for a list PUIDs, pids, sentinels using a single thread and Dragon's multi-join.
        We wake up from time to time to check if our parent disappeared, i.e. if we are lonely.

        :param puids: PUID objects to wait on
        :type puids: list of dragon.process.PUID
        """

        my_tid = threading.get_native_id()
        object_tuples = None
        joined = False
        lonely = False

        while not lonely and not joined:
            object_tuples, proc_status = multi_join(puids, self.wait_timeout)

            joined = object_tuples != None
            self.insert_lock.acquire()
            lonely = self.waiters[my_tid] == []
            self.insert_lock.release()

        if joined and not lonely:
            ready_objects = [i[0] for i in object_tuples]
            self._handle_wait_return(ready_objects)

        self.insert_lock.acquire()
        del self.objects[my_tid]
        del self.waiters[my_tid]
        self.insert_lock.release()

    def _wait_on_multiprocessing_objects(self, objs):
        """Paranoid wait for a list of standard multiprocessing objects.
        We wake up from time to time to check if someone else got ready.

        :param objs: mp objects to wait on
        :type objs: list of
                    multiprocessing.Connection,
                    socket.socket,
                    multiprocessing.Process.sentinel
        """

        my_tid = threading.get_native_id()
        joined = False
        lonely = False

        while not lonely and not joined:
            # saved this from monkey patching
            ready_objects = original_multiprocessing.connection_wait(objs, self.wait_timeout)

            joined = ready_objects != []
            self.insert_lock.acquire()
            lonely = self.waiters[my_tid] == []
            self.insert_lock.release()

        if joined and not lonely:
            self._handle_wait_return(ready_objects)

        self.insert_lock.acquire()
        del self.objects[my_tid]
        del self.waiters[my_tid]
        self.insert_lock.release()

    def _handle_wait_return(self, ready_objects):
        """Handle the ready objects returned by a wait call.

        :param ready_objs: list of objects that are "ready"
        :type ready_objs: list connections, puids, multiprocessing objects
        """

        my_tid = threading.get_native_id()
        my_pid = os.getpid()

        self.insert_lock.acquire()

        my_waiters = self.waiters[my_tid]

        # add their objects for pickup and notify
        for tid in my_waiters:
            self.ready_objects[tid] = ready_objects
            self.events[tid].set()

        self.insert_lock.release()

    def wait_for_me(self, puids, conns, queues, others):
        """Initiate a wait on puids, connections and multiprocessing objects.
        We spawn service threads for:
            * every list of new multiprocessing or CUID objects
            * every single new connection object
        We do not spawn service threads, if:
            * an object is already waited on

        :param puids: puid object handles to wait on
        :type puids: list
        :param conns: _description_
        :type conns: _type_
        :param queues: _description_
        :type queues: _type_
        :param others: _description_
        :type others: _type_
        :return: _description_
        :rtype: _type_
        """

        my_tid = threading.get_native_id()
        done_ev = threading.Event()

        # spin until we get the lock and can accept more waits
        self.insert_lock.acquire()

        self.events[my_tid] = done_ev

        # add myself as a waiter to existing threads

        my_objects = set(puids + conns + queues + others)

        for tid, objs in self.objects.items():
            if my_objects.intersection(set(objs)):
                self.waiters[tid].append(my_tid)

        # spawn threads for new objects not already waited on
        cur_objects = [obj for lst in self.objects.values() for obj in lst]  # flatten dict of lists

        new_puids = list(set(puids) - set(cur_objects))
        new_conns = list(set(conns) - set(cur_objects))
        new_queues = list(set(queues) - set(cur_objects))
        new_others = list(set(others) - set(cur_objects))

        if new_puids:
            t = threading.Thread(target=self._wait_on_puids, args=(new_puids,), daemon=True)
            t.start()
            self.waiters[t.native_id] = [my_tid]
            self.objects[t.native_id] = new_puids

        # TODO: multi_join for everyone, will break thread hoarding test

        if new_conns:
            for obj in new_conns:  # PE-43870 need to replace with multi-join on channels
                t = threading.Thread(target=self._wait_on_connection, args=(obj,), daemon=True)
                t.start()
                self.waiters[t.native_id] = [my_tid]
                self.objects[t.native_id] = [obj]

        if new_queues:
            for obj in new_queues:  # PE-43870 need to replace with multi-join on channels
                t = threading.Thread(target=self._wait_on_queue, args=(obj,), daemon=True)
                t.start()
                self.waiters[t.native_id] = [my_tid]
                self.objects[t.native_id] = [obj]

        if new_others:
            t = threading.Thread(
                target=self._wait_on_multiprocessing_objects, args=(new_others,), daemon=True
            )
            t.start()
            self.waiters[t.native_id] = [my_tid]
            self.objects[t.native_id] = new_others

        self.insert_lock.release()

        return done_ev

    def pickup_and_leave(self, my_objects, pickup):
        """Get the objects ready after the wait for this process.
        Clean internal data structures, even if we timed out.

        :param my_objects: list of objects that were waited on
        :param pickup: bool, that is false we only timed out and pickup ready objects.
        :return: list of objects ready after wait
        """

        my_tid = threading.get_native_id()
        my_ready_objects = []

        self.insert_lock.acquire()

        del self.events[my_tid]

        # remove myself from the waiters dict
        for tid, waiters in self.waiters.items():
            if my_tid in waiters:
                self.waiters[tid].remove(my_tid)

        if pickup:
            all_ready_objects = set(self.ready_objects.pop(my_tid))
            my_ready_objects = list(all_ready_objects.intersection(my_objects))

        self.insert_lock.release()

        return my_ready_objects


_wait_service = None

# these global variables store the original multiprocessing
# functions before monkeypatching

_original_cpu_count = None
_original_active_children = None
_original_current_process = None
_original_parent_process = None
_original_get_logger = None
_original_log_to_stderr = None

_original_Pipe = None
_original_Queue = None
_original_JoinableQueue = None
_original_Value = None
_original_RawValue = None
_original_Array = None
_original_RawArray = None
_original_Client = None
_original_SharedMemory = None
_original_ShareableList = None


class DragonContext(multiprocessing.context.BaseContext):
    """The Dragon context patches the Dragon Native API into
    Python Multiprocessing using the dragon.mpbridge modules,
    when the startmethod is set to 'dragon' and Dragon is imported.
    This is done in 2 ways:

    1. The Dragon context is selected and with it all of
    the replaced API (see dragon/__init__.py).

    2. The module is monkey patched during import, to ensure
    users that don't use the context to get parts of the API
    still get the dragon/mpbridge/*.py methods and classes.
    """

    Process = DragonProcess

    _name = DragonPopen.method

    # these allow us to enable parts of the API original multiprocessing API for debugging
    USE_MPFUNCTIONS = False  # switches all functions under context
    USE_MPQUEUE = False
    USE_MPJQUEUE = False
    USE_MPSQUEUE = False
    USE_MPPOOL = False
    USE_MPPIPE = False
    USE_MPMANAGER = True
    USE_MPLOCK = False
    USE_MPRLOCK = False
    USE_MPCONDITION = False
    USE_MPSEMAPHORE = False
    USE_MPBOUNDSEMAPHORE = False
    USE_MPEVENT = False
    USE_MPBARRIER = False
    USE_MPVALUE = False
    USE_MPRAWVALUE = False
    USE_MPARRAY = False
    USE_MPRAWARRAY = False
    USE_MPSHAREDMEMORY = False
    USE_MPSHAREABLELIST = False
    USE_MPLISTENER = False
    USE_MPCLIENT = False

    def cpu_count(self):
        """Returns the number of CPUs in the system"""
        if multiprocessing.get_start_method() == self._name:
            from ..native.machine import cpu_count

            if self.USE_MPFUNCTIONS:
                return _original_cpu_count()
            else:
                return cpu_count()
        else:
            return multiprocessing.get_context().cpu_count()

    def freeze_support(self):
        """Check whether this is a fake forked process in a frozen executable.
        If so then run code specified by commandline and exit.
        """
        if multiprocessing.get_start_method() == self._name:
            raise NotImplementedError(
                f"Multiprocessing.freeze_support() is not implemented in Dragon, as it is Windows only."
            )
        else:
            return multiprocessing.get_context().freeze_support()

    def log_to_stderr(self, level=None):
        """Turn on logging and add a handler which prints to stderr"""

        if multiprocessing.get_start_method() == self._name:
            from .util import log_to_stderr

            if self.USE_MPFUNCTIONS:
                return log_to_stderr(level=level, original=_original_log_to_stderr, use_base_impl=True)
            else:
                return log_to_stderr(level=level, original=_original_log_to_stderr, use_base_impl=False)
        else:
            return multiprocessing.get_context().log_to_stderr(level=level)

    def get_logger(self):
        """Return package logger -- if it does not already exist then
        it is created.
        """
        if multiprocessing.get_start_method() == self._name:
            from .util import get_logger

            if self.USE_MPFUNCTIONS:
                return get_logger(original=_original_get_logger, use_base_impl=True)
            else:
                return get_logger(original=_original_get_logger, use_base_impl=False)
        else:
            return multiprocessing.get_context().get_logger()

    @staticmethod
    def wait(object_list, timeout=None):
        """Implement a wait on a list of various Dragon and Multiprocessing objects.

        We aim to support the following types:
            * dragon.mpbridge.process.PUID,
            * dragon.infrastructure.connection.Connection,
            * dragon.infrastructure.connection.CUID, (broken)
            * dragon.mpbridge.queue.DragonQueue,
            * dragon.mpbridge.queue.Faker, (broken)
            * standard multiprocessing objects: sockets, mp.connection, mp.sentinel

        :param object_list: list of objects to wait on
        :type object_list: supported types
        :param timeout: time to wait before return, defaults to None
        :type timeout: float, optional
        :return: returns a list of objects that are "ready"
        :rtype: list of supported types
        """

        my_tid = threading.get_native_id()

        start = time.monotonic()

        if timeout == None:
            timeout = 1000000

        # check if we're already done
        ready_objects, puids, conns, queues, others = _sort_and_check_for_ready(object_list)
        if ready_objects or timeout <= 0:
            return ready_objects  # that was easy

        # use threading
        global _wait_service
        if _wait_service is None:
            _wait_service = _WaitService(timeout)

        _wait_service.insert_lock.acquire()
        if timeout == 1000000:
            _wait_service.wait_timeout = _wait_service.THREAD_WAIT # we need a value for the polling timeout
        else:
            # Multiply by 4 because he have previously divided by 4 the value of _wait_service.wait_timeout.
            # Setting the value of _wait_service.wait_timeout to a quarter of the minimum value among the
            # the timeouts we've seen so far is kind of random, but works good enough.
            # See PE-44688 for more info.
            _wait_service.wait_timeout = min(_wait_service.wait_timeout * 4, timeout) / 4
        _wait_service.insert_lock.release()

        done_ev = _wait_service.wait_for_me(puids, conns, queues, others)

        end = time.monotonic()

        if timeout:
            exact_timeout = max(0, timeout - (end - start))
        else:
            exact_timeout = None

        # don't wait longer than we have to
        pickup = done_ev.wait(timeout=exact_timeout)

        ready_objects = _wait_service.pickup_and_leave(object_list, pickup)

        return ready_objects

    def Manager(self):
        if multiprocessing.get_start_method() == self._name:
            from .managers import Manager

            m = Manager(ctx=self.get_context(), use_base_impl=self.USE_MPMANAGER)
        else:
            m = multiprocessing.get_context().Manager()

        m.start()

        return m

    def Pipe(self, duplex=True):
        if multiprocessing.get_start_method() == self._name:
            from .connection import Pipe

            return Pipe(duplex, original=_original_Pipe, use_base_impl=self.USE_MPPIPE)
        else:
            return _original_Pipe(duplex)

    def Lock(self):
        """Returns a non-recursive lock object"""
        if multiprocessing.get_start_method() == self._name:
            from .synchronize import Lock

            return Lock(ctx=self.get_context(), use_base_impl=self.USE_MPLOCK)
        else:
            return multiprocessing.get_context().Lock(ctx=self.get_context())

    def RLock(self):
        """Returns a recursive lock object"""
        if multiprocessing.get_start_method() == self._name:
            from .synchronize import RLock

            return RLock(ctx=self.get_context(), use_base_impl=self.USE_MPRLOCK)
        else:
            return multiprocessing.get_context().RLock(ctx=self.get_context())

    def Condition(self, lock=None):
        """Returns a condition object"""
        if multiprocessing.get_start_method() == self._name:
            from .synchronize import Condition

            return Condition(lock, ctx=self.get_context(), use_base_impl=self.USE_MPCONDITION)
        else:
            return multiprocessing.get_context().Condition(lock, ctx=self.get_context())

    def Semaphore(self, value=1):
        """Returns a semaphore object"""
        if multiprocessing.get_start_method() == self._name:
            from .synchronize import Semaphore

            return Semaphore(value, ctx=self.get_context(), use_base_impl=self.USE_MPSEMAPHORE)
        else:
            return multiprocessing.get_context().Semaphore(value, ctx=self.get_context())

    def BoundedSemaphore(self, value=1):
        """Returns a semaphore object"""
        if multiprocessing.get_start_method() == self._name:
            from .synchronize import BoundedSemaphore

            return BoundedSemaphore(value, ctx=self.get_context(), use_base_impl=self.USE_MPBOUNDSEMAPHORE)
        else:
            return multiprocessing.get_context().BoundedSemaphore(value, ctx=self.get_context())

    def Event(self):
        """Returns an event object"""
        if multiprocessing.get_start_method() == self._name:
            from .synchronize import Event

            return Event(ctx=self.get_context(), use_base_impl=self.USE_MPEVENT)
        else:
            return multiprocessing.get_context().Event(ctx=self.get_context())

    def Barrier(self, parties, action=None, timeout=None):
        """Returns a barrier object"""
        if multiprocessing.get_start_method() == self._name:
            from .synchronize import Barrier

            return Barrier(parties, action, timeout, ctx=self.get_context(), use_base_impl=self.USE_MPBARRIER)
        else:
            return multiprocessing.get_context().Barrier(
                parties, action=action, timeout=timeout, ctx=self.get_context()
            )

    def Queue(self, maxsize=0):
        if multiprocessing.get_start_method() == self._name:
            from .queues import Queue

            if self.USE_MPQUEUE:
                q = Queue(maxsize=maxsize, ctx=self.get_context(), use_base_impl=True)
                # patches internal dragon.infrastructure.connection.Connection
                # to not send EOT when it is closed.
                # q._writer.ghost = True
                return q
            else:
                return Queue(maxsize=maxsize, ctx=self.get_context(), use_base_impl=False)
        else:
            return multiprocessing.get_context().Queue(maxsize)

    def JoinableQueue(self, maxsize=0):
        if multiprocessing.get_start_method() == self._name:
            from .queues import JoinableQueue

            if self.USE_MPJQUEUE:
                q = JoinableQueue(maxsize=maxsize, ctx=self.get_context(), use_base_impl=True)
                # patches internal dragon.infrastructure.connection.Connection
                # to not send EOT when it is closed.
                # q._writer.ghost = True
                return q
            else:
                return JoinableQueue(maxsize=maxsize, ctx=self.get_context(), use_base_impl=False)
        else:
            return multiprocessing.get_context().JoinableQueue(maxsize)

    def SimpleQueue(self):
        if multiprocessing.get_start_method() == self._name:
            from .queues import SimpleQueue, Queue

            if self.USE_MPSQUEUE:
                q = SimpleQueue(ctx=self.get_context(), use_base_impl=True)
                # patches internal dragon.infrastructure.connection.Connection
                # to not send EOT when it is closed.
                # q._writer.ghost = True
                return q
            else:
                # TODO, PE-41041 we need a Dragon-native SimpleQueue (wrapped Pipe), but for now we can try
                # return DragonSimpleQueue(ctx=self.get_context(), use_base_impl=False)
                return SimpleQueue(ctx=self.get_context(), use_base_impl=False)
        else:
            return multiprocessing.get_context().SimpleQueue()

    def Pool(self, processes=None, initializer=None, initargs=(), maxtasksperchild=None):
        """Returns a process pool object"""
        if multiprocessing.get_start_method() == self._name:
            from .pool import Pool

            return Pool(
                processes,
                initializer,
                initargs,
                maxtasksperchild,
                context=self.get_context(),
                use_base_impl=self.USE_MPPOOL,
            )
        else:
            return multiprocessing.get_context().Pool(
                processes=processes,
                initializer=initializer,
                initargs=initargs,
                maxtasksperchild=maxtasksperchild,
                context=self.get_context(),
            )

    def RawValue(self, typecode_or_type, *args):
        """Returns a shared object"""
        if multiprocessing.get_start_method() == self._name:
            from .sharedctypes import RawValue

            return RawValue(
                typecode_or_type, *args, original=_original_RawValue, use_base_impl=self.USE_MPVALUE
            )
        else:
            return _original_RawValue(typecode_or_type, *args)

    def RawArray(self, typecode_or_type, size_or_initializer):
        """Returns a shared array"""
        if multiprocessing.get_start_method() == self._name:
            from .sharedctypes import RawArray

            return RawArray(
                typecode_or_type,
                size_or_initializer,
                original=_original_RawArray,
                use_base_impl=self.USE_MPARRAY,
            )
        else:
            return _original_RawArray(typecode_or_type, size_or_initializer)

    def Value(self, typecode_or_type, *args, lock=True, ctx=None):
        """Returns a synchronized shared object"""
        if multiprocessing.get_start_method() == self._name:
            from .sharedctypes import Value

            return Value(
                typecode_or_type,
                *args,
                lock=lock,
                ctx=ctx,
                original=_original_Value,
                use_base_impl=self.USE_MPRAWVALUE,
            )
        else:
            return _original_Value(typecode_or_type, *args, lock=lock, ctx=ctx)

    def Array(self, typecode_or_type, size_or_initializer, *, lock=True, ctx=None):
        """Returns a synchronized shared array"""
        if multiprocessing.get_start_method() == self._name:
            from .sharedctypes import Array

            return Array(
                typecode_or_type,
                size_or_initializer,
                lock=lock,
                ctx=ctx,
                original=_original_Array,
                use_base_impl=self.USE_MPRAWARRAY,
            )

        else:
            return _original_Array(typecode_or_type, size_or_initializer, lock=lock, ctx=ctx)

    def Listener(self, address=None, family=None, backlog=1, authkey=None):
        if multiprocessing.get_start_method() == self._name:
            from .connection import Listener

            return Listener(
                address=address,
                family=family,
                backlog=backlog,
                authkey=authkey,
                use_base_impl=self.USE_MPLISTENER,
            )

        else:
            return multiprocessing.get_context().Listener(
                address=address, family=family, backlog=backlog, authkey=authkey
            )

    def Client(self, address, family=None, authkey=None):
        if multiprocessing.get_start_method() == self._name:
            from .connection import Client

            return Client(
                address,
                family=family,
                authkey=authkey,
                original=_original_Client,
                use_base_impl=self.USE_MPCLIENT,
            )
        else:
            return _original_Client(address, family=family, authkey=authkey)

    def _apply_patches(self):
        """Additional monkeypatching.
        Unfortunately multiprocessing doesn't use the context for the whole API internally, so
        we have to force it to use our Dragon versions.
        Here we overwrite some of the internal multiprocessing classes explicitly with our own.
        Note that the corresponding files have to be explicitly imported at the top for this to work.
        """

        # save original multiprocessing functions
        # (and classes if their not in context)
        global _original_cpu_count
        global _original_active_children
        global _original_current_process
        global _original_parent_process
        global _original_get_logger
        global _original_log_to_stderr

        global _original_Pipe
        global _original_Value
        global _original_RawValue
        global _original_Array
        global _original_RawArray
        global _original_Client
        global _original_SharedMemory
        global _original_ShareableList
        global _original_Queue
        global _original_JoinableQueue

        if _original_cpu_count is None:
            _original_cpu_count = multiprocessing.cpu_count

        if _original_active_children is None:
            _original_active_children = multiprocessing.active_children

        if _original_current_process is None:
            _original_current_process = multiprocessing.current_process

        if _original_parent_process is None:
            _original_parent_process = multiprocessing.parent_process

        if _original_get_logger is None:
            _original_get_logger = multiprocessing.get_logger

        if _original_log_to_stderr is None:
            _original_log_to_stderr = multiprocessing.log_to_stderr

        if _original_Queue is None:
            _original_Queue = multiprocessing.queues.Queue

        if _original_JoinableQueue is None:
            _original_JoinableQueue = multiprocessing.queues.JoinableQueue

        if _original_Pipe is None:
            _original_Pipe = multiprocessing.connection.Pipe

        if _original_RawArray is None:
            _original_RawArray = multiprocessing.sharedctypes.RawArray

        if _original_Array is None:
            _original_Array = multiprocessing.sharedctypes.Array

        if _original_RawValue is None:
            _original_RawValue = multiprocessing.sharedctypes.RawValue

        if _original_Value is None:
            _original_Value = multiprocessing.sharedctypes.Value

        if _original_Client is None:
            _original_Client = multiprocessing.connection.Client

        if _original_SharedMemory is None:
            _original_SharedMemory = multiprocessing.shared_memory.SharedMemory

        if _original_ShareableList is None:
            _original_ShareableList = multiprocessing.shared_memory.ShareableList

    # this method is called by BaseContext.get_context() before
    # returning the actual context for the given startmethod
    def _check_available(self):
        pass

src/dragon/mpbridge/context.py
DragonHPC/dragon
"""Dragon's replacement for the Multiprocessing Heap based on Dragon managed memory.
We effectively back `multiprocessing.heap.Arena` with Dragons managed memory.
This is a single node only implementation.
"""

import os

import multiprocessing.heap

from ..managed_memory import MemoryAlloc, MemoryPool
from ..infrastructure.facts import default_pool_muid_from_index
from ..globalservices.pool import query
from ..infrastructure.parameters import this_process
import dragon.utils as du

_DEF_MUID = default_pool_muid_from_index(this_process.index)


class DragonArena(multiprocessing.heap.Arena):
    """Dragon's version of a file backed Arena uses
    the default managed memory pool to get a buffer.
    A custom `m_uid` can be given as well.
    """

    def __init__(self, size: int, fd: int = -1, _m_uid: int = _DEF_MUID):
        """Initialize the DragonArena.

        :param size: size in bytes
        :type size: int
        :param fd: a dummy file descriptor, defaults to -1
        :type fd: int, optional
        :param _m_uid: memory pool m_uid, defaults to _DEF_MUID
        :type _m_uid: int, optional
        """

        self._closed = True

        self.size = size
        self.fd = fd  # dummy

        if _m_uid == _DEF_MUID:
            sdesc = du.B64.str_to_bytes(os.environ.get("DRAGON_DEFAULT_PD", "-1"))
        else:
            pool_descr = query(_m_uid)
            sdesc = pool_descr.sdesc

        pool = MemoryPool.attach(sdesc)
        self._mem = pool.alloc(size)
        self.buffer = self._mem.get_memview()
        pool.detach()

        self._closed = False

    def __setstate__(self, state):
        self._closed = True

        mem_ser, self.size, self.fd = state
        self._mem = MemoryAlloc.attach(mem_ser)
        self.buffer = self._mem.get_memview()

        self._closed = False

    def __getstate__(self):
        return self._mem.serialize(), self.size, self.fd

    def __del__(self):
        if not self._closed:
            self._closed = True
            try:
                self._mem.detach()
            except Exception as e:
                pass

    def _destroy(self):
        if not self._closed:
            self._closed = True
            try:
                self._mem.destroy()
            except Exception as e:
                pass

src/dragon/mpbridge/heap.py
DragonHPC/dragon
"""Dragon's replacement of Multiprocessing Managers. Not implemented at the moment.
"""

import multiprocessing.managers


class DragonSyncManager(multiprocessing.managers.SyncManager):
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonSyncManager is not implemented, yet")
        super().__init__(*args, **kwargs)


class BaseImplSyncManager(multiprocessing.managers.SyncManager):  # Dummy for now
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class DragonBaseManager(multiprocessing.managers.BaseManager):
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonBaseManager is not implemented, yet")
        super().__init__(*args, **kwargs)


class DragonBaseProxy(multiprocessing.managers.BaseProxy):
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonBaseProxy is not implemented, yet")
        super().__init__(*args, **kwargs)


class DragonToken(multiprocessing.managers.Token):
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonToken is not implemented, yet")
        super().__init__(*args, **kwargs)


class DragonSharedMemoryManager(multiprocessing.managers.SharedMemoryManager):
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonSharedMemoryManager is not implemented, yet")
        super().__init__(*args, **kwargs)


def Manager(*, ctx=None, use_base_impl=True):
    if use_base_impl:
        return BaseImplSyncManager(ctx=ctx)
    else:
        return DragonSyncManager(ctx=ctx)

src/dragon/mpbridge/managers.py
DragonHPC/dragon
"""A class to monkeypatch Dragon classes into Python Multiprocessing.

Monkeypatching fixes direct imports and is necessary, if the user asks for
Multiprocessing classes and functions without the context. In particular, people
may use direct imports before the `start_method` is actually set. Hence we have to
do this in dragons __init__.py, i.e. before `import Multiprocessing` is executed
by the user for the first time.

Consider:

.. code-block:: Python

    import dragon
    import multiprocessing as mp
    from multiprocessing.connection import wait

    if __name__ == "__main__":
        mp.set_start_method("dragon")
        ctx = mp.get_context()
        print(f"{wait.__module__=}")
        print(f"{ctx.wait.__module__=}")


without correct patching in __init__.py this would print:

.. code-block:: Python    

    wait.__module__='multiprocessing.connection'
    ctx.wait.__module__='dragon.mpbridge.context'
"""

import dragon
import multiprocessing
import os


class Monkeypatcher:
    """Class managing the monkeypatching functionality.

    We replace the Multiprocessing API at the top level and at the submodule
    level to facilitate direct imports. The difference is that the top level import returns a
    function, while the submodule import returns the class.

    To understand this, consider:

    .. code-block:: Python

        import multiprocessing as mp
        from multiprocessing.pool import Pool

        p1 = mp.Pool # is a method
        p2 = Pool # is a class
        p3 = mp.Pool() # is the object
        p4 = Pool() # is also an object

    Note that multiprocessing is not designed in the same way everywhere.
    The standard pattern is that a function (`ctx.Queue()`) does `import multiprocessing.queues.Queue`
    and returns `Queue()`.
    However, some parts of the public API are not below the context, other parts have to be explicitly
    imported. A mechanism that works for the standard pattern, will leave other parts of the API
    unchanged.

    We store the original classes and functions in private attributes for the
    brave.

    If the switching is done before the user code is parsed, this replaces _all_
    mentions of Multiprocessing objects with Dragon ones, including the
    inheritance tree.
    """

    def __init__(self):
        """Store the original Multiprocessing classes"""
        from multiprocessing.managers import BaseManager, SyncManager, BaseProxy, Token, SharedMemoryManager
        from multiprocessing.shared_memory import SharedMemory, ShareableList
        from multiprocessing.util import get_logger, log_to_stderr
        from multiprocessing.heap import Arena

        # context functions not backed by classes
        self.cpu_count = multiprocessing.cpu_count
        self.active_children = multiprocessing.active_children
        self.parent_process = multiprocessing.parent_process
        self.current_process = multiprocessing.current_process
        self.log_to_stderr = multiprocessing.log_to_stderr
        self.get_logger = multiprocessing.get_logger
        self.freeze_support = multiprocessing.freeze_support

        # context functions backed by classes
        self.Process = multiprocessing.Process

        self.Queue = multiprocessing.Queue
        self.JoinableQueue = multiprocessing.JoinableQueue
        self.SimpleQueue = multiprocessing.SimpleQueue

        self.Event = multiprocessing.Event
        self.Lock = multiprocessing.Lock
        self.RLock = multiprocessing.RLock
        self.Barrier = multiprocessing.Barrier
        self.Condition = multiprocessing.Condition
        self.Semaphore = multiprocessing.Semaphore
        self.BoundedSemaphore = multiprocessing.BoundedSemaphore

        self.Manager = multiprocessing.Manager

        self.Pipe = multiprocessing.Pipe

        self.Pool = multiprocessing.Pool

        self.Array = multiprocessing.Array
        self.RawArray = multiprocessing.RawArray
        self.Value = multiprocessing.Value
        self.RawValue = multiprocessing.RawValue

        # submodule classes and functions
        self.process_current_process = multiprocessing.process.current_process
        self.process_parent_process = multiprocessing.process.parent_process
        self.process_active_children = multiprocessing.process.active_children

        self.queues_Queue = multiprocessing.queues.Queue
        self.queues_JoinableQueue = multiprocessing.queues.JoinableQueue
        self.queues_SimpleQueue = multiprocessing.queues.SimpleQueue

        self.synchronize_Event = multiprocessing.synchronize.Event
        self.synchronize_Barrier = multiprocessing.synchronize.Barrier
        self.synchronize_Lock = multiprocessing.synchronize.Lock
        self.synchronize_RLock = multiprocessing.synchronize.RLock
        self.synchronize_Condition = multiprocessing.synchronize.Condition
        self.synchronize_Semaphore = multiprocessing.synchronize.Semaphore
        self.synchronize_BoundedSemaphore = multiprocessing.synchronize.BoundedSemaphore

        self.managers_BaseManager = BaseManager  # only accessible via direct import :-S
        self.managers_SyncManager = SyncManager
        self.managers_BaseProxy = BaseProxy
        self.managers_Token = Token
        self.managers_SharedMemoryManager = SharedMemoryManager

        self.connection_wait = multiprocessing.connection.wait
        self.connection_Pipe = multiprocessing.connection.Pipe
        self.connection_Listener = multiprocessing.connection.Listener
        self.connection_Client = multiprocessing.connection.Client

        self.pool_Pool = multiprocessing.pool.Pool
        self.pool_wait = multiprocessing.pool.wait

        self.sharedctypes_Array = multiprocessing.sharedctypes.Array
        self.sharedctypes_RawArray = multiprocessing.sharedctypes.RawArray
        self.sharedctypes_Value = multiprocessing.sharedctypes.Value
        self.sharedctypes_RawValue = multiprocessing.sharedctypes.RawValue
        self.sharedctypes_copy = multiprocessing.sharedctypes.copy
        self.sharedctypes_synchronized = multiprocessing.sharedctypes.synchronized

        self.shared_memory_SharedMemory = SharedMemory
        self.shared_memory_ShareableList = ShareableList

        self.reduction_DupFd = multiprocessing.reduction.DupFd
        self.reduction_recvfds = multiprocessing.reduction.recvfds
        self.reduction_sendfds = multiprocessing.reduction.sendfds

        self.util_get_logger = get_logger
        self.util_log_to_stderr = log_to_stderr

        self.heap_Arena = Arena

    def switch_out(self, ctx) -> None:
        """Replace the standard Multiprocessing classes and functions outside the
        context with Dragons versions.

        To do so, we need the instantiated Dragon context class which is held next to
        the other context in `multiprocessing.context._concrete_contexts`.

        :param ctx: The actual Dragon context.
        :type ctx: dragon.mpbridge.context.DragonContext
        """

        from dragon.mpbridge.managers import (
            DragonBaseManager,
            DragonSyncManager,
            DragonBaseProxy,
            DragonToken,
            DragonSharedMemoryManager,
        )
        from dragon.mpbridge.shared_memory import DragonSharedMemory, DragonShareableList
        from dragon.mpbridge.util import get_logger, log_to_stderr
        from dragon.mpbridge.heap import DragonArena

        # context functions not backed by classes
        multiprocessing.cpu_count = ctx.cpu_count
        multiprocessing.active_children = ctx.active_children
        multiprocessing.parent_process = ctx.parent_process
        multiprocessing.current_process = ctx.current_process
        multiprocessing.log_to_stderr = ctx.log_to_stderr
        multiprocessing.get_logger = ctx.get_logger
        multiprocessing.freeze_support = ctx.freeze_support

        # context functions backed by classes
        multiprocessing.Process = ctx.Process

        multiprocessing.Queue = ctx.Queue
        multiprocessing.JoinableQueue = ctx.JoinableQueue
        multiprocessing.SimpleQueue = ctx.SimpleQueue

        multiprocessing.Event = ctx.Event
        multiprocessing.Lock = ctx.Lock
        multiprocessing.RLock = ctx.RLock
        multiprocessing.Barrier = ctx.Barrier
        multiprocessing.Condition = ctx.Condition
        multiprocessing.Semaphore = ctx.Semaphore
        multiprocessing.BoundedSemaphore = ctx.BoundedSemaphore

        multiprocessing.Pipe = ctx.Pipe

        multiprocessing.Manager = ctx.Manager

        multiprocessing.Pool = ctx.Pool

        multiprocessing.Array = ctx.Array
        multiprocessing.RawArray = ctx.RawArray
        multiprocessing.Value = ctx.Value
        multiprocessing.RawValue = ctx.RawValue

        # submodule classes and functions are directly replaced to facilitate
        # direct imports of classes via e.g. `from multiprocessing.pool import Pool`
        # and subclassing, i.e. inheritance from Multiprocessing classes.

        # multiprocessing.process.current_process = dragon.mpbridge.process.current_process # we use the original
        # multiprocessing.process.parent_process = dragon.mpbridge.process.parent_process # we use the original
        # multprocessing.active_children = dragon.mpbridge.process.active_children # we use the original

        multiprocessing.queues.Queue = dragon.mpbridge.queues.DragonQueue  # WithProcessesTestQueue.test_fork
        multiprocessing.queues.JoinableQueue = (
            dragon.mpbridge.queues.DragonJoinableQueue
        )  # WithProcessesTestQueue.test_task_done
        multiprocessing.queues.SimpleQueue = dragon.mpbridge.queues.DragonSimpleQueue

        multiprocessing.synchronize.Event = dragon.mpbridge.synchronize.DragonEvent
        multiprocessing.synchronize.Lock = (
            dragon.mpbridge.synchronize.DragonLock
        )  # WithProcessesTestConnection.test_spawn_close
        multiprocessing.synchronize.RLock = dragon.mpbridge.synchronize.DragonRLock
        multiprocessing.synchronize.Barrier = dragon.mpbridge.synchronize.DragonBarrier
        multiprocessing.synchronize.Condition = dragon.mpbridge.synchronize.DragonCondition
        multiprocessing.synchronize.Semaphore = dragon.mpbridge.synchronize.DragonSemaphore
        multiprocessing.synchronize.BoundedSemaphore = dragon.mpbridge.synchronize.DragonBoundedSemaphore

        multiprocessing.connection.wait = dragon.mpbridge.context.DragonContext.wait  # TODO: PE-41642
        multiprocessing.connection.Pipe = ctx.Pipe  ## ??
        multiprocessing.connection.Listener = ctx.Listener
        multiprocessing.connection.Client = ctx.Client

        dragon_base_pool = os.getenv("DRAGON_BASEPOOL", "NATIVE")
        if  dragon_base_pool == "PATCHED": 
            multiprocessing.pool.Pool = dragon.mpbridge.pool.DragonPoolPatched
        else:        
            multiprocessing.pool.Pool = dragon.mpbridge.pool.DragonPool
        multiprocessing.pool.wait = ctx.wait

        multiprocessing.sharedctypes.Array = dragon.mpbridge.sharedctypes.DragonArray
        multiprocessing.sharedctypes.RawArray = dragon.mpbridge.sharedctypes.DragonRawArray
        multiprocessing.sharedctypes.Value = dragon.mpbridge.sharedctypes.DragonValue
        multiprocessing.sharedctypes.RawValue = dragon.mpbridge.sharedctypes.DragonRawValue
        multiprocessing.sharedctypes.copy = dragon.mpbridge.sharedctypes.dragon_copy
        multiprocessing.sharedctypes.synchronized = dragon.mpbridge.sharedctypes.dragon_synchronized

        multiprocessing.reduction.DupFd = dragon.mpbridge.reduction.dragon_DupFd
        multiprocessing.reduction.recvfds = dragon.mpbridge.reduction.dragon_recvfds
        multiprocessing.reduction.sendfds = dragon.mpbridge.reduction.dragon_sendfds

        multiprocessing.util.get_logger = get_logger
        multiprocessing.util.log_to_stderr = log_to_stderr

        # Some submodules are declared public via __all__ but are not below the context
        # and are not accessible via `multiprocessing.module.Class`. Instead they must
        # be imported with `from multiprocessing.module import Class` or `import multiprocessing.module`
        # and `obj = multiprocessing.module.Class()`

        multiprocessing.managers.BaseManager = (
            DragonBaseManager  # mimic insane multiprocessing structure here
        )
        multiprocessing.managers.SyncManager = DragonSyncManager
        multiprocessing.managers.BaseProxy = DragonBaseProxy
        multiprocessing.managers.Token = DragonToken
        multiprocessing.managers.SharedMemoryManager = DragonSharedMemoryManager

        multiprocessing.shared_memory.SharedMemory = DragonSharedMemory
        multiprocessing.shared_memory.ShareableList = DragonShareableList

        multiprocessing.heap.Arena = DragonArena

    def switch_in(self) -> None:
        """Restore standard Multiprocessing classes and functions"""

        # context functions not backed by classes
        multiprocessing.cpu_count = self.cpu_count
        multiprocessing.active_children = self.active_children
        multiprocessing.parent_process = self.parent_process
        multiprocessing.current_process = self.current_process
        multiprocessing.log_to_stderr = self.log_to_stderr
        multiprocessing.get_logger = self.get_logger
        multiprocessing.freeze_support = self.freeze_support

        # context functions
        multiprocessing.Process = self.Process

        multiprocessing.Queue = self.Queue
        multiprocessing.JoinableQueue = self.JoinableQueue
        multiprocessing.SimpleQueue = self.SimpleQueue

        multiprocessing.Event = self.Event
        multiprocessing.Lock = self.Lock
        multiprocessing.RLock = self.RLock
        multiprocessing.Barrier = self.Barrier
        multiprocessing.Condition = self.Condition
        multiprocessing.Semaphore = self.Semaphore
        multiprocessing.BoundedSemaphore = self.BoundedSemaphore

        multiprocessing.Manager = self.Manager

        multiprocessing.Pool = self.Pool

        multiprocessing.Pipe = self.Pipe

        multiprocessing.Array = self.Array
        multiprocessing.RawArray = self.RawArray
        multiprocessing.Value = self.Value
        multiprocessing.RawValue = self.RawValue

        # submodule classes and functions
        multiprocessing.queues.Queue = self.queues_Queue
        multiprocessing.queues.JoinableQueue = self.queues_JoinableQueue
        multiprocessing.queues.SimpleQueue = self.queues_SimpleQueue

        multiprocessing.connection.wait = self.connection_wait
        multiprocessing.connection.Pipe = self.connection_Pipe
        multiprocessing.connection.Listener = self.connection_Listener
        multiprocessing.connection.Client = self.connection_Client

        multiprocessing.synchronize.Event = self.synchronize_Event
        multiprocessing.synchronize.Lock = self.synchronize_Lock
        multiprocessing.synchronize.RLock = self.synchronize_RLock
        multiprocessing.synchronize.Barrier = self.synchronize_Barrier
        multiprocessing.synchronize.Condition = self.synchronize_Condition
        multiprocessing.synchronize.Semaphore = self.synchronize_Semaphore
        multiprocessing.synchronize.BoundedSemaphore = self.synchronize_BoundedSemaphore

        multiprocessing.pool.Pool = self.pool_Pool
        multiprocessing.pool.wait = self.pool_wait

        multiprocessing.managers.BaseManager = self.managers_BaseManager
        multiprocessing.managers.SyncManager = self.managers_SyncManager
        multiprocessing.managers.BaseProxy = self.managers_BaseProxy
        multiprocessing.managers.Token = self.managers_Token
        multiprocessing.managers.SharedMemoryManager = self.managers_SharedMemoryManager

        multiprocessing.sharedctypes.Array = self.sharedctypes_Array
        multiprocessing.sharedctypes.RawArray = self.sharedctypes_RawArray
        multiprocessing.sharedctypes.Value = self.sharedctypes_Value
        multiprocessing.sharedctypes.RawValue = self.sharedctypes_RawValue
        multiprocessing.sharedctypes.copy = self.sharedctypes_copy
        multiprocessing.sharedctypes.synchronized = self.sharedctypes_synchronized

        multiprocessing.shared_memory.SharedMemory = self.shared_memory_SharedMemory
        multiprocessing.shared_memory.ShareableList = self.shared_memory_ShareableList

        multiprocessing.reduction.DupFd = self.reduction_DupFd
        multiprocessing.reduction.recvfds = self.reduction_sendfds
        multiprocessing.reduction.sendfds = self.reduction_sendfds

        multiprocessing.util.get_logger = self.util_get_logger
        multiprocessing.util.log_to_stderr = self.util_log_to_stderr

        multiprocessing.heap.Arena = self.heap_Arena


# This object can be imported by other modules to obtain the original
# Multiprocessing classes.
original_multiprocessing = Monkeypatcher()


class AugmentedDefaultContext(multiprocessing.context.DefaultContext):
    """This class inserts the Dragon start method 'dragon' in Multiprocessing
    and enables context switching.
    """

    def __init__(self, context, _actual_context):
        super().__init__(context)
        self._actual_context = _actual_context

    def get_all_start_methods(self):
        return [dragon.mpbridge.context.DragonContext._name] + super().get_all_start_methods()

    def set_start_method(self, method: str, *args, **kwargs) -> None:
        """Dragons replacement monkeypatches the class as well.

        :param method: name of the start method
        :type method: str
        """
        if method == dragon.mpbridge.process.DragonPopen.method:
            original_multiprocessing.switch_out(multiprocessing.context._concrete_contexts[method])
        else:
            original_multiprocessing.switch_in()

        super().set_start_method(method, *args, **kwargs)


def patch_multiprocessing():
    """Add Dragon to the list of Multiprocessing start methods.

    This function is called when a Python program is started using the Launcher, i.e.
    Dragon is invoked with the `dragon` command. 

    1. Insert Dragon's MPBridge context into the list of Multiprocessing contexts.
    2. Replace Multiprocessings default context with our own version to swap
       Dragon objects in and out when the start methods is changed.
    3. Replace the complete exported API (`multiprocessing.X`) with Dragon versions.
    4. Replace `multiprocessing.spawns` start method function with Dragon versions.
    5. Switch out the class hierarchy so the resultion of the inheritance tree can find
       Dragon's objects before the entrypoint is reached.
    """

    import multiprocessing.connection
    import multiprocessing.spawn

    # add Dragon context to multiprocessing concrete contexts
    multiprocessing.context._concrete_contexts[
        dragon.mpbridge.process.DragonPopen.method
    ] = dragon.mpbridge.context.DragonContext()

    # make our Augmented DefaultContext the new default context
    if not isinstance(multiprocessing.context._default_context, AugmentedDefaultContext):
        multiprocessing.context._default_context = AugmentedDefaultContext(
            multiprocessing.context._default_context._default_context,
            # Preserve whatever the normal system default context is.
            multiprocessing.context._default_context._actual_context
            # If _actual_context already set, must preserve it here.
        )

    # as it dynamically built out its own module namespace to begin with.

    # this collects the Multiprocessing public API
    _all_ = (x for x in dir(multiprocessing.context._default_context) if not x.startswith("_"))

    # this inserts most of the AugmentedDefaultContext public API into the multiprocessing public API
    mp_context_api_dict = {(name, getattr(multiprocessing.context._default_context, name)) for name in _all_}
    multiprocessing.__dict__.update(mp_context_api_dict)

    # Necessary but not alone sufficient for child processes to get/set the start method correctly;
    # must also preserve any already set context in AugmentedDefaultContext instantiation above.
    multiprocessing.spawn.get_start_method = multiprocessing.context._default_context.get_start_method
    multiprocessing.spawn.set_start_method = multiprocessing.context._default_context.set_start_method

    # Switch out the original modules with the new ones
    # This replaces Multiprocessing classes even if the user has inherited from them !!
    # I.e. we are modifying the class hierarchy. This not for the faint hearted.
    original_multiprocessing.switch_out(multiprocessing.context._concrete_contexts["dragon"])

src/dragon/mpbridge/monkeypatching.py
DragonHPC/dragon
"""Dragon's replacement for Multiprocessing Pool.

By default this uses a patched version of the dragon native pool and sets
`DRAGON_BASEPOOL="NATIVE"`. The private api for this class is still under 
development. To revert to the version based on the `multiprocessing.Pool` class 
with a patched terminate_pool method, set `DRAGON_BASEPOOL="PATCHED"` in the 
environment. 
"""
import multiprocessing.pool
from multiprocessing import get_start_method
from multiprocessing.pool import TERMINATE, IMapIterator, IMapUnorderedIterator
from multiprocessing.util import debug
from threading import current_thread
from typing import Iterable, Any

from ..native.pool import Pool as NativePool
from ..native.pool import job_counter, mapstar, ApplyResult, MapResult
from ..native.process_group import ProcessGroup
from ..globalservices.process import multi_join


import collections
import itertools
import time
import threading
import os


def starmapstar(args):
    return list(itertools.starmap(args[0], args[1]))


# XXX Override Pool._terminate_pool() in order to forcibly terminate pool
# XXX processes AFTER the task (and result) handler thread(s) have
# XXX been joined. This avoids the race condition where a process may
# XXX be terminated (via SIGTERM) while holding the task queue's channel OT
# XXX or UT lock, which effectively blocks the task handler when trying
# XXX to send sentinel messages (i.e., `None`) to processes that tell them
# XXX to exit.
# XXX
# XXX NOTE: This code was copied and modified from Python 3.9.10.
def dragon_terminate_pool(
    cls,
    taskqueue,
    inqueue,
    outqueue,
    pool,
    change_notifier,
    worker_handler,
    task_handler,
    result_handler,
    cache,
):
    # this is guaranteed to only be called once
    debug(f"finalizing pool, context: {get_start_method()}")

    # Notify that the worker_handler state has been changed so the
    # _handle_workers loop can be unblocked (and exited) in order to
    # send the finalization sentinel all the workers.
    worker_handler._state = TERMINATE
    change_notifier.put(None)

    task_handler._state = TERMINATE

    debug("helping task handler/workers to finish")
    cls._help_stuff_finish(inqueue, task_handler, len(pool))

    if (not result_handler.is_alive()) and (len(cache) != 0):
        raise AssertionError("Cannot have cache with result_hander not alive")

    result_handler._state = TERMINATE
    change_notifier.put(None)
    outqueue.put(None)  # sentinel

    # We must wait for the worker handler to exit before terminating
    # workers because we don't want workers to be restarted behind our back.
    debug("joining worker handler")
    if current_thread() is not worker_handler:
        worker_handler.join()

    # XXX Moved the loop to foribly terminate pool processes to after task
    # XXX and result handler threads have been joined.

    debug("joining task handler")
    if current_thread() is not task_handler:
        task_handler.join()

    debug("joining result handler")
    if current_thread() is not result_handler:
        result_handler.join()

    # XXX Terminate worker processes which have not already finished and
    # XXX then join on them.
    if pool and hasattr(pool[0], "terminate"):
        # TODO, since we're patching anyway, lets give the workers a chance
        # to clean themselves up before doing the dangerous thing of killing
        # them all
        puids = [p.sentinel for p in pool]

        debug("joining workers")
        object_tuples, proc_status = multi_join(puids, 2.0, join_all=True)
        if not object_tuples:
            debug(f"multi_join not able to join all pool workers")
            debug(f"proc_status = {proc_status}")

        debug("terminating workers")
        for p in pool:
            if p.exitcode is None:
                debug(f"pool worker has not exited. Terminating {p.pid}")
                p.terminate()

        debug("joining pool workers")
        for p in pool:
            if p.is_alive():
                # worker has not yet exited
                debug("cleaning up worker %d" % p.pid)
                p.join()


class DragonPoolPatched(multiprocessing.pool.Pool):  # Dummy
    """Dragon's patched implementation of Multiprocessing Pool."""

    _terminate_pool = classmethod(dragon_terminate_pool)

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class WrappedResult:
    """Returned by all functions that return a result. Wraps ApplyResult and MapResult so that correct timeout error can be raised."""

    def __init__(self, result: ApplyResult or MapResult = None):
        """Initializes wrapped result by saving input result

        :param result: returned result from native pool, defaults to None
        :type result: ApplyResult or MapResult, optional
        """
        self._result = result

    def get(self, timeout: float = None) -> Any:
        """Retrieves returned values of work submitted to pool

        :param timeout: timeout for getting result, defaults to None
        :type timeout: float, optional
        :raises multiprocessing.TimeoutError: raised if result is not ready in specified timeout
        :return: value returned by `func(*args, **kwargs)`
        :rtype: Any
        """
        try:
            return self._result.get(timeout=timeout)
        except TimeoutError:
            raise multiprocessing.TimeoutError

    def ready(self) -> bool:
        """Checks whether the result is ready.

        :return: returns True if the result is ready or False if it isn't
        :rtype: bool
        """
        return self._result.ready()

    def successful(self) -> bool:
        """Checks if the result is ready and returns whether the function call was successful.

        :raises ValueError: raised if result is not ready
        :return: returns True if function call was successful
        :rtype: bool
        """
        return self._result.successful()

    def wait(self, timeout: float = None) -> None:
        """Waits on event for result to be ready

        :param timeout: timeout indicating how long to wait for result to be ready, defaults to None
        :type timeout: float, optional
        """
        self._result.wait(timeout)


class WrappedPG:
    """Wraps a ProcessGroup and makes the interface compliant with how self._pool is used in the multiprocessing pool unittests."""

    def __init__(self, process_group: ProcessGroup = None):
        """Initializes wrapped process group

        :param process_group: ProcessGroup used by native pool implementation, defaults to None
        :type process_group: ProcessGroup, optional
        """
        self._pg = process_group

    def __len__(self) -> int:
        """Returns the number of processes in the process group

        :return: the number of processes in the process group
        :rtype: int
        """
        return self._pg.nproc


class DragonPool(NativePool):
    """Dragon's replacement for Multiprocessing Pool."""

    def __init__(self, *args, context=None, **kwargs):
        super().__init__(*args, **kwargs)
        self._pool = WrappedPG(self._pg)

    def apply_async(
        self,
        func: callable,
        args: tuple = (),
        kwargs: dict = {},
        callback: callable = None,
        error_callback: callable = None,
    ) -> WrappedResult:
        """Equivalent to calling `func(*args, **kwargs)` in a non-blocking way. A result is immediately returned and then updated when `func(*args, **kwargs)` has completed.

        :param func: user function to be called by worker function
        :type func: callable
        :param args: input args to func, defaults to ()
        :type args: tuple, optional
        :param kwargs: input kwargs to func, defaults to {}
        :type kwargs: dict, optional
        :param callback: user provided callback function, defaults to None
        :type callback: callable, optional
        :param error_callback: user provided error callback function, defaults to None
        :type error_callback: callable, optional
        :raises ValueError: raised if pool has already been closed or terminated
        :return: A result that has a `get` method to retrieve result of `func(*args, **kwargs)`
        :rtype: WrappedResult
        """
        return WrappedResult(
            super().apply_async(
                func=func, args=args, kwargs=kwargs, callback=callback, error_callback=error_callback
            )
        )

    def map_async(
        self,
        func: callable,
        iterable: Iterable,
        chunksize: int = None,
        callback: callable = None,
        error_callback: callable = None,
    ) -> WrappedResult:
        """Apply `func` to each element in `iterable`. The results
        are collected in a list that is returned immediately. It is equivalent to map if the `get` method is called immediately on the returned result.

        :param func: user provided function to call on elements of iterable
        :type func: callable
        :param iterable: input args to func
        :type iterable: iterable
        :param chunksize: size of work elements to be submitted to input work queue, defaults to None
        :type chunksize: int, optional
        :param callback: user provided callback function, defaults to None
        :type callback: callable, optional
        :param error_callback: user provided error callback function, defaults to None
        :type error_callback: callable, optional
        :return: A result that has a `get` method that returns an iterable of the output from applying `func` to each element of iterable.
        :rtype: WrappedResult
        """
        return WrappedResult(super()._map_async(func, iterable, mapstar, chunksize, callback, error_callback))

    def apply(self, func: callable, args: tuple = (), kwds: dict = {}) -> Any:
        """Equivalent to calling `func(*args, **kwargs)` in a blocking way. The function returns when `func(*args, **kwargs)` has completed and the result has been updated with the output

        :param func: user function to be called by worker
        :type func: callable
        :param args: input args to func, defaults to ()
        :type args: tuple, optional
        :param kwds: input kwargs to func, defaults to {}
        :type kwds: dict, optional
        :return: The result of `func(*args, **kwargs)`
        :rtype:
        """
        return WrappedResult(self.apply_async(func, args, kwds)).get()

    def map(self, func: callable, iterable: Iterable, chunksize: int = None) -> Iterable[Any]:
        """Apply `func` to each element in `iterable`, collecting the results
        in a list that is returned.

        :param func: user provided function to call on elements of iterable
        :type func: callable
        :param iterable: input args to func
        :type iterable: Iterable
        :param chunksize: size of work elements to be submitted to input work queue, defaults to None
        :type chunksize: int, optional
        :return: list of results from applying `func` to each element of input iterable
        :rtype: Iterable[Any]
        """
        return WrappedResult(self.map_async(func, iterable, chunksize)).get()

    def starmap(self, func: callable, iterable: Iterable[Iterable], chunksize: int = None) -> Iterable[Any]:
        """Like `map()` method but the elements of the `iterable` are expected to be iterables as well and will be unpacked as arguments. Hence `func` and (a, b) becomes func(a, b).


        :param func: user provided function to call on elements of iterable
        :type func: callable
        :param iterable: input iterable args to func
        :type iterable: Iterable[Iterable]
        :param chunksize: size of work elements to be submitted to input work queue, defaults to None
        :type chunksize: int, optional
        :return: results of applying `func` to input iterable args
        :rtype: Iterable[Any]
        """
        return WrappedResult(self._map_async(func, iterable, starmapstar, chunksize)).get()

    def starmap_async(
        self,
        func: callable,
        iterable: Iterable[Iterable],
        chunksize: int = None,
        callback: callable = None,
        error_callback: callable = None,
    ) -> WrappedResult:
        """Asynchronous version of `starmap()` method.

        :param func: user provided function to call on elements of iterable
        :type func: callable
        :param iterable: input iterable args to func
        :type iterable: Iterable[Iterable]
        :param chunksize: size of work elements to be submitted to input work queue, defaults to None
        :type chunksize: int, optional
        :param callback: user provided callback function, defaults to None
        :type callback: callable, optional
        :param error_callback: user provided error callback function, defaults to None
        :type error_callback: callable, optional
        :return: A result that has a `get` method that returns an iterable of the output from applying `func` to each element of iterable.
        :rtype: WrappedResult
        """
        return WrappedResult(
            self._map_async(func, iterable, starmapstar, chunksize, callback, error_callback)
        )

    def __enter__(self):
        self._check_running()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.terminate()
        self.join()

    def imap_unordered(
        self, func: callable, iterable: Iterable, chunksize: int = 1
    ) -> IMapUnorderedIterator or Any:
        """Like `imap()` method but ordering of results is arbitrary.

        :param func: user provided function to call on elements of iterable
        :type func: callable
        :param iterable: input args to func
        :type iterable: Iterable
        :param chunksize: size of work elements to be submitted to input work queue, defaults to 1
        :type chunksize: int, optional
        :raises ValueError: raised if chunksize is less than one
        :return: results of calling `func` on `iterable` elements
        :rtype: IMapUnorderedIterator or Any
        """
        self._check_running()
        if chunksize == 1:
            result = IMapUnorderedIterator(self)
            # wait to generate more tasks till last task generation thread has finished
            if self._map_launch_thread is not None:
                self._map_launch_thread.join()
            # start thread to put tasks into the input queue
            self._map_launch_thread = threading.Thread(
                target=NativePool._partition_tasks,
                args=(self._inqueue, result, func, iterable, self._end_threading_event, result._set_length),
            )
            self._map_launch_thread.start()
            return result
        else:
            if chunksize < 1:
                raise ValueError(f"Chunksize must be 1+, not {chunksize}")
            task_batches = NativePool._get_tasks(func, iterable, chunksize)
            result = IMapUnorderedIterator(self)
            self._map_launch_thread = threading.Thread(
                target=NativePool._partition_tasks,
                args=(
                    self._inqueue,
                    result,
                    mapstar,
                    task_batches,
                    self._end_threading_event,
                    result._set_length,
                ),
            )
            self._map_launch_thread.start()
            return (item for chunk in result for item in chunk)

    def imap(self, func: callable, iterable: Iterable, chunksize: int = 1) -> IMapIterator or Any:
        """Equivalent of `map()` -- can be MUCH slower than `Pool.map()`.
        Unlike `map()`, the iterable isn't immediately turned into a list.
        Rather, it is iterated over and items are put into the queue one at a
        time if chunksize=1. Unlike the non-imap variants of map, the returned
        object can then be iterated over and individual results will be
        returned as soon as they are available rather than when all work
        submitted via the iterable is done.

        :param func: user provided function to call on elements of iterable
        :type func: callable
        :param iterable: input args to func
        :type iterable: Iterable
        :param chunksize: size of work elements to be submitted to input work queue, defaults to 1
        :type chunksize: int, optional
        :raises ValueError: raised if chunksize is less than one
        :return: results of calling `func` on `iterable` elements
        :rtype: IMapIterator or Any
        """
        self._check_running()
        if chunksize == 1:
            result = IMapIterator(self)
            # wait to generate more tasks till last task generation thread has finished
            if self._map_launch_thread is not None:
                self._map_launch_thread.join()
            # start thread to put tasks into the input queue
            self._map_launch_thread = threading.Thread(
                target=NativePool._partition_tasks,
                args=(self._inqueue, result, func, iterable, self._end_threading_event, result._set_length),
            )
            self._map_launch_thread.start()
            return result
        else:
            if chunksize < 1:
                raise ValueError(f"Chunksize must be 1+, not {chunksize}")
            # generate batches of work based on chunksize
            task_batches = NativePool._get_tasks(func, iterable, chunksize)
            result = IMapIterator(self)
            # start thread that puts batches of work into input queue
            self._map_launch_thread = threading.Thread(
                target=NativePool._partition_tasks,
                args=(
                    self._inqueue,
                    result,
                    mapstar,
                    task_batches,
                    self._end_threading_event,
                    result._set_length,
                ),
            )
            self._map_launch_thread.start()
            return (item for chunk in result for item in chunk)


class BaseImplPool(multiprocessing.pool.Pool):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


def Pool(
    processes=None, initializer=None, initargs=(), maxtasksperchild=None, context=None, *, use_base_impl=True
):
    #    dragon_base_pool = int(os.getenv("DRAGON_BASEPOOL", "0"))
    dragon_base_pool = os.getenv("DRAGON_BASEPOOL", "NATIVE")
    if use_base_impl:
        return BaseImplPool(processes, initializer, initargs, maxtasksperchild, context=context)
    elif dragon_base_pool == "PATCHED":
        return DragonPoolPatched(processes, initializer, initargs, maxtasksperchild, context=context)
    else:
        return DragonPool(processes, initializer, initargs, maxtasksperchild, policy=None)

src/dragon/mpbridge/pool.py
DragonHPC/dragon
"""Dragon's replacement classes for Multiprocessing Process.
"""
import io
import logging
import os
import pickle
import multiprocessing
import multiprocessing.context
import multiprocessing.popen_spawn_posix
import multiprocessing.process
import multiprocessing.spawn
import multiprocessing.queues
import multiprocessing.connection
from multiprocessing.util import sub_debug

from ..globalservices.process import query, join, multi_join, create_with_argdata, kill, ProcessError
from ..infrastructure.parameters import this_process

LOG = logging.getLogger(__name__)


class PUID(int):
    """The purpose of this class is to use where such things like a fd (sentinel) and PID are found"""

    def thread_wait(self, timeout, done_ev, ready):
        """utility function for wait implementation

        :param timeout: _description_
        :type timeout: _type_
        :param done_ev: _description_
        :type done_ev: _type_
        :param ready: _description_
        :type ready: _type_
        """
        if not done_ev.is_set():
            done_res = join(int(self), timeout)
            if not done_ev.is_set():
                done_ev.set()
                if done_res is not None:
                    ready.append(self)


class DragonPopen(multiprocessing.popen_spawn_posix.Popen):
    """Dragon's class to spawn/open a new process."""

    method = "dragon"

    def _launch(self, process_obj):
        LOG.debug("launching")
        prep_data = multiprocessing.spawn.get_preparation_data(process_obj._name)
        fp = io.BytesIO()
        multiprocessing.context.set_spawning_popen(self)
        try:
            multiprocessing.context.reduction.dump(prep_data, fp)

            process_obj._p_p_uid = PUID(this_process.my_puid)

            multiprocessing.context.reduction.dump(process_obj, fp)
            sub_debug(f"{self} finished reduction.dump of process_obj self._fds={self._fds}")
        finally:
            multiprocessing.context.set_spawning_popen(None)

        # keep all this the same, and deliver fp as a BytesIO object for now
        cmd = multiprocessing.spawn.get_command_line()

        position = cmd.index("-c") + 1
        cmd[position] = "from dragon.mpbridge import dragon_spawn_main; dragon_spawn_main()"

        process_obj.proc_desc = create_with_argdata(
            exe=cmd[0], run_dir="", args=cmd[1:], env={}, argdata=pickle.dumps(fp, protocol=5)
        )

        self.p_uid = PUID(process_obj.proc_desc.p_uid)
        self.pid = self.p_uid
        self.sentinel = self.p_uid
        self.proc_desc = process_obj.proc_desc
        LOG.debug(f"launched {self.p_uid} as pid {self.pid}")

    def poll(self, flag=os.WNOHANG):
        if self.returncode is None:
            if flag & os.WNOHANG:
                current_desc = query(self.p_uid)
                if current_desc.state == current_desc.State.DEAD:
                    self.returncode = current_desc.ecode
            else:
                self.returncode = join(self.p_uid)

        return self.returncode

    def wait(self, timeout=None):
        if self.returncode is None:
            if timeout is None:
                return self.poll(flag=0)
            else:
                self.returncode = join(self.p_uid, timeout=timeout)

        return self.returncode

    def _send_signal(self, sig):
        if self.returncode is None:
            try:
                kill(self.p_uid, sig)
            except ProcessError:
                pass


class DragonProcess(multiprocessing.process.BaseProcess):
    """Dragon's Process uses most of the base process implementation.
    """
    _start_method = DragonPopen.method

    @staticmethod
    def _Popen(process_obj):
        return DragonPopen(process_obj)

    @property
    def ident(self) -> int:
        """Property method uniquely identifying the process.

        :return: The p_uid of this process, None if it has not been started.
        :rtype: int
        """
        if self._popen:
            return self._popen.p_uid
        else:
            return None

    pid = ident

    def _cleanup():
        if not multiprocessing.process._children:
            return
        # Check for processes which have finished
        p_uids = [int(p._popen.p_uid) for p in multiprocessing.process._children]
        dead, _ = multi_join(p_uids, timeout=0)
        if dead is None:  # All child processes are still alive
            return
        # Convert (p_uid, status) list of dead child processes to dict for convenience
        dead = dict(dead)
        # Remove dead processes from _children set
        for p in list(multiprocessing.process._children):
            if int(p._popen.p_uid) in dead:
                multiprocessing.process._children.discard(p)


class DragonPopenStub:
    """This Popen class replaces the actual DragonPopen in process objects returned
    to a child process by e.g. `mp.current_process()`. The parent will see the correct
    DragonPopen class, as they have actually started the child.  This class only contains
    what is necessary to make the process object work for the child with Multiprocessings
    internal methods that we have not modified.
    """

    def __init__(self, p_uid):
        self.p_uid = p_uid
        self.pid = p_uid
        self.returncode = None

    def poll(self, flag=os.WNOHANG):
        """So we can show the correct exitcode when printing a process
        object returned by mp.current_process()
        """
        if self.returncode is None:
            if flag & os.WNOHANG:
                current_desc = query(self.p_uid)
                if current_desc.state == current_desc.State.DEAD:
                    self.returncode = current_desc.ecode
            else:
                self.returncode = join(self.p_uid)

        return self.returncode

src/dragon/mpbridge/process.py
DragonHPC/dragon
"""Dragon's replacement of Multiprocessing Queue objects based on Channels.
""" 

import multiprocessing.queues
import logging
import pickle
import queue

from . import DragonPopen
from ..channels import ChannelError, Many2ManyWritingChannelFile
from dragon.native.queue import Queue as NativeQueue

LOG = logging.getLogger(__name__)


def restrict_public_api(public_api: list) -> callable:
    """ "Decorator that removes public methods of the decorated class
    not in public_api. To do so it changes the namespace the first parent & grandparent
    class. Breaks inheritance.

    :param public_api: contains the names of all valid public attributes and methods
    :type public_api: list
    :return: returns the cleaned class
    :rtype: type
    """

    def ancestor_method_remover(cls: type) -> type:

        parent_class = cls.__bases__[0]

        grandparent_class = parent_class.__bases__[0]

        grandparent_namespace = dict(grandparent_class.__dict__)

        attr = [item for item in grandparent_namespace.keys() if not item.startswith("_")]
        for name in attr:
            if name not in public_api:
                LOG.debug(f"Removing {name} from namespace of {cls.__name__}")
                _ = grandparent_namespace.pop(name)

        _ = grandparent_namespace.pop("__dict__")  # do not include a false __dict__

        new_grandparent_class = type(
            grandparent_class.__name__,
            grandparent_class.__bases__,
            grandparent_namespace,
        )

        parent_namespace = dict(parent_class.__dict__)

        attr = [item for item in parent_namespace.keys() if not item.startswith("_")]
        for name in attr:
            if name not in public_api:
                LOG.debug(f"Removing {name} from namespace of {cls.__name__}")
                _ = parent_namespace.pop(name)

        new_parent_class = type(
            parent_class.__name__,
            (new_grandparent_class,) + parent_class.__bases__[1:],
            parent_namespace,
        )

        cls.__bases__ = (new_parent_class,) + cls.__bases__[1:]

        return cls

    return ancestor_method_remover


class FakeConnectionHandle:
    """A placeholder connection handle for a queue that has been reset"""

    # would like to use types.SimpleNamespace, but this is unhashable.

    pass


class PatchedDragonNativeQueue(NativeQueue):
    """The Dragon native queue including patches for Python Multiprocessing
    The API has to exactly match the Multiprocessing.queues.Queue API.
    """

    def __init__(self, *args, **kwargs):  # need to mimic Multiprocessing API here
        NativeQueue.__init__(self, *args, **kwargs)  # cannot call super() here
        self._reset()

    def __setstate__(self, *args, **kwargs):
        NativeQueue.__setstate__(self, *args, **kwargs)
        self._reset()

    def qsize(self) -> int:
        """Return the approximate size of the queue. This number may not be reliable.

        :raises ValueError: If the queue is closed
        :return: approximate number of items in the queue
        :rtype: int
        """

        return self._size()

    def put(self, obj, block=True, timeout=None) -> None:
        """Puts the serialization of an object onto the queue.
        This version includes patches for the multiprocessing unit tests

        :param obj: object to serialize and put
        :param block: Whether to block
        :param timeout: Timeout, if blocking.  None means infinity, default
        :return: None
        """
        try:
            self._put(obj, block=block, timeout=timeout)
        except AttributeError as e:  # WithProcessesTestQueue.test_queue_feeder_on_queue_feeder_error
            # and WithProcessesTestQueue.test_queue_feeder_donot_stop_onexc
            self._on_queue_feeder_error(AttributeError(e), obj)

    def close(self) -> None:
        """Indicate that no more data will be put on this queue by the current process.
        The refcount of the background channel will be released and the channel might be
        removed, if no other processes are holding it.
        This version includes a patch for the multiprocessing unit tests.
        """

        self._close()

        # TestSimpleQueue.test_closed
        self._writer.closed = True
        self._reader.closed = True

    def join_thread(self) -> None:
        """This method is a no-op for Dragon based multiprocessing and exists
        solely for compatibility with the queues.Queue API.
        """
        pass

    def cancel_join_thread(self):
        """This method is a no-op for Dragon based multiprocessing and exists
        solely for compatibility with the queues.Queue API.
        """
        pass

    def _reset(self):
        """Add a fake connection object below the queue for Multiprocessing"""

        self._writer = FakeConnectionHandle()
        self._writer.close = lambda: None
        self._writer.closed = False
        self._writer.send = self.put

        self._reader = FakeConnectionHandle()
        self._reader.close = lambda: None
        self._reader.closed = False
        self._reader.poll = self._poll
        self._reader.recv = self.get
        self._reader.thread_wait = self._thread_wait

        self._rlock = FakeConnectionHandle()
        self._rlock.acquire = lambda: None

    @staticmethod
    def _on_queue_feeder_error(e, obj):
        """
        Private API hook called when feeding data in the background thread
        raises an exception.  For overriding by concurrent.futures.
        """
        # see WithProcessesTestQueue.test_queue_feeder_on_queue_feeder_error

        import traceback

        traceback.print_exc()


_QUEUE_API = [
    "cancel_join_thread",
    "close",
    "empty",
    "full",
    "get",
    "get_nowait",
    "join_thread",
    "put",
    "put_nowait",
    "qsize",
]


@restrict_public_api(_QUEUE_API)
class DragonQueue(PatchedDragonNativeQueue):
    def __init__(self, *args, ctx=None, **kwargs):
        super().__init__(*args, joinable=False, **kwargs)


_JOINABLE_QUEUE_API = [
    "cancel_join_thread",
    "close",
    "empty",
    "full",
    "get",
    "get_nowait",
    "join",
    "join_thread",
    "put",
    "put_nowait",
    "qsize",
    "task_done",
]


@restrict_public_api(_JOINABLE_QUEUE_API)
class DragonJoinableQueue(PatchedDragonNativeQueue):
    def __init__(self, *args, ctx=None, **kwargs):
        super().__init__(*args, joinable=True, **kwargs)


_SIMPLE_QUEUE_API = ["close", "empty", "get", "put"]


@restrict_public_api(_SIMPLE_QUEUE_API)
class DragonSimpleQueue(PatchedDragonNativeQueue):
    def __init__(self, *args, ctx=None, **kwargs):
        super().__init__(*args, joinable=False, **kwargs)


class BaseImplQueue(multiprocessing.queues.Queue):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # patches internal dragon.infrastructure.connection.Connection
        # to not send EOT when it is closed.
        self._writer.ghost = True

    def _poll(self, timeout=0):  # for dragon wait
        return not self.empty()


class BaseImplJoinableQueue(multiprocessing.queues.JoinableQueue):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # patches internal dragon.infrastructure.connection.Connection
        # to not send EOT when it is closed.
        self._writer.ghost = True

    def _poll(self, timeout=0):  # for dragon wait
        return not self.empty()


class BaseImplSimpleQueue(multiprocessing.queues.SimpleQueue):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        # patches internal dragon.infrastructure.connection.Connection
        # to not send EOT when it is closed.
        self._writer.ghost = True

    def _poll(self, timeout=0):  # for dragon wait
        return not self.empty()


def Queue(maxsize=0, *, ctx=None, use_base_impl=True):
    if use_base_impl:
        return BaseImplQueue(maxsize, ctx)
    else:

        if ctx is not None:
            assert DragonPopen.method == ctx._name

        # NOTE: when maxsize=0, local_opts.capacity is set to 100 (default)
        # whereas in the mp documentation, maxsize equal to zero means
        # infinite queue size
        if maxsize <= 0:
            maxsize = 100

        return DragonQueue(maxsize, ctx=ctx)


def JoinableQueue(maxsize=0, *, ctx=None, use_base_impl=True):
    if use_base_impl:
        return BaseImplJoinableQueue(maxsize, ctx)
    else:

        if ctx is not None:
            assert DragonPopen.method == ctx._name

        if maxsize <= 0:
            maxsize = 100

        return DragonJoinableQueue(maxsize, ctx=ctx)


def SimpleQueue(*, ctx=None, use_base_impl=True):
    if use_base_impl:
        return BaseImplSimpleQueue(ctx)
    else:
        if ctx is not None:
            assert DragonPopen.method == ctx._name

        return DragonSimpleQueue(ctx=ctx)

src/dragon/mpbridge/queues.py
DragonHPC/dragon
"""Dragon's replacements for Multiprocessings 'reduction' module.

To prevent users from sharing file descriptors between child processes, we always 
raise an exception.

See the documentation 'Multiprocessing with Dragon'.

:raises NotImplementedError: Always raise this.
"""

import array
import socket


def dragon_DupFd(fd: int) -> int:
    """Duplicates a file descriptor for a child process in Multiprocessing.

    Dragon does not support sharing file descriptors as it does not generalize to multi-node
    or federated systems.

    See the documentation 'Multiprocessing with Dragon'.

    :param fd: the file descriptor of the parent process to be duplicated
    :type fd: int
    :return: the duplicated file descriptor
    :rtype: int
    :raises NotImplementedError: Always raises an exception.
    """
    raise NotImplementedError(f"Dragon does not support duplicate file descriptors for child processes.")


def dragon_sendfds(sock: socket.socket, fds: array.array) -> None:
    """Send an array of file descriptors to another process via a socket.

    Dragon does not support this behaviour as it does not generalize to multi-node
    or federated systems.

    See the documentation 'Multiprocessing with Dragon'.

    :param sock: the socket to be used for sending
    :type sock: socket.socket
    :param fds: the array of file descriptors
    :type fds: array.array
    :raises NotImplementedError: Always raises an exception.
    """
    raise NotImplementedError(f"Dragon does not support sending file descriptors to other processes.")


def dragon_recvfds(sock: socket.socket, size: int) -> list:
    """Receive a list of file descriptors from another process via a socket.

    Dragon does not support this behaviour as it does not generalize to multi-node
    or federated systems.

    See the documentation 'Multiprocessing with Dragon'.

    :param sock: the socket to be used for sending
    :type sock: socket.socket
    :param size: the number of file descriptor to receive.
    :type size: int
    :raises NotImplementedError: Always raises an exception.
    :return: the file descriptors
    :rtype: list
    """
    raise NotImplementedError(f"Dragon does not support sending file descriptors to other processes.")

src/dragon/mpbridge/reduction.py
DragonHPC/dragon
"""Dragon's replacement for Multiprocessing Shared Memory.

Not Implemented.
""" 

import multiprocessing.shared_memory


class DragonSharedMemory(multiprocessing.shared_memory.SharedMemory):  # Dummy for now
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonSharedMemory has not been implemented with Dragon, yet")
        super().__init__(*args, **kwargs)


class DragonShareableList(multiprocessing.shared_memory.ShareableList):  # Dummy for now
    def __init__(self, *args, **kwargs):
        raise NotImplementedError(f"DragonShareableList has not been implemented with Dragon, yet")
        super().__init__(*args, **kwargs)

src/dragon/mpbridge/shared_memory.py
DragonHPC/dragon
"""Dragon's replacement for Multiprocessing's shared ctypes objects: Value and Array.
"""

import dragon
from ..native.lock import Lock
from ..native.value import Value
from ..native.array import Array
from ..globalservices.channel import create, get_refcnt, release_refcnt


def dragon_copy(obj: object) -> object:
    """make a copy of a shared object

    :param obj: original object
    :type obj: obj
    :return: copy of object
    :rtype: obj
    """
    raise NotImplementedError("dragon_copy not implemented yet")


def dragon_synchronized(obj: object, lock=None, ctx=None) -> object:
    """Synchronize an object by wrapping it.

    :param obj: the object to synchronize
    :type obj: object
    :param lock: multiprocessing lock for the synchronized object, defaults to None
    :type lock: multiprocessing.synchronize.Lock, optional
    :param ctx: multiprocessing context to use, defaults to None
    :type ctx: multiprocessing.context.Context, optional
    :return: A synchronized version of the object
    :rtype: object
    """
    raise NotImplementedError("dragon_synchronized not implemented yet")


def RawValue(typecode_or_type, *args, original=None, use_base_impl=True):
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(typecode_or_type, *args)
    else:
        return DragonRawValue(typecode_or_type, *args)


def RawArray(typecode_or_type, size_or_initializer, original=None, use_base_impl=True):
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(typecode_or_type, size_or_initializer)
    else:
        return DragonRawArray(typecode_or_type, size_or_initializer)


def Value(typecode_or_type, *args, lock=True, ctx=None, original=None, use_base_impl=False):
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(typecode_or_type, *args, lock=lock, ctx=ctx)
    else:
        return DragonValue(typecode_or_type, *args, lock=lock, ctx=ctx)


def Array(
    typecode_or_type, size_or_initializer, *args, lock=True, ctx=None, original=None, use_base_impl=False
):
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(typecode_or_type, size_or_initializer, *args, lock=lock, ctx=ctx)
    else:
        return DragonArray(typecode_or_type, size_or_initializer, *args, lock=lock, ctx=ctx)


class DragonRawValue(dragon.native.value.Value):
    """
    RawValue Class replacement for Value multiprocessing test cases
    """

    def __init__(self, typecode_or_type, value: int = 0, *, ctx: None = None, raw: bool = True):
        """Initialize the mpbridge RawValue object.
        :param typecode_or_type: the typecode or type is returned from the dictionary, typecode_to_type
        :type typecode_or_type: str or ctypes, required
        :param value: the value for the object
        :type value: int, optional
        :param raw: bool, optional
        :type raw: sets whether lock is used or not
        """
        super().__init__(typecode_or_type, value)


class DragonValue(dragon.native.value.Value):
    """
    Value Class replacement
    """

    def __del__(self):
        try:
            cuid = self._channel.cuid
            self._channel.detach()
            release_refcnt(cuid)
        except AttributeError:
            pass

    def __getstate__(self):
        return {
            "base_state": super().__getstate__(),
            "mp_state": (self.get_lock, self.get_obj, self._lock),
        }

    def __setstate__(self, state):
        super().__setstate__(state["base_state"])
        (self.get_lock, self.get_obj, self._lock) = state["mp_state"]

    def __repr__(self):
        return (
            f"Dragon Multiprocessing Value({self._type}, {self.value}, {self._channel.cuid}, {self._muid})"
        )

    def __init__(self, typecode_or_type, value: int = 0, *, ctx: None = None, lock: Lock = True):
        """Initialize the mpbridge value object.
        :param typecode_or_type: the typecode or type is returned from the dictionary, typecode_to_type
        :type typecode_or_type: str or ctypes, required
        :param value: the value for the object
        :type value: int, optional
        :param m_uid: memory pool to create the channel in and message to write value and typecode_or_type in managed memory, defaults to _DEF_MUID
        :type m_uid: int, optional
        :param lock: dragon.native.lock.Lock, optional
        :type lock: creates lock for synchronization for value
        """

        # if lock is False, return the subclass value
        if lock is False:
            super().__init__(typecode_or_type, value)
            return

        elif lock in (True, None) or isinstance(lock, dragon.mpbridge.synchronize.DragonLock):
            # set attributes for lock get_obj and get_lock
            if lock is True:
                lock = Lock(recursive=True)
            self._lock = lock
            super().__init__(typecode_or_type, value)
            self.get_lock = self._get_lock
            self.get_obj = self._type
            return

        # lock is not a valid type
        else:
            raise AttributeError

    def acquire(self):
        return self._lock.acquire()

    def release(self):
        return self._lock.release()

    def _get_lock(self):
        return self._lock

    def __enter__(self):
        self.acquire()

    def __exit__(self, *args):
        self.release()

class DragonRawArray(dragon.native.array.Array):
    """
    RawArray Class replacement for Array multiprocessing test cases
    """

    def __repr__(self):
        return f"{self.__class__.__name__}(typecode_or_type={self._type}, m_uid={self._muid})"

    def __init__(self, typecode_or_type, size_or_initializer, ctx: None = None, raw: bool = True):
        """Initialize the mpbridge RawArray object.
        :param typecode_or_type: the typecode or type is returned from the dictionary, typecode_to_type
        :type typecode_or_type: str or ctypes, required
        :param size_or_initializer: the array for the object
        :type size_or_initializer: range, int, list, required
        :param raw: bool, optional
        :type raw: sets whether lock is used or not
        """
        super().__init__(typecode_or_type, size_or_initializer, lock=False)

class DragonArray(dragon.native.array.Array):
    """
    Array Class replacement
    """

    def __getstate__(self):
        return {
            "base_state": super().__getstate__(),
            "mp_state": (self.get_lock, self.get_obj, self._lock),
        }

    def __setstate__(self, state):
        super().__setstate__(state["base_state"])
        (self.get_lock, self.get_obj, self._lock) = state["mp_state"]

    def __repr__(self):
        return f"{self.__class__.__name__}(typecode_or_type={self._type}, lock={self._lock}, m_uid={self._muid})"

    def __del__(self):
        try:
            cuid = self._channel.cuid
            self._channel.detach()
            release_refcnt(cuid)
        except AttributeError:
            pass

    def __init__(self, typecode_or_type,  size_or_initializer, lock: Lock, ctx: None = None):
        """Initialize the mpbridge array object.

        :param typecode_or_type: the typecode or type is returned from the dictionary, typecode_to_type
        :type typecode_or_type: str or ctypes, required
        :param size_or_initializer: the array for the object
        :type size_or_initializer: range, int, list, required
        :param lock: dragon.native.lock.Lock, optional
        :type lock: creates lock for synchronization for array
        """
        super().__init__(typecode_or_type, size_or_initializer, lock)
src/dragon/mpbridge/sharedctypes.py
DragonHPC/dragon
"""Dragon's replacements for the synchronization primitives in Multiprocessing.  
Except for `Condition`, all components are based on `dragon.native` component
and are implemented over one or more dragon channels.  `Lock` is patched with a
dummy `_SemLock` class, so `Condition` works out of the box.
"""

import multiprocessing.synchronize
from threading import BrokenBarrierError

import dragon.native.lock
import dragon.native.semaphore
import dragon.native.event
import dragon.native.barrier

# Lock


class _SemLock:
    """Fake Semlock to fool Multiprocessing"""

    def __init__(self, parent):
        self._parent = parent

    def _is_mine(self):
        return self._parent.is_mine()

    def _count(self):
        return self._parent._accesscount


class AugmentedDragonNativeLock(dragon.native.lock.Lock):
    """Augment the Dragon Native Lock so it can be used by Python Multiprocessing"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._semlock = _SemLock(self)

    def __setstate__(self, state):
        super().__setstate__(state)
        self._semlock = _SemLock(self)


class DragonLock(AugmentedDragonNativeLock):
    def __init__(self, *args, ctx, **kwargs):
        super().__init__(*args, **kwargs, recursive=False)


class DragonRLock(AugmentedDragonNativeLock):
    def __init__(self, *args, ctx, **kwargs):
        super().__init__(*args, **kwargs, recursive=True)


class BaseImplLock(multiprocessing.synchronize.Lock):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class BaseImplRLock(multiprocessing.synchronize.RLock):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


# Condition


class DragonCondition(multiprocessing.synchronize.Condition):
    """Dragons replacement for the Multiprocessing Condition removes
    all references to _semlock and the test for assert_spawning.
    The _semlock interface is mapped onto Dragon's native Lock
    implementation.
    """

    def __getstate__(self):
        return (self._lock, self._sleeping_count, self._woken_count, self._wait_semaphore)

    def __repr__(self):
        try:
            num_waiters = self._sleeping_count.get_value() - self._woken_count.get_value()
        except Exception:
            num_waiters = "unknown"
        return "<%s(%s, %s)>" % (self.__class__.__name__, self._lock, num_waiters)


class BaseImplCondition(multiprocessing.synchronize.Condition):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


# Semaphore
class DragonSemaphore(dragon.native.semaphore.Semaphore):
    def __init__(self, *args, ctx, **kwargs):
        super().__init__(*args, **kwargs, bounded=False)


class DragonBoundedSemaphore(dragon.native.semaphore.Semaphore):
    def __init__(self, *args, ctx, **kwargs):
        super().__init__(*args, **kwargs, bounded=True)


class BaseImplSemaphore(multiprocessing.synchronize.Semaphore):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


class BaseImplBoundedSemaphore(multiprocessing.synchronize.BoundedSemaphore):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


# Event
class DragonEvent(dragon.native.event.Event):
    def __init__(self, *args, ctx, **kwargs):
        super().__init__(*args, **kwargs)


class BaseImplEvent(multiprocessing.synchronize.Event):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


# Barrier


class DragonBarrier(dragon.native.barrier.Barrier):
    def __init__(self, parties: int, action: callable = None, timeout: float = None, *, ctx=None):
        if timeout is not None:
            if timeout < 0:
                timeout = 0
        super().__init__(parties, action, timeout)

    def wait(self, timeout: float = None):
        if timeout is not None:
            if timeout < 0:
                timeout = 0
        try:
            return super().wait(timeout)
        # multiprocessing uses the threading exception
        except dragon.native.barrier.BrokenBarrierError:
            raise BrokenBarrierError


class BaseImplBarrier(multiprocessing.synchronize.Barrier):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)


# Functions


def Lock(ctx=None, use_base_impl=False):
    if use_base_impl:
        return BaseImplLock(ctx=ctx)
    else:
        return DragonLock(ctx=ctx)


def RLock(ctx=None, use_base_impl=False):
    if use_base_impl:
        return BaseImplRLock(ctx=ctx)
    else:
        return DragonRLock(ctx=ctx)


def Condition(lock, ctx=None, use_base_impl=False):
    if use_base_impl:
        return BaseImplCondition(lock, ctx=ctx)
    else:
        return DragonCondition(lock, ctx=ctx)


def Semaphore(value, ctx=None, use_base_impl=False):
    if use_base_impl:
        return BaseImplSemaphore(value, ctx=ctx)
    else:
        return DragonSemaphore(value, ctx=ctx)


def BoundedSemaphore(value, ctx=None, use_base_impl=False):
    if use_base_impl:
        return BaseImplBoundedSemaphore(value, ctx=ctx)
    else:
        return DragonBoundedSemaphore(value, ctx=ctx)


def Event(ctx=None, use_base_impl=False):
    if use_base_impl:
        return BaseImplEvent(ctx=ctx)
    else:
        return DragonEvent(ctx=ctx)


def Barrier(parties, action=None, timeout=None, ctx=None, use_base_impl=False):
    if use_base_impl:
        print("Using base barrier and we should never see this", flush=True)
        return BaseImplBarrier(parties, action=action, timeout=timeout, ctx=ctx)
    else:
        return DragonBarrier(parties, action=action, timeout=timeout, ctx=ctx)

src/dragon/mpbridge/synchronize.py
DragonHPC/dragon
"""Dragon's replacement for Multiprocessing's logging methods. 
"""

import logging

from dragon.dlogging import util as dlog

LOG = logging.getLogger(__name__)


def get_logger(*args, original=None, use_base_impl=True):
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(*args)
    else:
        return LOG  # PE-41692
        # return log_to_dragon()


def log_to_stderr(*args, level=None, original=None, use_base_impl=True):
    if use_base_impl:
        if original == None:
            raise NameError(f"Dragon patch of Multiprocessing not correct.")
        else:
            return original(*args, level=level)
    else:
        raise NotImplementedError(f"log_to_stderr is not implemented, yet.")
        return dlog.dragon_logger(0)  # Returning DragonLogger() indexed at node 0

src/dragon/mpbridge/util.py
DragonHPC/dragon
""" This is the public API of dragon.native.
"""

from .semaphore import Semaphore
from .machine import cpu_count
from .lock import Lock
from .event import Event
from .barrier import Barrier
from .process import Process, Popen
from .value import Value
from .array import Array
from .pool import Pool 

src/dragon/native/__init__.py
DragonHPC/dragon
"""The Dragon native array is an ctypes object allocated from shared memory.
The return value is a synchronized wrapper for the object, and the object itself
can be accessed via the array attribute of a Array.

.. code-block:: python

   array = dragon.native.array.Array(typecode_or_type = ctypes.c_char, size_or_initializer = [b"0", b"1"], lock=True)
"""


import sys
import logging
import ctypes
import struct

import dragon
from ..channels import Channel, Message, ChannelError
from ..managed_memory import MemoryPool
import dragon.utils as du
from ..infrastructure.facts import default_pool_muid_from_index
from ..infrastructure.parameters import this_process
from ..infrastructure.channel_desc import ChannelOptions
from ..localservices.options import ChannelOptions as LSChannelOptions
from ..globalservices.channel import create, get_refcnt, release_refcnt
from ..dtypes import WHEN_DEPOSITED
from ..native.lock import Lock
import numpy as np

LOGGER = logging.getLogger(__name__)
_DEF_MUID = default_pool_muid_from_index(this_process.index)

_TYPECODE_TO_TYPE = {
    "c": ctypes.c_char,
    "u": ctypes.c_wchar,
    "b": ctypes.c_byte,
    "B": ctypes.c_ubyte,
    "h": ctypes.c_short,
    "H": ctypes.c_ushort,
    "i": ctypes.c_int,
    "I": ctypes.c_uint,
    "l": ctypes.c_long,
    "L": ctypes.c_ulong,
    "q": ctypes.c_longlong,
    "Q": ctypes.c_ulonglong,
    "f": ctypes.c_float,
    "d": ctypes.c_double,
}

_SUPPORTED_TYPES = frozenset(_TYPECODE_TO_TYPE.values())

class Array:
    """This class implements Dragon array resides in the Dragon channel."""

    def __getstate__(self):
        return (self._channel.serialize(), self._type, self._data_len)

    def __setstate__(self, state):
        (serialized_bytes, self._type, self._data_len) = state
        self._channel = dragon.channels.Channel.attach(serialized_bytes)
        self._reset()
        get_refcnt(self._channel.cuid)

    def __repr__(self):
        return f"{self.__class__.__name__}(typecode_or_type={self._type}, lock={self._lock}, m_uid={self._muid})"

    def __init__(self, typecode_or_type: str or type, size_or_initializer: int or range or list, lock: Lock or bool = True, m_uid: int = _DEF_MUID):
        """Initialize a array object.
        :param typecode_or_type: the typecode or type is returned from the dictionary, _TYPECODE_TO_TYPE
        :type typecode_or_type: str or ctypes, required
        :param size_or_initializer: the array for the object
        :type size_or_initializer: range, int, list, required
        :param lock: dragon.native.lock.Lock, optional
        :type lock: creates lock for synchronization for array
        :param m_uid: memory pool to create the channel in and message to write value and typecode_or_type in managed memory, defaults to _DEF_MUID
        :type m_uid: int, optional
        """

        if isinstance(typecode_or_type, str):
            if typecode_or_type in _TYPECODE_TO_TYPE:
                typecode_or_type = _TYPECODE_TO_TYPE[typecode_or_type]
            else:
                raise AttributeError

        else:
            if typecode_or_type not in _SUPPORTED_TYPES:
                raise AttributeError

        LOGGER.debug(
            f"Init Dragon Native Array with {typecode_or_type=}, {size_or_initializer=}, {lock=}, {m_uid=}"
        )

        # for repr method
        self._muid = m_uid

        sh_channel_options = LSChannelOptions(capacity=1)
        gs_channel_options = ChannelOptions(ref_count=True, local_opts=sh_channel_options)
        descriptor = create(m_uid, options=gs_channel_options)
        self._channel = Channel.attach(descriptor.sdesc)
        self._reset()

        # set array with type that is multiprocessing Array type
        self._type = typecode_or_type

        #remove the assignment to self._size_or_initializer
        if isinstance(size_or_initializer, int):
            size_or_initializer = [0] * size_or_initializer # convert to list

        self._data_len = len(size_or_initializer)

        msg = Message.create_alloc(self._mpool, max(self._data_len * 8, 8))
        mview = msg.bytes_memview()
        for idx, value in enumerate(size_or_initializer):
            start = idx*8
            stop = start + 8
            mview[start:stop] = self._array2bytes(value)

        #use send handle from reset
        self._sendh.send(msg)
        msg.destroy()

        # if lock is False, return the subclass value
        if lock in (True, None):
            lock = Lock(recursive=True)

        if isinstance(lock, bool):
            pass
        elif isinstance(lock, Lock): # don't need full qualifier here
            self._lock = lock
            self.get_lock = self._get_lock
            self.get_obj = self._type
        else:
            raise AttributeError(f"The Lock must be a bool or dragon.native.lock.Lock, but is  {type(lock)}")

    def __del__(self):
        try:
            self._sendh.close()
        except (AttributeError, ChannelError):
            pass
        try:
            self._recvh.close()
        except (AttributeError, ChannelError):
            pass
        try:
            self._channel.detach()
        except (AttributeError, ChannelError):
            pass
        try:
            release_refcnt(self._channel.cuid)
        except (AttributeError, ChannelError):
            pass

    def _reset(self) -> None:
        self._recvh = self._channel.recvh()
        self._sendh = self._channel.sendh(return_mode=WHEN_DEPOSITED)
        self._recvh.open()
        self._sendh.open()
        self._mpool = MemoryPool.attach(du.B64.str_to_bytes(this_process.default_pd))

    def __len__(self):
        return self._data_len

    def __setitem__(self, item, arr):

        #grab bytes from memory pool
        msg = self._recvh.recv()

        #overwrite the values of interest that are changed in the memory pool
        if isinstance(item, slice) or isinstance(arr, int):
            mview = msg.bytes_memview()
            value_info = mview[:].tobytes()
            msg_val = [self._bytes2array(value_info[i:i+8]) for i in range(0, len(value_info), 8)]
            msg_val[item] = arr
            arr = msg_val

        msg.destroy()
        #replace old value in memory with arr
        msg = Message.create_alloc(self._mpool, max(len(arr) * 8, 8))
        mview = msg.bytes_memview()
        for idx, value in enumerate(arr):
            start = idx * 8
            stop = start + 8
            mview[start:stop] = bytes(self._array2bytes(value))

        # puts array back into memory
        self._sendh.send(msg)
        msg.destroy()

    def __getitem__(self, item):

        #grab bytes from memory
        msg = self._recvh.recv()
        mview = msg.bytes_memview()
        value_info = mview[:].tobytes()
        #divide bytes into 8 and convert the message to the original array
        return_array = [self._bytes2array(value_info[i:i+8]) for i in range(0, len(value_info), 8)]
        self._sendh.send(msg)
        msg.destroy()
        return return_array[item]

    def _array2bytes(self, arr: list or bytes or bytearray) -> Message:
        valbytes = b""
        if self._type in [ctypes.c_char, ctypes.c_wchar]:
            # This is for ctypes.c_char as ctypes.c_wchar returns str
            if isinstance(arr, (bytes, bytearray)) and self._type is not ctypes.c_wchar:
                valbytes = bytes(arr)
            else:
                valbytes = bytes(str(arr).encode("utf-8"))
            # if there are more than one bytes in valbytes
            if len(valbytes) > 8:
                raise AttributeError
        # c_float and c_double return float type
        elif self._type in [ctypes.c_float, ctypes.c_double]:
            valbytes = bytes(struct.pack("d", arr))
        # int, longlong, and other ctypes return int type
        else:
            valbytes = int(arr).to_bytes(8, byteorder=sys.byteorder, signed=True)
        return valbytes + bytes(8 - len(valbytes))

    def _bytes2array(self, msg: Message) -> list or bytes or bytearray:
        if self._type in [ctypes.c_char, ctypes.c_wchar]:
            value_info = msg.decode("utf-8")
            # removes extraneous characters
            value_val = value_info.replace("\x00", "")
            if self._type is ctypes.c_char:
                # c_char returns bytes
                return str(value_val).encode()
            else:
                # c_wchar returns string
                return str(value_val)
        elif self._type in (ctypes.c_float, ctypes.c_double):
            # returns float for the float types
            (float_val,) = struct.unpack("d", msg)
            return float(float_val)
        else:
            # returns int for int types
            return int.from_bytes(msg, byteorder=sys.byteorder, signed=True)

    def acquire(self):
        """ Acquire the internal lock object """
        return self._lock.acquire()

    def release(self):
        """ Release the internal lock object """
        return self._lock.release()

    def _get_lock(self):
        return self._lock

    def __enter__(self):
        self.acquire()

    def __exit__(self, *args):
        self.release()

src/dragon/native/array.py
DragonHPC/dragon
"""The Dragon native barrier is an object for parties number of threads. The
class can be initialized with a number of threads, an action, and timeout. The
functions are wait, reset, and abort.
"""

import sys
import logging
import os
import time

import dragon
from ..channels import Channel, ChannelBarrierReady, ChannelBarrierBroken, EventType
from ..infrastructure.facts import default_pool_muid_from_index
from ..infrastructure.parameters import this_process
from ..infrastructure.channel_desc import ChannelOptions
from ..localservices.options import ChannelOptions as ShepherdChannelOptions
from ..globalservices.channel import create, get_refcnt, release_refcnt

LOGGER = logging.getLogger(__name__)
_DEF_MUID = default_pool_muid_from_index(this_process.index)


class BrokenBarrierError(ChannelBarrierBroken):
    pass


class Barrier:
    """This class implements Barrier objects with a Dragon channel.
    The Barrier handles parties of threads that are 2 and more, number of
    waiting threads, and whether or not the barrier is broken.
    """

    def __init__(
        self, parties: int = 2, action: callable = None, timeout: float = None, m_uid: int = _DEF_MUID
    ):

        """Initialize a barrier object
        :param parties: number of parties for the barrier to act upon
        :type parties: int, optional, > 2
        :param action: class that the barrier calls upon
        :type callable: optional
        :param timeout: timeout for barrier
        :type timeout: float, optional
        :param m_uid: memory pool to create the channel in, defaults to _DEF_MUID
        :type m_uid: int, optional
        """

        LOGGER.debug(f"Init Dragon Native Barrier with m_uid={m_uid}")

        if timeout is not None and timeout < 0:
            raise ValueError(
                "The timeout in init is less than 0 for the Dragon Native Barrier implementation."
            )

        if timeout is None:
            timeout = 1000000000  # so we can do math with it.

        self._timeout = timeout
        self._action = action
        self._channel = None

        if not isinstance(parties, int):
            raise ValueError("The parties value must be an integer.")

        if parties < 1:
            raise ValueError("The number of parties must be at least two.")

        if callable(action) is False and action is not None:
            raise ValueError("The action argument must be callable or None.")

        sh_channel_options = ShepherdChannelOptions(capacity=parties)
        gs_channel_options = ChannelOptions(ref_count=True, local_opts=sh_channel_options)
        descriptor = create(m_uid, options=gs_channel_options)
        self._channel = Channel.attach(descriptor.sdesc)
        self._parties = parties

    def __del__(self):
        try:
            cuid = self._channel.cuid
            self._channel.detach()
            release_refcnt(cuid)
        except AttributeError:
            pass

    def __getstate__(self):
        return (self._channel.serialize(), self._action, self._timeout, self._parties)

    def __setstate__(self, state):
        (serialized_bytes, self._action, self._timeout, self._parties) = state
        self._channel = Channel.attach(serialized_bytes)
        get_refcnt(self._channel.cuid)

    def _wait(self, timeout: float = None) -> int:

        # Use the barrier default timeout if None is provided.
        if timeout is None:
            timeout = self._timeout
        elif timeout < 0:
            raise ValueError(
                "The timeout in wait is less than 0 for the Dragon Native Barrier implementation."
            )

        try:
            recvh = self._channel.recvh()
            recvh.open()

            # What follows is our barrier implementation over channel.poll(). Every
            # waiter registers itself by calling poll with POLLBARRIER on the
            # channel.  The last waiter will raise a ChannelBarrierReady
            # exception to execute the action callable. The other waiters will
            # return and wait in `recvh.recv(timeout=rest_timeout)`.

            try:
                # Register the waiter with the timeout.
                # This poll may not return immediately. That is, if wait() is called twice,
                # we are in the second wait() already and other processes are still picking
                # up their messages from the first wait() (WithProcessesTestBarrier.test_thousand)
                beg = time.monotonic()
                if not self._channel.poll(event_mask=EventType.POLLBARRIER, timeout=timeout):
                    raise ChannelBarrierBroken(f"Channel poll on POLLBARRIER_WAIT return -1.")

            # When the last party polls, the channel barrier ready exception is generated.
            except ChannelBarrierReady:
                try:
                    # The action is called once by the last process to poll. The
                    # action is called before any of them released.
                    if self._action is not None:
                        self._action()
                except Exception as ex:
                    raise RuntimeError(f"There was an error calling the action: {ex}") from ex

                try:  # release the waiters by writing the messages
                    self._channel.poll(event_mask=EventType.POLLBARRIER_RELEASE, timeout=0)
                except ChannelBarrierBroken as ex:
                    raise BrokenBarrierError(f"Couldn't release barrier, it is broken.")

            # the timeout needs to be adjust by how long we waited registering in POLLBARRIER further up
            end = time.monotonic()
            rest_timeout = max(0, timeout - (end - beg))

            # all but the last wait to pickup a message from receive handle
            msg = recvh.recv(timeout=rest_timeout)
            msg_val = int.from_bytes(msg.bytes_memview(), byteorder=sys.byteorder, signed=True)

            # -1 can come back if the Barrier has been broken. If this happens,
            # another process raised the BrokenBarrierException and we do not need to call abort().
            if msg_val < 0 or msg_val >= self.parties:
                raise BrokenBarrierError(f"Got a value of {msg_val} from waiting on the Barrier.")

            return msg_val

        # We have gotten a timeout either during POLLBARRIER or during the recv. This breaks the barrier
        # and raises a BrokenBarrierError,
        except TimeoutError as ex:
            try:
                self.abort()
            except:
                pass
            raise BrokenBarrierError(f"The Dragon Barrier timed out with {ex}")

        # we might attempt to wait on a broken barrier. No need to call abort(),
        # Barrier is already broken.
        except (BrokenBarrierError, ChannelBarrierBroken) as ex:
            raise BrokenBarrierError(f"This Dragon Barrier is broken.")

        # Here we have to catch any exception, because for instance, once the
        # Barrier times out, it might get destroyed which would destroy
        # the underlying managed memory for the channel and then the channel
        # attach is going to fail. So catch any error here because severe
        # errors indicate the Barrier is not functioning. Try to abort, but
        # ignore that if it fails. Provide the exception back to user so they
        # see the text anyway.
        except Exception as ex:
            try:
                self.abort()
            except:
                pass
            raise BrokenBarrierError(f"There was an unexpected error of {repr(ex)} which broke the Barrier.")

        finally:
            recvh.close()

    def wait(self, timeout: float = None) -> int:
        """When all the processes party to the barrier have called wait, they are
        all released simultaneously.  The timeout for wait takes precedence over
        constructor.
        :param timeout: timeout for barrier
        :type timeout: float, optional
        :rtype: None
        """
        LOGGER.debug(f"Barrier {self!r} wait with timeout{timeout!r}")
        return self._wait(timeout=timeout)

    def reset(self) -> None:
        """Barrier is set to empty state. Any threads waiting receive
        BrokenBarrierException.
        :rtype: None
        """
        LOGGER.debug(f"Barrier {self!r} Reset")

        try:
            self._channel.poll(event_mask=EventType.POLLRESET)
        except Exception as ex:
            raise BrokenBarrierError(f"The reset failed with the error of {repr(ex)}.")

    def abort(self) -> None:
        """The barrier is broken. Any threads that are in wait state will
        receive BrokenBarrierError.
        :rtype: None
        """
        LOGGER.debug(f"Barrier {self!r} abort")

        try:
            self._channel.poll(event_mask=EventType.POLLBARRIER_ABORT, timeout=0)
        except Exception as ex:
            raise BrokenBarrierError(f"The abort failed with the error of {repr(ex)}.")

    @property
    def parties(self):
        return self._parties

    @property
    def n_waiting(self):
        return self._channel.barrier_count

    @property
    def broken(self):
        return self._channel.broken_barrier

src/dragon/native/barrier.py
DragonHPC/dragon
"""The Dragon native event is an object that is used for signaling between threads or processes.

One thread or process signals an event and the other waits. When the waiting thread is signalled,
the waiting thread is activated.
"""

import logging

from ..channels import Channel, ChannelEmpty, ChannelFull, ChannelRecvH, ChannelSendH
from ..dtypes import WHEN_DEPOSITED
from ..globalservices.channel import release_refcnt, create, get_refcnt

from ..infrastructure.channel_desc import ChannelOptions
from ..infrastructure.parameters import this_process
from ..infrastructure.facts import default_pool_muid_from_index

LOGGER = logging.getLogger(__name__)

_DEF_MUID = default_pool_muid_from_index(this_process.index)


class Event:
    """This class implements Event objects with a Dragon channel.
    The event sends and receives bytes of the message with the respective functions set and clear.
    Wait is the function that waits for bytes to be written into the channel
    """

    _BYTEORDER = "big"

    def _cleanup(self) -> None:
        if not self._closed:
            self._closed = True
            try:
                self._reader.close()
                self._writer.close()
                self._channel.detach()
                release_refcnt(self._cuid)
            except Exception:
                pass  # all of those may fail if the channel has disappeared during shutdown

    def __del__(self):
        self._cleanup()

    def __init__(self, m_uid: int = _DEF_MUID):
        """Initialize an event object
        :param m_uid: memory pool to create the channel in, defaults to _DEF_MUID
        :type m_uid: int, optional
        """

        LOGGER.debug(f"Init Event with  m_uid={m_uid}")

        self._closed = True

        the_options = ChannelOptions(ref_count=True)
        the_options.local_opts.capacity = 1

        descr = create(m_uid=m_uid, options=the_options)
        self._channel = Channel.attach(descr.sdesc)
        get_refcnt(descr.c_uid)

        self._cuid = descr.c_uid
        self._reset()
        self._closed = False

    def _reset(self):

        self._lmsg = this_process.my_puid.to_bytes(8, byteorder=self._BYTEORDER)
        self._reader = ChannelRecvH(self._channel)
        self._writer = ChannelSendH(self._channel, return_mode=WHEN_DEPOSITED)
        self._reader.open()
        self._writer.open()

    def __getstate__(self) -> tuple:

        if self._closed:
            raise ValueError(f"Event {self!r} is closed")

        return self._channel.serialize(), self._cuid

    def __setstate__(self, state) -> None:

        self._closed = True

        sdesc, self._cuid = state

        get_refcnt(self._cuid)
        self._channel = Channel.attach(sdesc)

        self._reset()

        self._closed = False

    def __enter__(self):
        self.set()

    def __exit__(self, *args):
        self.clear()

    def set(self) -> None:
        """Set the event.

        Set writes a messsage into the channel.

        :return: None
        :rtype: NoneType
        """

        LOGGER.debug(f"Set Event {self!r}")
        try:
            self._writer.send_bytes(self._lmsg, timeout=0)

        except ChannelFull:
            pass

    def is_set(self) -> bool:
        """Returns if the event has been set. Polls the channel for the message.

        :return: Results from a channel poll with a timeout set to 0.
        :rtype: bool
        """

        LOGGER.debug(f"Is Set")
        return self._channel.poll(timeout=0)

    def wait(self, timeout: float = None) -> bool:
        """Wait for the timeout to elapse and poll for the message in the channel.

        :param timeout: how long to block in seconds, defaults to None
        :type timeout: float, optional
        :return: if the event has been set
        :rtype: bool
        """

        LOGGER.debug(f"Clear Event{self!r} and timeout{timeout!r}")
        if timeout is not None and timeout < 0:
            timeout = 0
        return self._channel.poll(timeout=timeout)

    def clear(self) -> None:
        """Clears the event once the message from the channel has been received.

        Clear receives a messsage from the channel

        :return: None
        :rtype: NoneType
        """
        LOGGER.debug(f"Clear Event{self!r}")

        try:
            _ = self._reader.recv_bytes(timeout=0)

        except ChannelEmpty:
            pass

src/dragon/native/event.py
DragonHPC/dragon
"""The Dragon native lock provides basic synchronization functionality between processes and threads.

Importantly, a `Lock()` can be released by any process or thread, while a recursive
`RLock()` needs to be released by the process & thread that called it.

Our implementation also extends the API by a few methods that are internal
in Python Multiprocessing, but are used very frequently, thus appear useful for everyone.
"""


import logging
import os
from sys import byteorder
from threading import current_thread

import dragon

from dragon.channels import Channel, ChannelRecvH, ChannelSendH, ChannelFull, ChannelSendTimeout, ChannelEmpty
from dragon.dtypes import WHEN_DEPOSITED

from ..globalservices.channel import release_refcnt, create, get_refcnt

from ..infrastructure.channel_desc import ChannelOptions
from ..infrastructure.facts import default_pool_muid_from_index
from ..infrastructure.parameters import this_process

LOGGER = logging.getLogger(__name__)

_DEF_MUID = default_pool_muid_from_index(this_process.index)


class Lock:
    """Dragon native Lock implementation based on a channel with a capacity of 1.

    Once a process or thread has acquired a lock, subsequent attempts to acquire it from any process or thread
    will block until it is released; any process may release it.

    Recursive use is supported by means of the 'recursive' parameter.  A recursive lock must be released by
    the process or thread that acquired it. Once a process or thread has acquired a recursive lock, the same
    process or thread may acquire it again without blocking; that process or thread must release it once for
    each time it has been acquired.

    Lock supports the context manager protocol and thus may be used in 'with' statements.
    """

    def _close(self) -> None:
        if not self._closed:
            self._closed = True
            try:
                self._reader.close()
                self._writer.close()
                self._channel.detach()
                release_refcnt(self._cuid)
            except Exception:
                pass  # all of those may fail if the channel has disappeared during shutdown

    def __del__(self):
        self._close()

    def __init__(self, m_uid: int = _DEF_MUID, recursive: bool = False):
        """Initialize the lock object

        :param m_uid: memory pool to create the channel in, defaults to _DEF_MUID
        :type m_uid: int, optional
        :param recursive: if the Lock should be recursive, defaults to False
        :type recursive: bool, optional
        """

        self._closed = True

        self._recursive = recursive

        the_options = ChannelOptions(ref_count=True)
        the_options.local_opts.capacity = 1

        descr = create(m_uid=m_uid, options=the_options)
        self._channel = Channel.attach(descr.sdesc)
        get_refcnt(descr.c_uid)

        self._cuid = descr.c_uid
        self._reset()
        self._closed = False

    def _reset(self):

        self._lmsg = this_process.my_puid.to_bytes(8, byteorder)
        self._reader = ChannelRecvH(self._channel)
        self._writer = ChannelSendH(self._channel, return_mode=WHEN_DEPOSITED)
        self._reader.open()
        self._writer.open()
        self._accesscount = 0

    def __getstate__(self) -> tuple:

        if self._closed:
            raise ValueError(f"Lock {self!r} is closed")

        puid = this_process.my_puid
        tid = current_thread

        return self._channel.serialize(), self._recursive, self._cuid, self._accesscount, puid, tid

    def __setstate__(self, state) -> None:

        self._closed = True

        sdesc, self._recursive, self._cuid, accesscount, puid, tid = state

        get_refcnt(self._cuid)
        self._channel = Channel.attach(sdesc)

        self._reset()

        # if sent to myself, set accesscount correctly
        if this_process.my_puid == puid:
            if current_thread == tid:
                self._accesscount = accesscount

        self._closed = False

    def acquire(self, block: bool = True, timeout: float = None) -> bool:
        """Acquire a lock, blocking or non-blocking.

        If the lock is not recursive:

        With the block argument set to True (the default), the method call will block until the lock is in an
        unlocked state, then set it to locked and return True.

        With the block argument set to False, the method call does not block.  If the lock is currently in a
        locked state, return False; otherwise set the lock to a locked state and return True.  When invoked
        with a positive, floating-point value for timeout, block for at most the number of seconds specified
        by timeout as long as the lock can not be acquired. Invocations with a negative value for timeout are
        equivalent to a timeout of zero. Invocations with a timeout value of None (the default) set the
        timeout period to infinite.  The timeout argument has no practical implications if the block argument
        is set to False and is thus ignored. Returns True if the lock has been acquired or False if the
        timeout period has elapsed.

        If the lock is recursive, repeated calls to a lock already owned by the current process return True
        and increment the recursion level. `Release()` has then be called an equal amount of times by the same
        process to free the lock.

        :param block: should the call block at all, defaults to True
        :type block: bool, optional
        :param timeout: how long to block in seconds, ignored if block=False, defaults to None
        :type timeout: float, optional
        :return: if the lock was acquired
        :rtype: bool
        """

        LOGGER.debug(f"Acquire Lock {self!r} with blocking={block}, timeout={timeout}")

        if timeout is not None and timeout < 0:
            timeout = 0

        if self._recursive:
            if self._accesscount > 0:
                self._accesscount += 1
                return True

        # NOTE: If we had an rwlock() on the channel, we might just that and make the lock
        # even lighter.

        try:
            self._writer.send_bytes(self._lmsg, timeout=timeout, blocking=block)

            self._accesscount = 1
            return True
        except (ChannelFull, ChannelSendTimeout) as e:
            return False

    def release(self) -> None:
        """Release a lock.

        If the lock is not recursive: This can be called from any process or thread,
        not only the process or thread which originally acquired the lock.

        If the lock is recursive: Decrement the recursion level. If after the decrement the recursion level is
        zero, reset the lock to unlocked (not owned by any process or thread) and if any other processes or
        threads are blocked waiting for the lock to become unlocked, allow exactly one of them to proceed. If
        after the decrement the recursion level is still nonzero, the lock remains locked and owned by the
        calling dragon process and thread.

        :raises ValueError: if the lock has not been acquired before
        :raises AssertionError: if the lock is recursive and not held by the caller
        """

        LOGGER.debug(f"Release Lock {self!r}")

        if self._recursive:
            assert self._accesscount > 0, "attempt to release recursive lock not owned by process"
            self._accesscount -= 1
            if self._accesscount > 0:
                return

        # this will remove any message, thus unblocking the next process to acquire
        try:
            _ = self._reader.recv_bytes(timeout=0)
            self._accesscount = 0
        except ChannelEmpty as e:
            raise ValueError("Cannot release unlocked lock.")

    # API extention
    # NOTE: These methods are private methods on _Semlock in Python Multiprocessing,
    # but used very frequently internally. They are very useful, so I think we want
    # them for our users and I am pretty sure packages on top of Multiprocessing
    # are using them too.

    def is_recursive(self) -> bool:
        """Return True if this lock is recursive, i.e. can be called multiple times by the
        same thread/process.

        :return: recursiveness of the lock
        :rtype: bool
        """
        return self._recursive

    def is_mine(self) -> bool:
        """Return True if this rlock is owned by the caller process & thread

        :return: ownership of the lock
        :rtype: bool
        """
        return self._accesscount > 0

    def count(self) -> int:
        """How often this lock has been acquired.
        Will return only 0 or 1 for a non-recursive Lock.
        :return: The number of times this lock has been acquired
        :rtype: int
        """
        return self._accesscount

    def __enter__(self):
        self.acquire()

    def __exit__(self, *args):
        self.release()


# Extensions.  TODO, should we plug into the mp.util.Finalize stuff?

src/dragon/native/lock.py
DragonHPC/dragon
""" The Dragon native classes to log messages from the runtime.

This is a stub.

"""

import os
from logging import Handler, Logger, LogRecord, getLogger
from dragon.dlogging.logger import DragonLogger


def log_to_stderr(level=None) -> None:
    raise NotImplementedError(f"{__name__} is not implemented, yet.")


class DragonHandler(Handler):
    """Logging handler to enable python logging to be sent to the Dragon Logging channel."""

    def __init__(self, serialized_log_descr: bytes) -> None:
        Handler.__init__(self)
        self._dlog = DragonLogger.attach(serialized_log_descr)

    def emit(self, record: LogRecord) -> None:
        """Emit a Dragon log record

        :param record: the record to emit
        :type record: LogRecord
        """

        try:
            msg = self.format(record)
            self._dlog.put(msg, record.levelno)
        except Exception:
            self.handleError(record)


def log_to_dragon(serialized_log_descr: bytes, level: int = None) -> Logger:
    """Turn on logging and add a handler which routes log messages to Dragon logging channels

    :param serialized_log_descr: The serialized descriptor
    :type serialized_log_descr: bytes
    :param level: the log level, defaults to None
    :type level: int, optional
    :return: the Logger
    :rtype: Logger
    """

    logger = getLogger()
    handler = DragonHandler(serialized_log_descr)
    logger.addHandler(handler)

    if level:
        logger.setLevel(level)
    return logger

src/dragon/native/logging.py
DragonHPC/dragon
""" The Dragon native machine interface provides hardware related functionality.

These calls return various hardware properties of the system(s), like the number of processors
or a list of all nodes.
"""

import logging

from ..globalservices.node import query, query_total_cpus, get_list
from ..infrastructure.parameters import this_process
from ..utils import host_id

LOG = logging.getLogger(__name__)

# TODO: Decide on a model for system architecture that generalizes
# well to HPC and cloud environments alike


def cpu_count() -> int:
    """Return the number of CPUs in this system.

    :return: integer representing the number of cpus
    """

    return query_total_cpus()


class Node:
    def __init__(self, ident: str or int):
        """A class abstracting the current hardware node.

        :param ident: name or h_uid of the node to request.
        :type ident: str or int
        """

        self._descr = None
        self._update_descriptor(ident)

    @property
    def name(self) -> str:
        """Return the Dragon name of this node.
        Use Node.hostname for the hostname.

        :return: The name of the node
        :rtype: str
        """
        return self._descr.name

    @property
    def h_uid(self) -> int:
        """The unique global identifier of this node.

        :return: the h_uid
        :rtype: in
        """

        return self._descr.h_uid

    @property
    def ip(self) -> str:
        """Return the IPv4 address of this node

        :return: The addres
        :rtype: str
        """
        return self._descr.ip_addrs

    @property
    def num_cpus(self) -> int:
        """Return the number of CPUs of this node.

        :return: The number of CPUs
        :rtype: int
        """
        return self._descr.num_cpus

    @property
    def physical_mem(self) -> int:
        """Return the physical memory of this node.

        :return: Physical memory in bytes
        :rtype: int
        """
        return self._descr.physical_mem

    @property
    def hostname(self) -> str:
        """Host name of the node in the network.

        This is not the Dragon internal name (Node.name) for this node.

        :return: The hostname
        :rtype: str
        """
        return self._descr.host_name

    def _update_descriptor(self, ident=None):
        if ident != None:
            self._descr = query(ident)
        else:
            self._descr = query(self._descr.h_uid)


def current() -> Node:
    """Return a node object for this process.

    :return: Node object
    :rtype: Node
    """
    h_uid = host_id()
    return Node(h_uid)

class System:

    def __init__(self):
      """A stub of a system abstraction"""
      self.nodes = get_list()
    
    def nodes(self): 
        return self.nodes
    
    def nnodes(self) -> int:
        return len(self.nodes)
src/dragon/native/machine.py
DragonHPC/dragon
"""The Dragon native pool manages a pool of child processes. 
"""

from __future__ import annotations
import logging
import queue
import socket
import itertools
import traceback
import threading
import types
import signal
from typing import Iterable, Any
import time

import dragon

from .queue import Queue
from .process import TemplateProcess
from .process_group import ProcessGroup
from .event import Event
from .process import current as current_process
from .machine import System

from ..infrastructure.policy import Policy

# helper functions copied from multiprocessing.pool.Pool

job_counter = itertools.count()

POLL_FREQUENCY = 0.2


class RemoteTraceback(Exception):
    def __init__(self, tb):
        self.tb = tb

    def __str__(self):
        return self.tb


class ExceptionWithTraceback:
    def __init__(self, exc, tb):
        tb = traceback.format_exception(type(exc), exc, tb)
        tb = "".join(tb)
        self.exc = exc
        self.tb = '\n"""\n%s"""' % tb

    def __reduce__(self):
        return rebuild_exc, (self.exc, self.tb)


def rebuild_exc(exc, tb):
    exc.__cause__ = RemoteTraceback(tb)
    return exc


class MaybeEncodingError(Exception):
    """Wraps possible unpickleable errors, so they can be
    safely sent through the socket.
    """

    def __init__(self, exc, value):
        self.exc = repr(exc)
        self.value = repr(value)
        super(MaybeEncodingError, self).__init__(self.exc, self.value)

    def __str__(self):
        return "Error sending result: '%s'. Reason: '%s'" % (self.value, self.exc)

    def __repr__(self):
        return "<%s: %s>" % (self.__class__.__name__, self)


def _helper_reraises_exception(ex):
    "Pickle-able helper function for use by _guarded_task_generation."
    raise ex


class ApplyResult:
    """Returned by `apply_async` and has `get` method to retrieve result."""

    def __init__(self, pool: Pool, callback: callable, error_callback: callable):
        """Initializes a result

        :param pool: the pool where work has been submitted
        :type pool: dragon.native.pool.Pool
        :param callback: function called on returned result
        :type callback: callable
        :param error_callback: function called on returned result if an error was raised
        :type error_callback: callable
        """
        self._pool = pool
        self._event = threading.Event()
        self._job = next(job_counter)
        self._cache = pool._cache
        self._callback = callback
        self._error_callback = error_callback
        self._cache[self._job] = self
        self._success = None
        self._value = None

    def ready(self) -> bool:
        """Checks whether the result is ready.

        :return: returns True if the result is ready or False if it isn't
        :rtype: bool
        """
        return self._event.is_set()

    def successful(self) -> bool:
        """Checks if the result is ready and returns whether the function call was successful.

        :raises ValueError: raised if result is not ready
        :return: returns True if function call was successful
        :rtype: bool
        """
        if not self.ready():
            raise ValueError("{0!r} not ready".format(self))
        return self._success

    def wait(self, timeout: float = None) -> None:
        """Waits on event for result to be ready

        :param timeout: timeout indicating how long to wait for result to be ready, defaults to None
        :type timeout: float, optional
        """
        self._event.wait(timeout)

    def get(self, timeout: float = None) -> Any:
        """Retrieves returned values of work submitted to pool

        :param timeout: timeout for getting result, defaults to None
        :type timeout: float, optional
        :raises TimeoutError: raised if result is not ready in specified timeout
        :raises self._value: raises error if returned work was unsuccessful
        :return: value returned by `func(*args, **kwargs)`
        :rtype: Any
        """
        self.wait(timeout)
        if not self.ready():
            raise TimeoutError
        if self._success:
            return self._value
        else:
            raise self._value

    def _set(self, i, obj):
        self._success, self._value = obj
        if self._callback and self._success:
            self._callback(self._value)
        if self._error_callback and not self._success:
            self._error_callback(self._value)
        self._event.set()
        del self._cache[self._job]
        self._pool = None

    __class_getitem__ = classmethod(types.GenericAlias)


AsyncResult = ApplyResult  # create alias -- see #17805

#
# Class whose instances are returned by `Pool.map_async()`
#


class MapResult(ApplyResult):
    def __init__(self, pool: Pool, chunksize: int, length: int, callback: callable, error_callback: callable):
        """Initialization method

        :param pool: the pool where work is submitted
        :type pool: dragon.native.pool.Pool
        :param chunksize: size that work is supposed to be broken up into
        :type chunksize: int
        :param length: number of items in iterable
        :type length: int
        :param callback: function to be called upon return of result
        :type callback: callable
        :param error_callback: function to be called upon return of result if an error was raised.
        :type error_callback: callable
        """

        ApplyResult.__init__(self, pool, callback, error_callback=error_callback)
        self._success = True
        self._value = [None] * length
        self._chunksize = chunksize
        if chunksize <= 0:
            self._number_left = 0
            self._event.set()
            del self._cache[self._job]
        else:
            self._number_left = length // chunksize + bool(length % chunksize)

    def _set(self, i, success_result):
        self._number_left -= 1
        success, result = success_result
        if success and self._success:
            self._value[i * self._chunksize : (i + 1) * self._chunksize] = result
            if self._number_left == 0:
                if self._callback:
                    self._callback(self._value)
                del self._cache[self._job]
                self._event.set()
                self._pool = None
        else:
            if not success and self._success:
                # only store first exception
                self._success = False
                self._value = result
            if self._number_left == 0:
                # only consider the result ready once all jobs are done
                if self._error_callback:
                    self._error_callback(self._value)
                del self._cache[self._job]
                self._event.set()
                self._pool = None


def mapstar(args):
    return list(map(*args))


LOGGER = logging.getLogger(__name__)


class Pool:
    """A Dragon native Pool relying on native Process Group and Queues

    The interface resembles the Python Multiprocessing.Pool interface and focuses on the most generic functionality. Here we directly support the asynchronous API. The synchronous version of these calls are supported by calling get on the objects returned from the asynchronous functions. By using a Dragon native Process Group to coordinate the worker processes this implementation addresses scalability limitations with the patched base implementation of Multiprocessing.Pool.

    At this time, both `terminate` and `close` send a `signal.SIGTERM` that the workers catch. Using `close` guarantees that all work submitted to the pool is finished before the signal is sent while `terminate` sends the signal immediately. The user is expected to call `join` following both of these calls. If `join` is not called, zombie processes may be left.

    """

    def __init__(
        self,
        processes: int = None,
        initializer: callable = None,
        initargs: tuple = (),
        maxtasksperchild: int = None,
        policy: Policy = None,
    ):
        """Init method

        :param processes: number of worker processes, defaults to None
        :type processes: int, optional
        :param initializer: initializer function, defaults to None
        :type initializer: callable, optional
        :param initargs: arguments for initializer function, defaults to ()
        :type initargs: tuple, optional
        :param maxtasksperchild: maximum tasks each worker will perform, defaults to None
        :type maxtasksperchild: int, optional
        :param policy: determines the placement of the processes
        :type policy: dragon.infrastructure.policy.Policy
        :raises ValueError: raised if number of worker processes is less than 1
        """
        myp = current_process()
        LOGGER.debug(f"pool init on node {socket.gethostname()} by process {myp.ident}")
        self._inqueue = Queue()
        self._outqueue = Queue()
        self._end_threading_event = threading.Event()
        self._maxtasksperchild = maxtasksperchild
        self._initializer = initializer
        self._initargs = initargs
        self._can_join = False

        if processes is None:
            my_sys = System()
            self._processes = my_sys.nnodes
            LOGGER.warning(
                f"Number of processes not specified. Setting processes={self._processes}, equal to the number of nodes."
            )
        else:
            self._processes = processes

        if self._processes < 1:
            raise ValueError("Number of processes must be at least 1")

        # starts a process group with nproc workers
        self._template = TemplateProcess(
            self._worker_function,
            args=(self._inqueue, self._outqueue, self._initializer, self._initargs, self._maxtasksperchild),
        )
        self._pg = ProcessGroup(restart=True, ignore_error_on_exit=True)
        self._pg.add_process(self._processes, self._template)
        self._pg.init()
        self._pg.start()

        # dict that holds jobs submitted via apply_async and map_async
        self._cache = {}

        # thread that handles getting results from outqueue and putting in dict
        self._results_handler = threading.Thread(
            target=self._handle_results, args=(self._outqueue, self._cache, self._end_threading_event)
        )
        self._results_handler.start()

        # thread used by map_async to chunk input list
        self._map_launch_thread = None
        # thread used by close to wait for results before sending shutdown signals
        self._close_thread = None

    def _check_running(self):
        status = self._pg.status
        if not (status == "Running" or status == "Maintain"):
            raise ValueError("Pool not running")

    def terminate(self) -> None:
        """This sends the signal.SIGTERM and sets the threading event immediately. Calling this method before blocking on results may lead to some work being undone. If all work should be done before stopping the workers, `close` should be used. If `join` is not called after terminate zombie processes may be left over and interfere with future pool or process group use."""

        LOGGER.debug(f"terminating pool")
        # setting end event to stop all threads
        self._end_threading_event.set()
        # send SIGTERM signal to process group. Worker functions will return on their next pass through the while loop. See CIRRUS-1467 for reasons why sending SIGKILL doesn't work here.
        if not self._pg.status == "Stop":
            self._pg.kill(signal.SIGTERM)
        self._can_join = True

    @staticmethod
    def _close(cache, end_threading_event, pool):
        def work_done(cache):
            return len(cache) == 0

        LOGGER.debug(f"closing pool and setting end events in thread {threading.get_native_id()}")

        # waits until all submitted jobs are done.
        while not work_done(cache):
            if end_threading_event.is_set():
                return

        # send SIGTERM to worker_functions
        if not pool.status == "Stop":
            pool.kill(signal.SIGTERM)
        end_threading_event.set()

    def close(self) -> None:
        """This method starts a thread that waits for all submitted jobs to finish. This thread then sends signal.SIGTERM to the process group and sets the threading event to close the handle_results thread. Waiting and then sending the signals within the thread allows close to return without all work needing to be done.

        :raises ValueError: raised if `close` or `terminate` have been previously called
        """
        if self._can_join:
            raise ValueError("Trying to close pool that has already been closed or terminated")
        # defines and starts thread that waits for jobs to be done before sending shutdown signals
        self._close_thread = threading.Thread(
            target=self._close, args=(self._cache, self._end_threading_event, self._pg)
        )
        self._close_thread.start()
        self._can_join = True

    def join(self) -> None:
        """Waits for all workers to return. The user must have called close or terminate prior to calling join.

        :raises ValueError: raised if `close` or `terminate` were not previously called
        """
        LOGGER.debug(f"joining pool")
        if not self._can_join:
            raise ValueError("Trying to join pool that has not been closed or terminated")
        # waits on close_thread before joining workers
        if self._close_thread is not None:
            self._close_thread.join()
        if self._map_launch_thread is not None:
            self._map_launch_thread.join()

        if not self._pg.status == "Stop":
            self._pg.join()
            self._pg.stop()
        self._results_handler.join()

    @staticmethod
    def _worker_function(inqueue, outqueue, initializer=None, initargs=(), maxtasks=None):
        myp = current_process()

        # handles shutdown signal
        termflag = False

        def handler(signum, frame):
            LOGGER.debug("_worker_function SIGTERM handler saw signal")
            nonlocal termflag
            termflag = True

        signal.signal(signal.SIGTERM, handler)

        if (maxtasks is not None) and not (isinstance(maxtasks, int) and maxtasks >= 1):
            raise AssertionError("Maxtasks {!r} is not valid".format(maxtasks))

        if initializer is not None:
            initializer(*initargs)

        completed_tasks = 0

        while not termflag and (maxtasks is None or (maxtasks and completed_tasks < maxtasks)):
            # get work item
            LOGGER.debug(f"getting work from inqueue on node {socket.gethostname()}")
            try:
                # still need timeout here to check signal occasionally
                task = inqueue.get(timeout=0)
            except queue.Empty:
                time.sleep(POLL_FREQUENCY)
                continue
            except TimeoutError:
                time.sleep(POLL_FREQUENCY)
                continue

            LOGGER.debug(f"{task} task received by {myp.ident} on {socket.gethostname()}")

            job, i, func, args, kwargs = task
            LOGGER.debug(f"job={job}, i={i}, func={func}, args={args}, kwargs={kwargs}")

            # perform work
            try:
                output = (True, func(*args, **kwargs))
            except Exception as e:
                LOGGER.debug(f"\nDragon pool._worker_function threw an exception: \n{e}")
                if func is not _helper_reraises_exception:
                    e = ExceptionWithTraceback(e, e.__traceback__)
                output = (False, e)
            # put result in output queue
            LOGGER.debug(f"putting result in outqueue = {(job, i, output)}")
            try:
                outqueue.put((job, i, output))
            except Exception as e:
                wrapped = MaybeEncodingError(e, output[1])
                LOGGER.debug("Possible encoding error while sending result: %s" % (wrapped))
                outqueue.put((job, i, (False, wrapped)))
            completed_tasks += 1

        LOGGER.debug(f"{myp.ident} returning from worker_function")
        

    @classmethod
    def _handle_results(cls, outqueue, cache, end_event):
        LOGGER.debug(
            f"handle_results on node {socket.gethostname()} with thread id {threading.get_native_id()}"
        )
        while not end_event.is_set():

            # timeout with some frequency so we check if end_event is set
            try:
                task = outqueue.get(timeout=POLL_FREQUENCY)
            except queue.Empty:
                continue
            except TimeoutError:
                continue

            job, i, obj = task

            # place result into dict holding different jobs
            try:
                cache[job]._set(i, obj)
            except KeyError:
                pass

            task = job = obj = None

    def apply_async(
        self,
        func: callable,
        args: tuple = (),
        kwargs: dict = {},
        callback: callable = None,
        error_callback: callable = None,
    ) -> ApplyResult:
        """Equivalent to calling `func(*args, **kwargs)` in a non-blocking way. A result is immediately returned and then updated when `func(*args, **kwargs)` has completed.

        :param func: user function to be called by worker function
        :type func: callable
        :param args: input args to func, defaults to ()
        :type args: tuple, optional
        :param kwargs: input kwargs to func, defaults to {}
        :type kwargs: dict, optional
        :param callback: user provided callback function, defaults to None
        :type callback: callable, optional
        :param error_callback: user provided error callback function, defaults to None
        :type error_callback: callable, optional
        :raises ValueError: raised if pool has already been closed or terminated
        :return: A result that has a `get` method to retrieve result of `func(*args, **kwargs)`
        :rtype: ApplyResult
        """

        if self._can_join:
            raise ValueError("Pool closed or terminated. Cannot add work to a closed pool.")

        result = ApplyResult(self, callback, error_callback)
        self._inqueue.put((result._job, 0, func, args, kwargs))
        return result

    @staticmethod
    def _get_tasks(func, it, size):
        it = iter(it)
        while True:
            x = tuple(itertools.islice(it, size))
            if not x:
                return
            yield (func, x)

    @staticmethod
    def _guarded_task_generation(result_job, func, iterable):
        try:
            i = -1
            for i, x in enumerate(iterable):
                yield (result_job, i, func, (x,), {})
        except Exception as e:
            yield (result_job, i + 1, _helper_reraises_exception, (e,), {})

    @staticmethod
    def _partition_tasks(inqueue, result, mapper, task_batches, end_event, set_length=None):
        # this is here in case the iterable is empty that way the set_length
        # can be computed correctly
        task = None

        try:
            for task in Pool._guarded_task_generation(result._job, mapper, task_batches):
                if end_event.is_set():
                    return
                inqueue.put(task)
            # only used by imap and imap_unordered
            else:
                if set_length:
                    idx = task[1] if task else -1
                    set_length(idx + 1)
        finally:
            task = taskseq = job = None

    def _map_async(self, func, iterable, mapper, chunksize=None, callback=None, error_callback=None):
        if self._can_join:
            raise ValueError("Pool closed or terminated. Cannot map to a closed pool.")
        self._check_running()

        if not hasattr(iterable, "__len__"):
            iterable = list(iterable)

        num_items = len(iterable)

        if num_items == 0:
            chunksize = 0

        if chunksize is None:
            chunksize, remainder = divmod(num_items, self._processes * 4)
            if remainder:
                chunksize += 1

        # Generate the tasks and place result reference in dict
        task_batches = Pool._get_tasks(func, iterable, chunksize)
        result = MapResult(self, chunksize, len(iterable), callback, error_callback=error_callback)

        # wait to generate more tasks till last task generation thread has finished
        if self._map_launch_thread is not None:
            self._map_launch_thread.join()
        # start thread to put tasks into the input queue
        self._map_launch_thread = threading.Thread(
            target=Pool._partition_tasks,
            args=(self._inqueue, result, mapper, task_batches, self._end_threading_event),
        )
        self._map_launch_thread.start()

        return result

    def map_async(
        self,
        func: callable,
        iterable: Iterable,
        chunksize: int = None,
        callback: callable = None,
        error_callback: callable = None,
    ) -> MapResult:
        """Apply `func` to each element in `iterable`. The results
        are collected in a list that is returned immediately. It is equivalent to map if the `get` method is called immediately on the returned result.

        :param func: user provided function to call on elements of iterable
        :type func: callable
        :param iterable: input args to func
        :type iterable: iterable
        :param chunksize: size of work elements to be submitted to input work queue, defaults to None
        :type chunksize: int, optional
        :param callback: user provided callback function, defaults to None
        :type callback: callable, optional
        :param error_callback: user provided error callback function, defaults to None
        :type error_callback: callable, optional
        :return: A result that has a `get` method that returns an iterable of the output from applying `func` to each element of iterable.
        :rtype: MapResult
        """
        return self._map_async(func, iterable, mapstar, chunksize, callback, error_callback)

src/dragon/native/pool.py
DragonHPC/dragon
""" The Dragon native process object provides process management across one or
multiple distributed systems. `TemplateProcess` can hold a blueprint for a process
that can be used to generate many similar processes.
"""

import os
import sys
import logging
import signal
import cloudpickle
import threading
from shutil import which
from dataclasses import dataclass


from ..globalservices.process import (
    ProcessError,
    create as process_create,
    query as process_query,
    join as process_join,
    kill as process_kill,
    create_with_argdata as process_create_with_argdata,
)

from ..infrastructure.parameters import this_process
from ..infrastructure.facts import ARG_IMMEDIATE_LIMIT
from ..infrastructure.process_desc import ProcessOptions
from ..infrastructure.messages import PIPE as MSG_PIPE, STDOUT as MSG_STDOUT, DEVNULL as MSG_DEVNULL

LOG = logging.getLogger(__file__)

class Popen:
    PIPE = MSG_PIPE
    STDOUT = MSG_STDOUT
    DEVNULL = MSG_DEVNULL

    def __init__(
        self,
        args,
        executable=None,
        stdin=None,
        stdout=None,
        stderr=None,
        preexec_fn=None,
        shell=False,
        cwd=None,
        env=None,
        group=None,
        user=None,
        proc_name=None,
        umask=-1,
        pipesize=-1,
    ):
        """
        Emulate the subprocess.Popen interface using Dragon infrastructure.

        The DragonPopen constructor mirrors the interface of the subprocess.Popen
        class. Not every field is used, but is provided for signature
        compatability. Unsupported options may be ignored. See subprocess.Popen
        description for a description of arguments. Here are notable differences.

        When PIPE is specified, stdout and stderr will be a Connection. When STDOUT
        is specified for stderr, they will both be set to the same Connection object.
        If DEVNULL is specified for stdout or stderr, then the output is thrown away.

        When bufpool is provided it will be used for any allocations of connections for
        this process.

        :param args: _description_
        :type args: _type_
        :param executable: _description_, defaults to None
        :type executable: _type_, optional
        :param stdin: _description_, defaults to None
        :type stdin: _type_, optional
        :param stdout: _description_, defaults to None
        :type stdout: _type_, optional
        :param stderr: _description_, defaults to None
        :type stderr: _type_, optional
        :param preexec_fn: _description_, defaults to None
        :type preexec_fn: _type_, optional
        :param shell: _description_, defaults to False
        :type shell: bool, optional
        :param cwd: _description_, defaults to None
        :type cwd: _type_, optional
        :param env: _description_, defaults to None
        :type env: _type_, optional
        :param group: _description_, defaults to None
        :type group: _type_, optional
        :param user: _description_, defaults to None
        :type user: _type_, optional
        :param proc_name: _description_, defaults to None
        :type proc_name: _type_, optional
        :param umask: _description_, defaults to -1
        :type umask: int, optional
        :param pipesize: _description_, defaults to -1
        :type pipesize: int, optional
        :raises ex: _description_
        """

        self._args = args
        self._executable = executable
        self._stdin_requested = stdin
        self._stdout_requested = stdout
        self._stderr_requested = stderr
        self._preexec_fn = preexec_fn
        self._shell = shell
        self._cwd = cwd
        self._env = env
        self._group = group
        self._user = user
        self._proc_name = proc_name
        self._umask = umask
        self._pipesize = pipesize
        self._puid = None
        self._pid = None
        self._theproc = None

        if cwd is None or cwd == ".":
            cwd = os.getcwd()

        if not os.path.isfile(executable):
            executable = cwd + "/" + executable

        if executable.strip().endswith(".py"):
            args = [executable] + args
            executable = sys.executable

        options = ProcessOptions(make_inf_channels=True)

        # Create the manager owner process
        self._theproc = process_create(
            exe=executable,
            run_dir=cwd,
            args=args,
            env=env,
            options=options,
            stdin=stdin,
            stdout=stdout,
            stderr=stderr,
            group=group,
            user=user,
            umask=umask,
            pipesize=pipesize,
            user_name=proc_name,
        )

        self._stdin = self._theproc.stdin_conn
        self._stdout = self._theproc.stdout_conn
        self._stderr = self._theproc.stderr_conn

    def __del__(self):
        try:
            if self._stdin is not None:
                self._stdin.close()
        except:
            pass

        try:
            if self._stdout is not None:
                self._stdout.close()
        except:
            pass

        try:
            if self._stderr is not None:
                self._stderr.close()
        except:
            pass

    def poll(self):
        """Check if process has terminated"""
        raise NotImplementedError

    def wait(self, timeout=None):
        """wait for process to terminate"""
        raise NotImplementedError

    def communicate(self, input=None, timeout=None):
        """receive all output and transmit all input and wait for process to terminate."""
        raise NotImplementedError

    def send_signal(self, signal):
        """just like it sounds"""
        raise NotImplementedError

    def terminate(self):
        """Send SigTERM"""
        raise NotImplementedError

    def kill(self):
        """Send SigKILL"""
        raise NotImplementedError

    @property
    def args(self):
        """actual args passed to process"""
        raise NotImplementedError

    @property
    def stdin(self):
        """the stdin connection when PIPE was specified. None otherwise."""
        if self._theproc:
            return self._stdin

    @property
    def stdout(self):
        """The stdout connection when PIPE was specified. None otherwise."""
        if self._theproc:
            return self._stdout

    @property
    def stderr(self):
        """The stderr connection when PIPE was specified. None otherwise."""
        if self._theproc:
            return self._stderr

    @property
    def puid(self):
        """Dragon process puid. Globally unique"""
        if self._theproc:
            return self._theproc.p_uid

    @property
    def pid(self):
        """OS process pid. Node unique"""
        raise NotImplementedError

    @property
    def returncode(self):
        """When it has terminated, exit code. None otherwise."""
        raise NotImplementedError


class TemplateProcess:
    """This class provides a template for a Dragon process."""

    def __init__(
        self,
        target: str or callable,
        args: tuple = None,
        kwargs: dict = None,
        cwd: str = ".",
        env: dict = None,
        stdin: int = None,
        stdout: int = None,
        stderr: int = None,
        # policy: Policy = DefaultPolicy,
    ):
        """Generic Dragon process template object defining a process based on a
        binary executable or a Python callable.

        If the target is a callable, define the Python process executing the
        callable.
        If the target is a string, assume it's the filename of a binary to be
        started.  If the target file is not found, the object tries to resolve
        the command via the PATH environment variable. If that fails, it assumes
        the file is in the current working directory.

        This template cannot be started, but it can be used to create many
        similar processes like this:

        ```Python
            t = TemplateProcess(myfunc)
            p1 = Process.from_template(t, ident="NewProcess1")
            p2 = Process.from_template(t, ident="NewProcess2")
        ```

        The template stores the absolute path of the target if the target is an executable.
        If the target is a Python function, the template stores the Python executable in
        its `target` attribute, the command line parameters in its `args` attribute and
        the Python data in its `argdata` attribute. To obtain the original attributes in
        Python use

        ```Python
            target, args, kwargs = p.get_original_python_parameters()
        ```

        :param target: The binary or Python callable to execute.
        :type target: str or callable
        :param args: arguments to pass to the binary or callable, defaults to None
        :type args: tuple, optional
        :param kwargs: keyword arguments to pass to the Python callable, defaults to None
        :type kwargs: tuple, optional
        :param cwd: current working directory, defaults to "."
        :type cwd: str, optional
        :param env: environment variables to pass to the process environment, defaults to None
        :type env: dict, optional
        :param stdin: Stdin file handling. Valid values are PIPE and None.
        :type stdin: int, optional
        :param stdout: Stdout file handling. Valid values are PIPE, STDOUT and None.
        :type stdout: int, optional
        :param stderr: Stderr file handling. Valid values are PIPE, STDOUT and None.
        :type stderr: int, optional
        """

        self.is_python = callable(target)

        # store only the modified targets. We want to be able to pickle the
        # TemplateProcess. The user has to use a method to get the original arguments.

        if self.is_python:
            self.target, self.args, self.argdata = self._get_python_process_parameters(target, args, kwargs)
        else:  # binary
            self.args = args
            self.kwargs = None
            self.target = self._find_target(target, cwd)

        self.cwd = cwd
        self.env = env

        self.stdout = stdout
        self.stderr = stderr
        self.stdin = stdin

    @staticmethod
    def _find_target(target, cwd) -> str:

        new_target = target

        if not os.path.isfile(target):
            if which(target):
                new_target = which(target)
            else:
                new_target = os.path.join(cwd, target)

                if not os.path.isfile(new_target):
                    raise AttributeError(
                        f"Cannot find executable {target} directly, in PATH or in working directory {cwd}"
                    )

        return new_target

    def _get_python_process_parameters(self, target, args, kwargs) -> tuple:

        new_target = sys.executable
        new_args = [
            "-c",
            "from dragon.native.process import _dragon_native_python_process_main; _dragon_native_python_process_main()",
        ]

        argdata = cloudpickle.dumps((target, args or (), kwargs or {}))

        return new_target, new_args, argdata

    def get_original_python_parameters(self) -> tuple[callable, tuple, dict]:
        """Return the original parameters of a Python process from the object.

        :return: function, args and kwargs of the templated process
        :rtype: tuple[callable, tuple, dict]
        """
        return cloudpickle.loads(self.argdata)

    def start(self) -> None:
        raise NotImplementedError(f"You need to create a Process object from this template to start it.")


class Process(TemplateProcess):
    """This object describes a process managed by the Dragon runtime."""

    def __new__(cls, *args, **kwargs):
        """Prevent users from subclassing `Process`.
        :raises NotImplementedError: If a subclass of `Process` is being instantiated
        :return: A Process object
        :rtype: Process
        """

        if cls != Process:
            raise NotImplementedError("Subclassing process not supported by the Dragon runtime.")

        return super(Process, cls).__new__(cls)

    def __init__(
        self,
        target: str or callable,
        args: tuple = None,
        kwargs: dict = None,
        cwd: str = ".",
        env: dict = None,
        ident: str or int = None,
        _pmi_enabled: bool = False,
        stdin: int = None,
        stdout: int = None,
        stderr: int = None,
        # policy: Policy = DefaultPolicy,
    ):
        """Generic Dragon process object executing a binary executable or a
        Python callable.

        Processes are named and can be looked up by specifying the `ident`
        keyword: `Process(None, ident="Apple")`

        If the target is a callable, start a new Python process executing the
        callable.  The class uses cloudpickle to transfer the callable to the
        child process. See the cloudpickle documentation for possible
        limitations.

        If the target is a string, assume it's the filename of a binary to be
        started.  If the target file is not found, the object tries to resolve
        the command via the PATH environment variable. If that fails, it assumes
        the file is in the current working directory.

        The Process class should not be subclassed, because the Dragon runtime
        only holds on to attributes of the Process super class, breaking the name
        lookup.

        :param target: The binary or Python callable to execute.
        :type target: str or callable
        :param args: arguments to pass to the binary or callable, defaults to None
        :type args: tuple, optional
        :param kwargs: keyword arguments to pass to the Python callable, defaults to None
        :type kwargs: tuple, optional
        :param cwd: current working directory, defaults to "."
        :type cwd: str, optional
        :param env: environment variables to pass to the process environment, defaults to None
        :type env: dict, optional
        :param ident: unique id or name of this process
        :type ident: int or str
        :param stdin: Stdin file handling. Valid values are PIPE and None.
        :type stdin: int, optional
        :param stdout: Stdout file handling. Valid values are PIPE, STDOUT and None.
        :type stdout: int, optional
        :param stderr: Stderr file handling. Valid values are PIPE, STDOUT and None.
        :type stderr: int, optional
        """

        self.started = False
        self.ident = ident
        self._pmi_enabled = _pmi_enabled

        # strip the name/uid from the parameters, as TemplateProcess cannot have it by definition.
        if ident:
            try:
                self._update_descriptor(ident=ident)
            except ProcessError:
                pass
            else:
                if target:
                    LOG.warning(f"Returning named process from infrastructure with valid target arg in call.")
                return  # we're done

        return super().__init__(target, args, kwargs, cwd, env, stdin, stdout, stderr)

    @classmethod
    def from_template(
        cls, template: TemplateProcess, ident: str = None, _pmi_enabled: bool = False
    ) -> object:
        """A classmethod that creates a new process object from a template.

        :param template: the template to base the process on
        :type template: TemplateProcess
        :param ident: intended name of the process, defaults to None
        :type ident: str, optional
        :return: The new process object
        :rtype: dragon.native.Process
        """
        # is there a way to write this less explicitly ?

        if template.is_python:
            target, args, kwargs = template.get_original_python_parameters()
        else:
            target = template.target
            args = template.args
            kwargs = template.kwargs

        if ident:  # subtle detail, a process from a template _has_ to be new !
            try:
                _ = process_query(ident)
            except ProcessError:
                pass
            else:
                raise ProcessError(f"A process '{ident}' already exists within the Dragon runtime.")

        return cls(target, args, kwargs, template.cwd, template.env, ident=ident, 
                   _pmi_enabled=_pmi_enabled, stdin=template.stdin, stdout=template.stdout, stderr=template.stderr)

    def start(self) -> None:
        """Start the process represented by the underlying process object."""

        if self.started:
            raise RuntimeError(f"This Process has already been started with puid {self.puid}")

        options = ProcessOptions(make_inf_channels=True)

        if self.is_python:

            self._descr = process_create_with_argdata(
                exe=self.target,
                run_dir=self.cwd,
                args=self.args,
                argdata=self.argdata,
                env=self.env,
                options=options,
                user_name=self.ident,
            )
        else:  # binary

            self._descr = process_create(
                exe=self.target,
                run_dir=self.cwd,
                args=self.args,
                env=self.env,
                options=options,
                user_name=self.ident,
                stdin=self.stdin,
                stdout=self.stdout,
                stderr=self.stderr
            )

        self.started = True

    @property
    def is_alive(self) -> bool:
        """Check if the process is still running

        :return: True if the process is running, False otherwise
        :rtype: bool
        """
        try:
            self._update_descriptor()
        except AttributeError:
            return False
        else:
            return self._descr.state == self._descr.State.ACTIVE

    def join(self, timeout: float = None) -> int:
        """Wait for the process to finish.

        :param timeout: timeout in seconds, defaults to None
        :type timeout: float, optional
        :return: exit code of the process, None if timeout occurs
        :rtype: int
        :raises: ProcessError
        """
        ecode = process_join(self._descr.p_uid, timeout=timeout)

    def terminate(self) -> None:
        """Send SIGTERM signal to the process, terminating it.

        :return: None
        :rtype: NoneType
        """
        process_kill(self._descr.p_uid, sig=signal.SIGTERM)

    def kill(self) -> None:
        """Send SIGKILL signal to the process, killing it.

        :return: None
        :rtype: NoneType
        """
        process_kill(self._descr.p_uid, sig=signal.SIGKILL)

    @property
    def name(self) -> str:
        return self._descr.name

    @property
    def puid(self) -> int:
        """Process puid. Globally unique"""
        return self._descr.p_uid

    @property
    def node(self) -> int:
        """Return the unique host id of the node the process is running on.

        :return: huid of host node.
        :rtype: int
        """
        return self._descr.node

    @property
    def returncode(self) -> int:
        """When the process has terminated, return exit code. None otherwise."""
        self._update_descriptor()
        return self._descr.ecode

    def children(self) -> list[object]:
        """Find all active children of the current process.

        :returns: a list of p_uids of active children of the current process
        :return type: list(int)
        """
        self._update_descriptor()
        return [Process(None, ident=puid) for puid in self._descr.live_children]

    def parent(self) -> object:
        """Get the process object of the parent process.
        :returns: The puid or the parent of the current process
        :return type: int
        """
        return Process(None, ident=self._descr.p_p_uid)

    def _update_descriptor(self, ident=None):
        self._descr = process_query(ident or self._descr.p_uid)

    @property
    def stdin_conn(self):
        return self._descr.stdin_conn

    @property
    def stdout_conn(self):
        return self._descr.stdout_conn

    @property
    def stderr_conn(self):
        return self._descr.stderr_conn

def current() -> Process:
    """Get the PythonProcess object of the current process.

    :returns: The current puid
    :return type: int
    """

    return Process(None, ident=this_process.my_puid)


# What follows is the Python process implementation. It is a little ... involved.


def _dragon_native_python_process_main():
    """This function is called by the Dragon runtime using the Python executable
    when a new Python process is started.

    It returns an exit code of 0, if the Python function exited without raising
    an exception or called `sys.exit()` with an error code of `None` or `0`.  If
    `sys.exit()` is called with a non-zero error code, it returns that error
    code.  Otherwise it returns an exit code of 1.
    """

    from dragon.globalservices.api_setup import connect_to_infrastructure

    connect_to_infrastructure()

    myp = current()

    LOG.debug(f"New Dragon Python Process puid={myp.puid}, pid={os.getpid()}, node={myp.node}")

    try:
        from dragon.globalservices.api_setup import _ARG_PAYLOAD  # import as late as possible

        target, args, kwargs = cloudpickle.loads(_ARG_PAYLOAD)  # unpack argdata from GS
    except Exception as e:
        LOG.error(f"Could not unpickle payload data:\n{e}")
        sys.exit(1)

    try:
        if threading._HAVE_THREAD_NATIVE_ID:
            threading.main_thread()._set_native_id()

        target(*args, **kwargs)  # run the payload
        ecode = 0
    except Exception as e:
        LOG.error(f"\nDragon Python Process puid={myp.puid} threw an exception: \n{e}")
        ecode = 1
    except SystemExit as e:  # catch sys.exit
        if e.code is None:
            ecode = 0
        elif isinstance(e.code, int):
            ecode = e.code
        else:
            ecode = 1
        if ecode != 0:
            LOG.error(f"\nDragon Python Process puid={myp.puid} exited abnormally: {ecode}")
    finally:
        threading._shutdown()
        try:
            sys.stdout.flush()
        except (AttributeError, ValueError):
            pass
        try:
            sys.stderr.flush()
        except (AttributeError, ValueError):
            pass

    sys.exit(ecode)

src/dragon/native/process.py
DragonHPC/dragon
"""The Dragon native class for managing the life-cycle of a group of Dragon processes.

The intent is for the class to be agnostic about what the processes are doing,
it only maintains their lifecycle.

This file implements a client API class and a Manager process handling all
groups on the node. The manager holds a list of GroupContext classes and gets
signals from the client classes using a queue. The group of processes undergoes
state transitions depending on the signals the client send to the manager.

Due to transparency constraints the manager cannot send information back to the client.
Instead we are using shared state

The underlying state machine looks like this:

.. image:: ./images/dragon_worker_pool.svg
    :scale: 75%
"""


import logging
import time
import enum
import signal
from queue import Empty

from abc import ABC, abstractmethod

import dragon

from .process import Process, TemplateProcess
from .queue import Queue
from .value import Value
from .lock import Lock
from .machine import current as current_node

from ..globalservices.process import multi_join, kill as process_kill, get_create_message, get_create_message_with_argdata, query as process_query
from ..globalservices.group import create, kill as group_kill, add_to, create_add_to, remove_from, destroy
from ..infrastructure.policy import Policy

LOG = logging.getLogger(__name__)

# exit code returned by cython for sigterm
# we also mod by 256 for unsigned char repr
CYTHON_SIGTERM_ECODE = -15


class DragonProcessGroupError(Exception):
    """Exceptions raised by the Dragon Pool Workers implementation"""

    pass


class DragonProcessGroupAbnormalExit(DragonProcessGroupError):
    """Exception raised by the Dragon Pool Workers implementation"""

    pass


# TODO: Extend API to control distribution over multiple nodes.


@enum.unique
class PGSignals(enum.IntEnum):
    ERROR = enum.auto()  # for testing purposes
    NEW = enum.auto()  # Start manager but no workers
    START = enum.auto()  # start all processes/workers
    JOIN = enum.auto()  # Wait for all the processes to complete
    SHUTDOWN = enum.auto()  # stop all processes/workers via SIGTERM
    KILL = enum.auto()  # forcefully stop all processes/workers via SIGKILL
    STOP = enum.auto()  # kill all Dragon processes and exit the manager


class BaseState(ABC):
    """This class declares methods that all concrete State classes should
    implement and also provides a backreference to the Context object,
    associated with the State. This backreference can be used by States to
    transition the Context to another State.
    It also defines common methods and data structures to all states.
    """

    forbidden: list[int] = None

    @property
    def context(self):
        """Link back to it's own context.

        :return: The group context holding this state
        :rtype: GroupContext
        """

        return self._context

    @context.setter
    def context(self, new_context) -> None:
        self._context = new_context

    @abstractmethod
    def run(self, prior_state, signal: PGSignals, sig_id: int) -> None:
        """This method runs the state it is a part of on the context."""
        pass

    def __str__(self):
        return self.__class__.__name__


# concrete states of the process group


class Error(BaseState):
    """This is the fallback state if an issue with the group occurs.
    The state of the processes is undefined here.
    """

    forbidden: list = [s for s in PGSignals if s not in [PGSignals.KILL, PGSignals.STOP]]

    def run(self, prior_state: BaseState, signal: PGSignals, sig_id: int) -> None:

        try:
            LOG.error(f"Process Group {self.context} is in error state.")
        except Exception:
            pass

        time.sleep(self.context.update_interval_sec * 10)


class Idle(BaseState):
    """This state kills existing processes and does nothing otherwise."""

    forbidden = [PGSignals.SHUTDOWN, PGSignals.KILL]

    def run(self, prior_state: BaseState, signal: PGSignals, sig_id: int) -> None:
        """The idle state just does nothing except making sure all processes are gone."""

        if self.context.guid:
            group_kill(self.context.guid)
            self.context.guid = self.context._group_descr = None

        self._start_time = None


class Maintain(BaseState):
    """This state starts missing processes and restarts processes that are not
    alive anymore.

    :raises DragonProcessGroupError: If one of the processes could not be (re)started.
    """

    forbidden = [PGSignals.NEW, PGSignals.START]

    def run(self, prior_state: BaseState, signal: PGSignals, sig_id: int) -> None:

        nretries = self.context.num_restart_retries

        if not self.context.guid:
            self.context._start_group_once()

        guid = self.context.guid

        group_descr = None
        for i, lst in enumerate(self.context._group_descr.sets):
            for descr in lst:
                puid = descr.uid
                # check the status of this process by asking GS
                proc_desc = process_query(puid)
                if proc_desc.state == descr.desc.State.ACTIVE:
                    continue

                # remove process from group
                group_descr = remove_from(guid, [puid])

                # create a new one
                msg = self.context.messages[self.context.puid_to_message_map[puid]]
                nrestarts = 0

                while nrestarts < nretries:
                    try:
                        group_descr = create_add_to(guid, [(1, msg.serialize())], self.context.policy)
                        # update the puid_to_message_map dict with the newly added process
                        puids = [descr.uid for lst in group_descr.sets for descr in lst]
                        for new_puid in puids:
                            if new_puid not in self.context.puid_to_message_map:
                                self.context.puid_to_message_map[new_puid] = self.context.puid_to_message_map[puid]
                                # since we added only one process, we can safely assume that we found it
                                break
                        break
                    except Exception as e:
                        nrestarts += 1

                if nrestarts == nretries:
                    raise DragonProcessGroupError(f"Unable to start process {i} using message {msg}.")

        # we need to update the group descriptor after all the above additions/removals
        if group_descr is not None:
            self.context._group_descr = group_descr

class Running(BaseState):

    # user needs to wait for group to become IDLE
    forbidden = [
        s for s in PGSignals if s not in [PGSignals.ERROR, PGSignals.KILL, PGSignals.SHUTDOWN, PGSignals.JOIN]
    ]

    def run(self, prior_state: BaseState, pgsignal: PGSignals, sig_id: int) -> None:

        # this is more complicated as it needs to be, because we're using
        # process objects and multi_join wants puids.

        timeout = self.context.update_interval_sec
        ignore_err = self.context.ignore_error_on_exit

        if prior_state == Idle:  # if we started with restart == False from Idle
            self.context._start_group_once()
            self.context.update_status(sig_id)

        if pgsignal == PGSignals.SHUTDOWN:  # have processes exit
            group_kill(self.context.guid, sig=signal.SIGTERM)

        # collect puids
        puids = self.context.puids

        # join on them
        ready = multi_join(puids, join_all=True, timeout=timeout)

        if ready[0] != None:  # no timeout

            # catch bad ecodes
            if not ignore_err:
                for puid, ecode in ready[0]:
                    if ecode not in {0, CYTHON_SIGTERM_ECODE, CYTHON_SIGTERM_ECODE % 256}:
                        LOG.debug(f"Bad exit code {ecode} for puid {puid} in ProcessGroup")
                        raise DragonProcessGroupAbnormalExit("Some processes in group exited abnormally !")

            # move on to Idle
            prior_state = self.context.transition_to(Idle)
            self.context.run(prior_state, pgsignal, sig_id)
            self.context.update_status(None)

class Stop(BaseState):
    """Stops all processes of the group and removes the group from the Manager. The
    group cannot be restarted anymore.
    """

    forbidden = [s for s in PGSignals]  # end of line

    def run(self, prior_state: BaseState, signal: PGSignals, sig_id: int) -> None:

        if self.context.guid != None:
            destroy(self.context.guid)
        self.context.guid = self.context._group_descr = None

# end concrete state classes

class GroupContext:
    """The Context defines the group interface for the manager and the client.
    In particular, it handles signals and state changes. It maintains a
    reference to an instance of a State subclass, which represents the current
    state of the group of processes.
    """

    _state: BaseState = None

    update_interval_sec: float = 0.5
    num_kill_retries: int = 2
    num_restart_retries: int = 2

    # why have more than 1 target status ?
    # because JOIN transitions to Join, and then auto-transitions to Idle.
    # so the target state for the manager is Join, but the client has to
    # wait for Signal completion on Join and Idle to not introduce a race condition.
    # The same is true for SHUTDOWN, which is JOIN with a SIGTERM.

    target_states: dict = {
        PGSignals.ERROR: [Stop],
        PGSignals.NEW: [Idle],
        PGSignals.START: [Maintain],
        PGSignals.JOIN: [Running, Idle],
        PGSignals.SHUTDOWN: [Running, Idle],
        PGSignals.KILL: [Idle],
        PGSignals.STOP: [Stop],
    }

    forbidden: dict = {
        str(Error()): Error.forbidden,
        str(Idle()): Idle.forbidden,
        str(Maintain()): Maintain.forbidden,
        str(Running()): Running.forbidden,
        str(Stop()): Stop.forbidden,
    }

    def __init__(
        self,
        templates: list[tuple],
        nproc: int,
        restart: bool,
        ignore_error_on_exit: bool,
        pmi_enabled: bool,  # MPI jobs
        walltime: float,
        policy: Policy,
    ) -> None:
        """This class represents a group of processes and exposes an interface
        to the Manager to handle signals and state transitions.

        :param templates: a list of tuples where each tuple contains a replication factor `n` and a Dragon TemplateProcess object specifing the properties of the process to start. The processes can hold the same or different attributes in any way and will be numbered in order.
        :type templates: list[tuple(int, dragon.native.process.TemplateProcess),]
        :param nproc: total number of processes that belong to the ProcessGroup
        :type nproc: int
        :param restart: wether to restart processes that exited prematurely, defaults to True
        :type restart: bool
        :param ignore_error_on_exit: whether to ignore errors when the group exists, defaults to False
        :type ignore_error_on_exit: bool
        :param pmi_enabled: wether to instruct Dragon to enable MPI support, defaults to False
        :type pmi_enabled: bool
        :param walltime: time in seconds to run processes before killing them.
        :type walltime: float
        :param policy: determines the placement of the group resources
        :type policy: dragon.infrastructure.policy.Policy
        """

        self.nproc = nproc
        self.templates = templates # list of tuples
        self.messages = {} # keys are the indices of tuples in self.templates

        # use a dict to make restarting easy and order-safe
        self.puid_to_message_map = {} # keys are the puids, values are the keys in self.messages{}

        for i, tup in enumerate(templates):
            t = tup[1]
            if t.is_python:
                self.messages[i] = get_create_message_with_argdata(t.target, t.cwd, t.args, t.env, t.argdata, pmi_required=pmi_enabled, stdin=t.stdin, stdout=t.stdout, stderr=t.stderr)
            else:
                self.messages[i] = get_create_message(t.target, t.cwd, t.args, t.env, pmi_required=pmi_enabled, stdin=t.stdin, stdout=t.stdout, stderr=t.stderr)

        self.guid = None
        self._group_descr = None

        self.restart = restart
        self.ignore_error_on_exit = ignore_error_on_exit
        self.pmi_enabled = pmi_enabled
        self.walltime = walltime
        if policy:
            self.policy = policy
        else:
            self.policy = Policy()

        self._start_time = 0

        self._state = None

        # count signals so we can report completion
        self._signal_counter = Value("i", value=0)
        self._signal_counter_lock = Lock()

        self._status_queue = Queue()
        self._status_queue.put((-1, -1, -1))

        self.transition_to(Idle)
        self.update_status(None)

    def __del__(self):

        try:
            self._status_queue.put(True)
        except Exception:
            pass

    def handle_signal(self, signal: PGSignals) -> tuple:
        """This method takes a signal, checks if the signal is allowed, given
        the current state, and returns the new state.
        If the walltime is up, it overwrites the signal with kill.

        :param signal: The signal to consider
        :type signal: Signal
        :return: target state
        :rtype: BaseState
        :raises RuntimeError: raised, if the transition is not allowed
        """

        if signal in self._state.forbidden:
            LOG.error(f"Signal not accepted {signal} for state {self._state} !")
            return

        if self._walltime_expired():
            signal = PGSignals.KILL  # goodbye

        if signal:
            next_state = self.target_states[signal][0]
        else:
            next_state = self.current_state

        return next_state, self.current_state

    def transition_to(self, new_state: BaseState) -> BaseState:
        """Transition the context to a new state. This does NOT run
        the new state or updates the status.

        :param new_state: state to transition to
        :type new_state: BaseState
        :return: previous state
        :rtype: BaseState
        """

        prior_state = self._state

        self._state = new_state
        self._state.context = self

        return prior_state

    @property
    def current_state(self) -> BaseState:
        """Return the current state object.

        :return: the state
        :rtype: BaseState
        """

        return self._state

    def run(self, prior_state: BaseState, signal: PGSignals, sig_id: int):
        """Execute the current (!) state in self._state.

        :param prior_state: prior state
        :type prior_state: BaseState
        """

        try:
            self._state.run(self._state, prior_state, signal, sig_id)  # Why ?
        except Exception as e:  # automagically transition to Error if something goes wrong
            LOG.error(f"Exception in transition {prior_state} -> {self._state}: {e} ")
            failed_state = self._state
            self.transition_to(Error)
            self.update_status(None)
            self.run(self._state, None, sig_id)

    def get_next_sig_id(self) -> int:
        """Obtain the next unique signal ID for this group context

        :return: next signal id
        :rtype: int
        """

        # we need to communicate if a signal completed, so we need an ID for it.

        with self._signal_counter_lock:
            sig_id = self._signal_counter.value
            self._signal_counter.value = 1 + self._signal_counter.value

        return sig_id

    @property
    def status(self) -> str:
        """The current status of the Group.

        :return: state name
        :rtype: str
        """

        item = self._status_queue.get(timeout=None)
        self._status_queue.put(item)
        state_name, _, __ = item

        return state_name

    @property
    def puids(self) -> list[int or None]:
        """The puids maintained by the GroupContext.

        :return: A list of puids if the process is started or None if not.
        :rtype: list[int or None]
        """

        item = self._status_queue.get(timeout=None)
        self._status_queue.put(item)
        _, puids, __ = item

        return puids

    @property
    def last_completed_signal(self) -> int:
        """Return the ID of the last successfully completed signal.

        :return: the signal id
        :rtype: int
        """

        item = self._status_queue.get(timeout=None)
        self._status_queue.put(item)
        _, __, compl_sig_id = item

        return compl_sig_id

    def _start_group_once(self):
        if not self._start_time:
            self._start_time = time.monotonic()

        group_descr = create([ (tup[0], self.messages[i].serialize()) for i, tup in enumerate(self.templates) ], self.policy)

        self._group_descr = group_descr
        self.guid = group_descr.g_uid

        # construct the puid_to_message_map dict now that the processes are created
        puids = [descr.uid for lst in group_descr.sets for descr in lst]
        puid_idx = 0
        for i, tup in enumerate(self.templates):
            for _ in range(tup[0]):
                self.puid_to_message_map[puids[puid_idx]] = i
                puid_idx += 1

    def _walltime_expired(self) -> bool:
        """Check if the walltime has expired

        :return: True if the walltime has expired, False otherwise
        :rtype: bool
        """

        if self.walltime and self._start_time > 0:
            expired = time.monotonic() - self._start_time >= self.walltime
        else:
            expired = False

        return expired

    # this is a hack using a queue to communicate the status. We would want to
    # use a few Array objects and a single lock to share the puids and status
    # string.  I need to report the last completed signal ID here, so that the
    # corresponding client can be be sure his request has been completed.

    def update_status(self, compl_sig_id: int):
        """update the global status of this context.

        :param compl_sig_id: signal just completed. If None, reuse the last value from the queue.
        :type compl_sig_id: int
        """

        _, __, last_sig_id = self._status_queue.get(timeout=None)

        if compl_sig_id == None:  # this transition was automatic
            compl_sig_id = last_sig_id

        state_name = str(self._state.__name__)

        if self._group_descr:
            puids = [descr.uid for lst in self._group_descr.sets for descr in lst]
        else:
            puids = [None for _ in range(self.nproc)]

        item = (state_name, puids, compl_sig_id)
        self._status_queue.put(item)


class Manager:
    def __init__(self):
        """The Manager class holds a process that handles the life-cycle of all
        process groups on this node.
        We handle the group using a group_context class that holds all necessary
        methods and attributes shared by manager and clients.
        """

        # I am not yet using a single node-wide process here, because I need to
        # discover it and its communication queue, but I cannot ask Dragon for that
        # named queue yet.

        node = current_node()

        # missing shared memory implementation, so I am using a Queue
        self.group_context = None

        ident = f"_DPoolQ1-{node.h_uid}-{id(self)}"
        self.queue = Queue()  # ident=ident) # need a named queue here

        ident = f"_DPoolMan-{node.h_uid}-{id(self)}"
        self._proc = Process(
            self._signal_handler,
            (self.queue,),
            # ident=ident,
        )
        self._proc.start()

    def __del__(self):

        # Kill processes and remove context
        try:
            self.queue.put((PGSignals.STOP, None))
        except Exception:
            pass

        # wait for manager to exit
        try:
            self._proc.join(timeout=self.update_interval_sec)
        except Exception:
            pass

        # make manager exit
        try:
            self._proc.kill()  # please go away
        except Exception:
            pass

    def _signal_handler(self, queue: Queue):
        """Get a new signal from the queue, obtain the new state, transition the
        context, run the new context, update the status.
        If the queue is empty, run the current context again.
        """

        signal = None

        while signal != PGSignals.STOP:

            try:
                signal, payload, sig_id = queue.get(timeout=0)
            except Empty:
                signal = payload = sig_id = None  # no change

            if signal == PGSignals.NEW and payload:  # new context !
                self.group_context = payload

            # this looks overengineered, but I need to cover Join :
            # * Join.run depends on the prior state (Idle or Maintain), need to save it
            # * transitioning is not followed by a run for Join->Idle
            # * update_status has to be run multiple times by a Join
            # * we update the status without running the state in Join
            if self.group_context:
                new_state, prior_state = self.group_context.handle_signal(signal)
                self.group_context.transition_to(new_state)
                self.group_context.run(prior_state, signal, sig_id)
                self.group_context.update_status(sig_id)

                time.sleep(self.group_context.update_interval_sec)  # do not spin too hot

        self.group_context = None


class ProcessGroup:
    """Robustly manage the lifecycle of a group of Dragon processes using Dragon Global Services.

    This is really a state machine of the underlying processes. We should always
    be able to "ask" the manager process for the state of the group and send it a
    signal to make a state transition.
    """

    def __init__(
        self,
        restart: bool = True,
        ignore_error_on_exit: bool = False,
        pmi_enabled: bool = False,
        walltime: float = None,
        policy: Policy = None
    ):
        """Instantiate a number of Dragon processes.

        :param restart: if True, restart worker processes that exit unexpectedly and suppress any errors from them, defaults to True.
        :type restart: bool
        :param ignore_error_on_exit: If to ignore errors coming from processes when they exit from Join state.
        :type ignore_error_on_exit: bool
        :param flags: optional flags that affect the handling of a worker process.
        :type flags: Worker.Flags
        :param pmi_enabled: Instruct the runtime to setup the environment so that the binary can use MPI for inter-process communication.
        :type pmi_enabled: Bool
        :param walltime: Time in seconds until the processes in the group get killed
        :type walltime: float
        :param policy: determines the placement of the processes
        :type policy: dragon.infrastructure.policy.Policy
        """

        self.templates = [] # this will be a list of tuples that will be sent to the GSGroup API
        self.nproc = 0
        self.restart = restart
        self.ignore_error_on_exit = ignore_error_on_exit
        self.pmi_enabled = pmi_enabled
        self.walltime = walltime
        self.policy = policy
        self._group_context = None

    def add_process(self, nproc: int, template: TemplateProcess) -> None:
        """Add processes to the ProcessGroup.

        :param template: single template processes, i.e. unstarted process objects
        :type template: dragon.native.process.TemplateProcess
        :param nproc: number of Dragon processes to start that follow the provided template
        :type nproc: int
        """

        # if add_process is called after the ProcessGroup is initialized, then we raise
        if self._group_context:
            raise DragonProcessGroupError("You cannot call add_process() to already initialized ProcessGroup. Please use ProcessGroup.create_add_to() instead to add more processes.")

        self.templates.append((nproc, template))
        self.nproc += nproc

    def init(self) -> None:
        """Initialize the GroupContext and Manager."""

        self._group_context = GroupContext(self.templates, self.nproc, self.restart, self.ignore_error_on_exit, self.pmi_enabled, self.walltime, self.policy)
        self._manager = Manager()
        self._send_signal(PGSignals.NEW)

    def start(self) -> None:
        """Starts up all processes according to the templates. If `restart ==
        False`, transition to 'Running', otherwise transition to 'Maintain'.
        """

        if not self.restart:
            self._send_signal(PGSignals.JOIN)
        else:
            self._send_signal(PGSignals.START)

    def join(self, timeout: float = None) -> None:
        """Wait for all processes to complete and the group to transition to
        Idle state. If the group status is 'Maintain', transition to 'Running'.

        Raises TimeoutError, if the timeout occurred.

        :param timeout: Timeout in seconds, optional defaults to None
        :type timeout: float
        :return: True if the timeout occured, False otherwise
        :retype: bool
        """

        start = time.monotonic()
        if self.status == str(Maintain()):
            self._send_signal(PGSignals.JOIN)
        stop = time.monotonic()

        if timeout == None:
            timeout = 100000000

        timeout = max(0, timeout - (stop - start))

        return self._wait_for_status(str(Idle()), timeout=timeout)

    def kill(self, signal: signal.Signals = signal.SIGKILL) -> None:
        """Send a signal to each process of the process group.

        The signals SIGKILL and SIGTERM have the following side effects:

        * If the signal is SIGKILL, the group will transition to 'Idle'. It can then be reused.
        * If the group status is 'Maintain', SIGTERM will transition it to 'Running'.
        * If the group status is 'Error', SIGTERM will raise a `DragonProcessGroupError`.

        :param signal: the signal to send, defaults to signal.SIGKILL
        :type signal: signal.Signals, optional
        """

        if signal == signal.SIGTERM:
            self._send_signal(PGSignals.SHUTDOWN)
        elif signal == signal.SIGKILL:
            self._send_signal(PGSignals.KILL)
        else:
            for puid in self.puids:
                process_kill(puid, sig=signal)

    def stop(self) -> None:
        """Forcibly terminate all workers by sending `SIGKILL` from any state,
        transition to `Stop`. This also removes the group from the manager process
        and marks the end of the group life-cycle.
        """

        self._send_signal(PGSignals.STOP)

    @property
    def puids(self) -> list[int]:
        """Return the puids of the processes contained in this group.

        :return: a list of puids
        :rtype: list[int or None]
        """

        return self._group_context.puids

    @property
    def status(self) -> str:
        """Get the current status of the process group handled by this instance.

        :returns: current status of the group
        :rtype: str
        """

        return self._group_context.status

    # Private interface

    def _send_signal(self, signal: PGSignals) -> bool:
        """Send the signal to the manager and wait for the response.
        The method guarantees completion of the signal by the manager,
        nothing more. I.e. the processes may have been started, but not
        actually executed any useful code yet.
        In case of sending IDLE, include the group context as well.
        """

        status = self.status

        if signal in self._group_context.forbidden[status]:
            raise DragonProcessGroupError(f"Signal {str(signal)} is not a valid transition from {status}")

        if signal == PGSignals.NEW:
            payload = self._group_context
        else:
            payload = None

        sig_id = self._group_context.get_next_sig_id()

        self._manager.queue.put((signal, payload, sig_id))

        while self._group_context.last_completed_signal < sig_id:
            if self.status == str(Error()):
                if signal not in [PGSignals.KILL, PGSignals.STOP]:
                    if self._group_context.last_completed_signal == sig_id - 1:
                        raise DragonProcessGroupError(
                            f"Signal {str(signal)} was not successful. Group in ERROR state."
                        )
                    else:
                        raise DragonProcessGroupError(
                            f"Signal {str(signal)} cannot be completed. Group in ERROR state"
                        )

            time.sleep(self._group_context.update_interval_sec)

    def _wait_for_status(self, status: str, timeout: float = None) -> None:

        sleep_time = self._group_context.update_interval_sec

        start = time.monotonic()

        while not self._group_context.status == status:

            dt = timeout - time.monotonic() + start

            if dt <= 0:
                raise TimeoutError(f"Timeout waiting for status {status}")

            time.sleep(min(sleep_time, dt))

src/dragon/native/process_group.py
DragonHPC/dragon
""" The Dragon native queue provides ordered object communication between processes.

This is Dragon's specialized implementations of a Queue, working
in both single and multi node settings. The API is similar to Multiprocessing.Queue
but has a number of extentions and simplifications.

The implementation relies on Dragon GS Client API and the Dragon runtime libraries.
"""

import cloudpickle as cp
import queue
import time
import logging

from ..channels import (
    Channel,
    ChannelEmpty,
    ChannelFull,
    ChannelError,
    Many2ManyReadingChannelFile,
    Many2ManyWritingChannelFile,
)
from ..managed_memory import MemoryPool

from ..globalservices.channel import create, destroy, query, release_refcnt, get_refcnt
from ..infrastructure.facts import default_pool_muid_from_index
from ..infrastructure.parameters import this_process, POLICY_USER, Policy
from ..infrastructure.channel_desc import ChannelOptions
import dragon.utils as du

_DEF_MUID = default_pool_muid_from_index(this_process.index)

LOG = logging.getLogger(__name__)


class QueueError(Exception):
    pass


class Queue(object):
    """A Dragon native Queue relying on Dragon Channels.

    The interface resembles Python Multiprocessing.Queues. In particular, the
    queue can be initialized as joinable, which allows to join on the completion
    of a items. We have simplified the interface a little and changed a few oddities
    like renaming `qsize()` into `size()`.

    The usual `queue.Empty` and `queue.Full` exceptions from
    the Python standard library’s queue module are raised to signal timeouts.

    This implementation is intended to be a simple and high performance
    implementation for use on both single node and multiple nodes.
    """

    def __init__(
        self,
        maxsize: int = 100,
        *,
        m_uid: int = _DEF_MUID,
        block_size: int = 2**16,  # 64 kbytes
        joinable: bool = False,
        policy: Policy = POLICY_USER,
        _ext_channel: object = None,
    ):
        """Init method

        :param maxsize: sets the upperbound limit on the number of items that can be placed in the queue, defaults to 100
        :type maxsize: int, optional
        :param m_uid: The m_uid of the memory pool to use, defaults to _DEF_MUID
        :type m_uid: int, optional
        :param block_size: Block size for the underlying channel, defaults to 64 kbytes
        :type block_size: int, optional
        :param joinable: If this queue should be joinable, defaults to False
        :type joinable: bool, optional
        :param policy: policy object, defaults to POLICY_USER
        :type policy: object, optional
        :param _ext_channel: non-local externally managed channel for testing purposes only, defaults to None
        :type _ext_channel: instance of dragon.channel, optional
        :raises QueueError: If something goes wrong during creation
        :raises ValueError: If a parameter is wrong
        :raises NotImplementedError: If a joinable queue is used with an external channel
        """

        if not isinstance(policy, Policy):
            raise ValueError("The class of service must be a dragon.infrastructure.parameters.Policy value.")

        self._policy = policy
        self._closed = False
        self._ext_channel = False
        self._buffer_pool = None
        self._read_adapter = None
        self._write_adapter = None
        self._channel = None
        # Helper variable to know when we need to attach/detach to/from an externally managed channel.
        # In the case of an externally managed channel, we need to attach to the channel only when the
        # queue instance is created from unpickling. When the queue instance is created by the user
        # via the constructor, then we do not need to attach/detach.
        self._unpickled_instance = False

        # if we are joinable, create channels to hold a counter and a event
        self._joinable = joinable

        self._read_adapter_event = None
        self._write_adapter_event = None
        self._read_adapter_count = None
        self._write_adapter_count = None

        if _ext_channel is None:  # create the channel internally

            the_options = ChannelOptions(ref_count=True)
            the_options.local_opts.block_size = block_size

            if maxsize > 0:
                the_options.local_opts.capacity = int(maxsize)
            else:
                LOG.debug(f"Maxsize was {maxsize}. Using default queue capacity of 100.")
                maxsize = 100

            try:
                descr = create(m_uid=m_uid, options=the_options)
                self._channel = Channel.attach(descr.sdesc)
            except Exception as e:
                self._closed = True
                raise QueueError(f"Could not create queue\n{e}")

        else:
            self._channel = _ext_channel
            if self._channel.is_local:
                raise ValueError(
                    f"Channel {self._channel} is an externally managed channel and is not considered as remote."
                )
            self._ext_channel = True
            self._buffer_pool = m_uid  # only in this case, m_uid is a MemoryPool object

        if self._joinable:

            if self._ext_channel:
                raise NotImplementedError(f"Joinable queues do not support testing with external channels")

            the_options.local_opts.capacity = 1
            try:
                cnt_descr = create(m_uid=m_uid, options=the_options)
                self._cnt_channel = Channel.attach(cnt_descr.sdesc)
                ev_descr = create(m_uid=m_uid, options=the_options)
                self._ev_channel = Channel.attach(ev_descr.sdesc)
            except Exception as e:
                self._closed = True
                raise QueueError(f"Could not create queue\n{e}")

            self._reset_task_counter()
            self._reset_event_channel()

            LOG.debug("Created queue {self!r}")

    def __getstate__(self):

        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")

        if self._joinable:
            cnt_chan_cuid = self._cnt_channel.cuid
            ev_chan_cuid = self._ev_channel.cuid
        else:
            cnt_chan_cuid = None
            ev_chan_cuid = None

        return (
            self._channel.serialize(),
            self._channel.cuid,
            self._joinable,
            self._policy,
            self._ext_channel,
            cnt_chan_cuid,
            ev_chan_cuid,
        )

    def __setstate__(self, state):
        (
            ser_ch,
            q_chan_cuid,
            self._joinable,
            self._policy,
            self._ext_channel,
            cnt_chan_cuid,
            ev_chan_cuid,
        ) = state

        self._read_adapter = None
        self._write_adapter = None
        self._read_adapter_event = None
        self._write_adapter_event = None
        self._read_adapter_count = None
        self._write_adapter_count = None
        self._buffer_pool = None
        self._unpickled_instance = False
        self._channel = None
        self._closed = False

        try:
            if self._ext_channel:
                self._unpickled_instance = True
            else:
                get_refcnt(q_chan_cuid)

            self._channel = Channel.attach(ser_ch)

        except ChannelError as e:
            self._closed = True
            raise QueueError(f"Could not deserialize queue\n{e}")

        if self._joinable:  # attach to counter and event channels
            try:
                descr = query(cnt_chan_cuid)  # implicitly increases refcount
                self._cnt_channel = Channel.attach(descr.sdesc)

                descr = query(ev_chan_cuid)
                self._ev_channel = Channel.attach(descr.sdesc)

            except ChannelError as e:
                try:
                    release_refcnt(q_chan_cuid)
                    self._channel.detach()
                except:
                    pass
                self._closed = True
                raise QueueError(f"Could not deserialize joinable queue\n{e}")

        # figure out what should be the pool in which we will allocate the messages to send:
        # if the channel is local, then the pool is the channel pool
        # else the pool is the default pool of the caller
        try:
            if not self._channel.is_local:
                # if it is an externally managed channel, we have already provided a pool
                if not self._ext_channel:
                    self._buffer_pool = MemoryPool.attach(du.B64.str_to_bytes(this_process.default_pd))
        except Exception as e:
            try:
                release_refcnt(q_chan_cuid)
            except:
                pass
            self._closed = True
            raise QueueError(f"Could not deserialize queue\n{e}")


    def __del__(self):
        self._close()

    def close(self) -> None:
        """Indicate that no more data will be put on this queue by the current process.
        The refcount of the background channel will be released and the channel might be
        removed, if no other processes are holding it.
        """

        self._close()

    def full(self) -> bool:
        """Return True if the queue is full, False otherwise.

        :raises ValueError: If the queue is closed
        :return: Wether or not the queue is full
        :rtype: bool
        """

        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")

        return self._channel.num_msgs == self._channel.capacity

    def empty(self) -> bool:
        """Return True if the queue is empty, False otherwise.
        This might not be reliable

        :raises ValueError: If the queue is closed
        :return: Wether or not the queue is empty
        :rtype: bool
        """

        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")

        return not self._channel.poll(timeout=0)

    def size(self) -> int:
        """Return the approximate size of the queue. This number may not be reliable.

        :raises ValueError: If the queue is closed
        :return: approximate number of items in the queue
        :rtype: int
        """

        return self._size()

    def get(self, block: bool = True, timeout: float = None) -> object:
        """Remove and return an item from the queue.

        :param block: Make this call blocking, defaults to True
        :type block: bool, optional
        :param timeout: number of seconds to block, defaults to None
        :type timeout: float, optional
        :raises ValueError: If queue is closed
        :raises queue.Empty: If queue is empty
        :return: The next item in the queue
        :rtype: object
        """

        return self._get(block=block, timeout=timeout)

    def get_nowait(self) -> object:
        """Equivalent to get(False).

        :return: The next item in the queue
        :rtype: object
        """
        return self.get(False)

    def put(self, obj, block: bool = True, timeout: float = None) -> None:
        """Puts the serialization of an object onto the queue.
        If the queue is joinable, require one more call to task_done(), for
        join() to unblock.

        :param obj: object to serialize and put
        :param block: Whether to block
        :param timeout: Timeout, if blocking.  None means infinity, default
        :return: None
        """

        self._put(obj, block=block, timeout=timeout)

    def put_nowait(self, obj) -> None:
        """Equivalent to put(obj, False).

        :param obj: object to serialize and put
        :type obj: object
        :return: None
        """
        return self.put(obj, False)

    def task_done(self) -> None:
        """Indicate that a formerly enqueued task is complete. Used by queue consumers.
        Only for joinable queues.

        :raises QueueError: If this queue is not joinable or there were no pending tasks.
        :raises ValueError: If called more times than there were items in the queue.
        """

        if not self._joinable:
            raise QueueError("Queue is not joinable.")

        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")

        counter = self._decrease_task_counter()

        if counter < 0:
            self._reset_task_counter()
            raise QueueError("All tasks are already done")

        if counter == 0:  # release all processes currently joining the queue
            self._set_event_channel()

    def join(self, *args, timeout: float = None, **kwargs) -> None:
        """Block until all items in the queue have been gotten and processed.

        If a join() is blocking, it will resume when all items have been processed
        (meaning that a task_done() call was received for every item that had been put() into
        the queue).

        :param timeout: time to wait on the join in sec, optional
        :type timeout: float
        :return: None
        """

        if not self._joinable:
            raise QueueError(f"Queue is not joinable.")

        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")

        self._ev_channel.poll(timeout=timeout)

    @property
    def closed(self) -> bool:
        return self._closed

    # for dragon wait

    def poll(self, timeout: float = 0) -> None:
        """Block if the queue is empty
        :param timeout: time to wait in seconds, defaults to 0
        :type timeout: float, optional
        :return: None
        """
        return self._poll(timeout=timeout)

    def destroy(self) -> None:
        """Remove underlying channel right now."""
        self._close()
        try:
            destroy(self._channel.cuid)
        except:
            pass

    def thread_wait(self, timeout: float, done_ev: object, ready: list) -> None:
        """Thread waiter signaling with an ev."""

        return self._thread_wait(timeout, done_ev, ready)

    # private methods

    def _setup_write_adapter(self):
        self._write_adapter = Many2ManyWritingChannelFile(
            self._channel,
            self._buffer_pool,
            wait_mode=self._policy.wait_mode,
            return_mode=self._policy.return_when,
        )

    def _setup_read_adapter(self):
        self._read_adapter = Many2ManyReadingChannelFile(self._channel, wait_mode=self._policy.wait_mode)

    def _setup_write_adapter_event(self):
        self._write_adapter_event = Many2ManyWritingChannelFile(
            self._ev_channel,
            self._buffer_pool,
            wait_mode=self._policy.wait_mode,
            return_mode=self._policy.return_when,
        )

    def _setup_read_adapter_event(self):
        self._read_adapter_event = Many2ManyReadingChannelFile(
            self._ev_channel, wait_mode=self._policy.wait_mode
        )

    def _setup_write_adapter_count(self):
        self._write_adapter_count = Many2ManyWritingChannelFile(
            self._cnt_channel,
            self._buffer_pool,
            wait_mode=self._policy.wait_mode,
            return_mode=self._policy.return_when,
        )

    def _setup_read_adapter_count(self):
        self._read_adapter_count = Many2ManyReadingChannelFile(
            self._cnt_channel, wait_mode=self._policy.wait_mode
        )

    def _call_read_adapter(self, adapter, block, timeout):
        try:
            adapter.open()
            adapter.set_adapter_timeout(block, timeout)
            payload = cp.load(adapter)
        except (ChannelEmpty, TimeoutError):
            raise queue.Empty
        except Exception as e:
            raise QueueError(f"Could not get the object from queue.\n{e}")
        finally:
            adapter.close()

        return payload

    def _call_write_adapter(self, obj, adapter, block, timeout):
        try:
            adapter.open()
            adapter.set_adapter_timeout(block, timeout)
            cp.dump(obj, file=adapter, protocol=5)
        except (ChannelFull, TimeoutError):
            raise queue.Full
        except AttributeError:
            # needed for:
            # WithProcessesTestQueue.test_queue_feeder_on_queue_feeder_error
            # and WithProcessesTestQueue.test_queue_feeder_donot_stop_onexc
            raise AttributeError
        except Exception as e:
            raise QueueError(f"Could not put the object to queue.\n{e}")
        finally:
            adapter.close()

    def _get(self, block=True, timeout=None):

        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")

        try:
            if self._read_adapter is None:
                self._setup_read_adapter()
            payload = self._call_read_adapter(self._read_adapter, block, timeout)
        except (ChannelEmpty, TimeoutError):
            raise queue.Empty

        return payload

    def _put(self, obj, block=True, timeout=None):

        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")
        try:
            if self._write_adapter is None:
                self._setup_write_adapter()
            self._call_write_adapter(obj, self._write_adapter, block, timeout)
        except (ChannelFull, TimeoutError):
            raise queue.Full

        if self._joinable:
            counter = self._increase_task_counter()

            if counter == 1:
                self._reset_event_channel()

    def _poll(self, timeout: float = 0):
        return self._channel.poll(timeout=timeout)

    def _close(self):
        if not self._closed:
            self._closed = True

            if not self._ext_channel:
                try:
                    release_refcnt(self._channel.cuid)
                    self._channel.detach()
                except Exception as e:
                    LOG.debug(f'Could not complete release of refcount or detach. {e=}')

            elif self._ext_channel and self._unpickled_instance:
                try:
                    self._channel.detach()  # this should be revisited once refcounting fully works as it should not be necessary then
                except Exception as e:
                    LOG.debug(f"We couldn't detach from externally managed channel. {e=}")
            if self._joinable:
                try:
                    release_refcnt(self._cnt_channel.cuid)
                    self._cnt_channel.detach()
                    release_refcnt(self._ev_channel.cuid)
                    self._ev_channel.detach()
                except Exception as e:
                    LOG.debug(
                      f'Joinable queue: there was a problem with releasing channel refcount or while detaching from channel. {e=}')

        if self._buffer_pool is not None:
            try:
                self._buffer_pool.detach()
            except Exception as e:
                LOG.debug(f"We couldn't detach from queue's buffer pool. {e=}")

    def _size(self) -> int:
        if self._closed:
            raise ValueError(f"Queue {self!r} is closed")

        return self._channel.num_msgs

    def _thread_wait(self, timeout, done_ev, ready):
        if timeout is None:
            timeout = 1000000
        start = time.monotonic()
        delta = min(0.01, timeout)  # TODO: figure out the value for delta
        timed_out = False

        if not done_ev.is_set():
            while not timed_out and not done_ev.is_set():
                done_res = self.poll(delta)
                if not done_ev.is_set():
                    if done_res:
                        ready.append(self)
                        done_ev.set()
                timed_out = (time.monotonic() - start) > timeout
            done_ev.set()  # if not set here, timed out.

    def _reset_event_channel(self) -> None:
        if self._ev_channel.num_msgs > 0:
            if self._read_adapter_event is None:
                self._setup_read_adapter_event()
            _ = self._call_read_adapter(self._read_adapter_event, False, None)

    def _set_event_channel(self) -> None:
        if self._write_adapter_event is None:
            self._setup_write_adapter_event()
        self._call_write_adapter(1, self._write_adapter_event, False, None)

    def _reset_task_counter(self) -> None:
        if self._cnt_channel.num_msgs > 0:
            if self._read_adapter_count is None:
                self._setup_read_adapter_count()
            counter = self._call_read_adapter(self._read_adapter_count, True, None)

        counter = 0

        if self._write_adapter_count is None:
            self._setup_write_adapter_count()
        self._call_write_adapter(counter, self._write_adapter_count, True, None)

    def _increase_task_counter(self) -> int:
        if self._read_adapter_count is None:
            self._setup_read_adapter_count()
        counter = self._call_read_adapter(self._read_adapter_count, True, None)

        counter += 1

        if self._write_adapter_count is None:
            self._setup_write_adapter_count()
        # with self._write_adapter_count as adapter:
        #     adapter.set_adapter_timeout(True, None)
        #     pickle.dump(counter, file=adapter, protocol=5)
        self._call_write_adapter(counter, self._write_adapter_count, True, None)

        return counter

    def _decrease_task_counter(self) -> int:
        if self._read_adapter_count is None:
            self._setup_read_adapter_count()
        counter = self._call_read_adapter(self._read_adapter_count, True, None)

        counter -= 1
        if self._write_adapter_count is None:
            self._setup_write_adapter_count()
        self._call_write_adapter(counter, self._write_adapter_count, True, None)
        return counter

src/dragon/native/queue.py
DragonHPC/dragon
"""The Dragon native implementation of the classical semaphore synchronization object.

The API follows the Python Multiprocessing API of Semaphore and extends it with
a `get_value()` method.
"""

import logging
import time
import os

from .lock import Lock
from ..channels import Channel, ChannelRecvTimeout, ChannelEmpty, ChannelRecvH, ChannelSendH
from ..dtypes import WHEN_DEPOSITED
from ..globalservices.channel import get_refcnt, release_refcnt, create
from ..infrastructure.channel_desc import ChannelOptions
from ..infrastructure.parameters import this_process
from ..infrastructure.facts import default_pool_muid_from_index

LOGGER = logging.getLogger(__name__)

_DEF_MUID = default_pool_muid_from_index(this_process.index)


class Semaphore:
    """This class implements Semaphore objects on top of Dragon channels.
    We use a Dragon native Lock to protect a channel with a single entry that
    is used to count the number of acquire and release calls.
    We extend the standard API with a get_value method to ask the
    Semaphore for the current value.
    """

    _BYTEORDER = "big"
    _POLL_FREQ = 0.01

    def __del__(self):
        if not self._closed:
            self._closed = True
            try:
                self._writer.close()
                self._reader.close()
                release_refcnt(self._channel.cuid)
                self._channel.detach()
            except Exception:
                pass

    def __enter__(self):
        self.acquire()

    def __exit__(self, *args):
        self.release()

    def __init__(self, value: int = 1, *, m_uid: int = _DEF_MUID, bounded=False):
        """Create a Semaphore object that can be shared across Dragon processes.

        :param value: initial value for the counter, defaults to 1
        :type value: int, optional
        :param m_uid: dragon managed memory pool id to use, defaults to the standard pool.
        :type m_uid: int, optional
        :param bounded: set if the value of the Semaphore cannot exceed its initial value, defaults to False
        :type bounded: bool, optional
        """

        LOGGER.debug(f"Init Semaphore with value={value}, m_uid={m_uid}, bounded={bounded}")

        self._closed = True

        self._bounded = bounded
        self._initial_value = value

        self._lock = Lock()

        the_options = ChannelOptions(ref_count=True)
        the_options.local_opts.capacity = 1

        descr = create(m_uid=m_uid, options=the_options)

        self._channel_sdesc = descr.sdesc
        self._channel = Channel.attach(descr.sdesc)
        # get_refcnt(self._channel.cuid)

        self._reset()

        self._lmsg = value.to_bytes(8, byteorder=self._BYTEORDER)

        if value > 0:
            self._writer.send_bytes(self._lmsg)

        self._closed = False

    def __getstate__(self) -> tuple:

        if self._closed:
            raise ValueError(f"Semaphore {self!r} is closed and cannot be serialized.")

        sdesc = self._channel.serialize()

        return sdesc, self._bounded, self._initial_value, self._lock

    def __setstate__(self, state) -> None:

        self._closed = True

        ch_sdesc, self._bounded, self._initial_value, self._lock = state

        self._channel = Channel.attach(ch_sdesc)
        get_refcnt(self._channel.cuid)

        self._reset()

        value = self.get_value()
        self._lmsg = value.to_bytes(8, byteorder=self._BYTEORDER)

        self._closed = False

    def acquire(self, blocking: bool = True, timeout: float = None) -> bool:
        """Acquire the Semaphore, decrementing the counter, blocking the other
        processes, if the counter is 0.

        :param blocking: block the process if the Semaphore cannot be acquired immediately, defaults to True
        :type blocking: bool, optional
        :param timeout: timeout for a block in seconds, defaults to None
        :type timeout: float, optional
        """

        LOGGER.debug(f"Acquire Semaphore with blocking={blocking}, timeout={timeout}")

        # We need to protect the channel with a lock to stop other processes from
        # reading or writing into it while we are incrementing the counter.
        # This is ugly and not very performant
        # TODO: use dragon_channel_write_lock instead of a poor mans wait here.
        # That should get rid of the self._lock

        if timeout is None:
            timeout = 1000000000

        if timeout < 0 or not blocking:
            timeout = 0

        sleep_time = min(self._POLL_FREQ, timeout)

        start = time.monotonic()
        stop = start

        success = False
        while (stop - start) <= timeout:

            if not self._lock.acquire(block=False, timeout=None):
                time.sleep(sleep_time)
                stop = time.monotonic()
                continue

            try:
                self._lmsg = self._reader.recv_bytes(timeout=None, blocking=False)
                success = True
                break
            except ChannelEmpty:  # empty channel - keep waiting
                pass

            self._lock.release()

            time.sleep(sleep_time)

            stop = time.monotonic()

        if not success:
            return False

        # still holding the lock here
        value = int.from_bytes(self._lmsg, byteorder=self._BYTEORDER)
        value -= 1

        if value > 0:  # only write back to channel if value is not 0, else leave it empty
            self._lmsg = value.to_bytes(8, self._BYTEORDER)
            self._writer.send_bytes(self._lmsg, timeout=None, blocking=True)

        self._lock.release()

        return True

    def release(self, n: int = 1) -> None:
        """Release the Semaphore, incrementing the internal value by n,
        thus unblocking up to n processes.

        :param n: how much to increment the internal value by, defaults to 1
        :type n: int, optional
        """

        LOGGER.debug(f"Release Semaphore n={n}")

        self._lock.acquire(block=True, timeout=None)

        try:
            self._lmsg = self._reader.recv_bytes(timeout=None, blocking=False)
        except ChannelEmpty:
            value = 0
        else:  # use read value
            value = int.from_bytes(self._lmsg, byteorder=self._BYTEORDER)

        value += n

        if self._bounded:
            if value > self._initial_value:
                self._writer.send_bytes(self._lmsg, timeout=None, blocking=True)
                raise ValueError(
                    f"BoundedSemaphore value of {value} would exceed initial value of {self._initial_value}"
                )

        self._lmsg = value.to_bytes(8, self._BYTEORDER)

        self._writer.send_bytes(self._lmsg, timeout=None, blocking=True)

        self._lock.release()

    def get_value(self) -> int:
        """Get the value of the internal counter without acquiring the Semaphore."""

        self._lock.acquire(block=True, timeout=None)

        try:
            self._lmsg = self._reader.recv_bytes(timeout=None, blocking=False)
        except ChannelEmpty:
            value = 0
        else:
            self._writer.send_bytes(self._lmsg, timeout=None, blocking=False)
            value = int.from_bytes(self._lmsg, byteorder=self._BYTEORDER)

        self._lock.release()

        return value

    def _reset(self) -> None:

        self._reader = ChannelRecvH(self._channel)
        self._writer = ChannelSendH(self._channel, return_mode=WHEN_DEPOSITED)
        self._reader.open()
        self._writer.open()

src/dragon/native/semaphore.py
DragonHPC/dragon
"""The Dragon native value is an ctypes object allocated from shared memory.
The return value is a synchronized wrapper for the object, and the object itself
can be accessed via the value attribute of a Value.

.. code-block:: python

   value = dragon.native.value.Value(typecode_or_type = ctypes.c_char, value = 0)
"""


import sys
import logging
import ctypes
import struct

import dragon
from ..channels import Channel, Message, ChannelError
from ..dtypes import WHEN_DEPOSITED
from ..infrastructure.facts import default_pool_muid_from_index
from ..infrastructure.parameters import this_process
from ..infrastructure.channel_desc import ChannelOptions
from ..localservices.options import ChannelOptions as ShepherdChannelOptions
from ..globalservices.channel import create, get_refcnt, release_refcnt

LOGGER = logging.getLogger(__name__)
_DEF_MUID = default_pool_muid_from_index(this_process.index)

_TYPECODE_TO_TYPE = {
    "c": ctypes.c_char,
    "u": ctypes.c_wchar,
    "b": ctypes.c_byte,
    "B": ctypes.c_ubyte,
    "h": ctypes.c_short,
    "H": ctypes.c_ushort,
    "i": ctypes.c_int,
    "I": ctypes.c_uint,
    "l": ctypes.c_long,
    "L": ctypes.c_ulong,
    "q": ctypes.c_longlong,
    "Q": ctypes.c_ulonglong,
    "f": ctypes.c_float,
    "d": ctypes.c_double,
}

_SUPPORTED_TYPES = frozenset(_TYPECODE_TO_TYPE.values())


class Value:
    """This class implements Dragon value resides in the Dragon channel."""

    def __getstate__(self):
        return (self._channel.serialize(), self._type)

    def __setstate__(self, state):
        (serialized_bytes, self._type) = state
        self._channel = dragon.channels.Channel.attach(serialized_bytes)
        self._reset()
        get_refcnt(self._channel.cuid)

    def __repr__(self):
        return f"Dragon Value(type={self._type}, value={self.value}, cuid={self._channel.cuid})"

    def __init__(self, typecode_or_type, value: int = 0, m_uid: int = _DEF_MUID):
        """Initialize a value object.
        :param typecode_or_type: the typecode or type is returned from the dictionary, _TYPECODE_TO_TYPE
        :type typecode_or_type: str or ctypes, required
        :param value: the value for the object
        :type value: int, optional
        :param m_uid: memory pool to create the channel in and message to write value and typecode_or_type in managed memory, defaults to _DEF_MUID
        :type m_uid: int, optional
        """

        if isinstance(typecode_or_type, str):
            if typecode_or_type in _TYPECODE_TO_TYPE.keys():
                typecode_or_type = _TYPECODE_TO_TYPE[typecode_or_type]
            else:
                raise AttributeError(f"typecode not found, has to be one of {list(_TYPECODE_TO_TYPE.keys())}")

        else:
            if typecode_or_type not in _SUPPORTED_TYPES:
                raise AttributeError(f"type not found, has to be one of {_SUPPORTED_TYPES}")

        LOGGER.debug(
            f"Init Dragon Native Value with typecode_or_type={typecode_or_type}, value={value}, m_uid={m_uid}"
        )

        # for repr method
        self._muid = m_uid

        sh_channel_options = ShepherdChannelOptions(capacity=1)
        gs_channel_options = ChannelOptions(ref_count=True, local_opts=sh_channel_options)
        descriptor = create(m_uid, options=gs_channel_options)
        self._channel = Channel.attach(descriptor.sdesc)
        self._reset()

        # set value with type that is multiprocessing Value type
        self._type = typecode_or_type
        valbytes = self._value2valbytes(value)

        # create value in shared memory
        mpool = self._channel.get_pool()
        msg = Message.create_alloc(mpool, 8)
        mview = msg.bytes_memview()
        mview[: len(valbytes)] = valbytes

        self._sendh.send(msg)

        msg.destroy()

    def __del__(self):
        try:
            self._sendh.close()
            self._recvh.close()
            self._channel.detach()
            release_refcnt(self._channel.cuid)
        except (AttributeError, ChannelError):
            pass

    @property
    def value(self) -> object:
        """Get the current value.

        :return: The current value
        :rtype: object
        """

        # grab value from the channel
        msg = self._recvh.recv()
        mview = msg.bytes_memview()
        value_info = mview[:].tobytes()

        self._sendh.send(msg)
        msg.destroy()
        return self._valbytes2value(value_info)

    @value.setter
    def value(self, val: object) -> None:
        """Set the value

        :param val: the new value
        :type val: object
        :return: None
        :rtype: NoneType
        """

        # grabs value from the channel
        msg = self._recvh.recv()

        # assigns the changed value to bytes string
        valbytes = self._value2valbytes(val)

        # puts value back into memory
        mview = msg.bytes_memview()
        mview[: len(valbytes)] = valbytes

        self._sendh.send(msg)
        msg.destroy()

    # private methods

    def _reset(self) -> None:
        self._recvh = self._channel.recvh()
        self._sendh = self._channel.sendh(return_mode=WHEN_DEPOSITED)
        self._recvh.open()
        self._sendh.open()

    def _value2valbytes(self, value: object) -> bytes:

        valbytes = b""
        if self._type in [ctypes.c_char, ctypes.c_wchar]:
            # This is for ctypes.c_char as ctypes.c_wchar returns str
            if isinstance(value, (bytes, bytearray)) and self._type is not ctypes.c_wchar:
                valbytes = bytes(value)
            else:
                valbytes = bytes(str(value).encode("utf-8"))
            # if there are more than one bytes in valbytes
            if len(valbytes) > 8:
                raise AttributeError
        # c_float and c_double return float type
        elif self._type in [ctypes.c_float, ctypes.c_double]:
            valbytes = bytes(struct.pack("d", value))
        # int, longlong, and other ctypes return int type
        else:
            valbytes = int(value).to_bytes(8, byteorder=sys.byteorder, signed=True)
        return valbytes

    def _valbytes2value(self, valbytes: bytes) -> object:

        if self._type in [ctypes.c_char, ctypes.c_wchar]:
            value_info = valbytes.decode("utf-8")
            # removes extraneous characters
            value_val = value_info.replace("\x00", "")
            if self._type is ctypes.c_char:
                # c_char returns bytes
                return str(value_val).encode()
            else:
                # c_wchar returns string
                return str(value_val)
        elif self._type in (ctypes.c_float, ctypes.c_double):
            # returns float for the float types
            (float_val,) = struct.unpack("d", valbytes)
            return float(float_val)
        else:
            # returns int for int types
            return int.from_bytes(valbytes, byteorder=sys.byteorder, signed=True)

src/dragon/native/value.py
DragonHPC/dragon
from collections import ChainMap
from collections.abc import Iterable
import logging
import os
import shutil
import subprocess
from typing import Union

from ..channels import Channel
from ..infrastructure import facts
from ..utils import B64
from ..dlogging.util import setup_BE_logging, DragonLoggingServices as dls


def resolve_args(args=None):

    # If no transport agent settings are set in launch parameters
    if not args:
        args = facts.DEFAULT_TRANSPORT_AGENT
    alias = facts.TRANSPORT_AGENT_ALIASES.get(args[0].lower())
    if alias:
        args = alias + args[1:]
    # Resolve executable
    cmd = shutil.which(args[0])
    if cmd:  # Only replace if resolved
        args[0] = cmd
    return args


def start_transport_agent(
    node_index: Union[int, str],
    in_ch_sdesc: B64,
    log_ch_sdesc: B64 = None,
    args: list[str] = None,
    env: dict[str, str] = None,
    gateway_channels: Iterable[Channel] = None,
    infrastructure: bool = False
) -> subprocess.Popen:
    """Start the backend transport agent

    :param node_index: node index for this agent. should be in 0, N-1,
                       where N is number of backend compute nodes
    :type node_index: Union[int, str]
    :param in_ch_sdesc: Channel to be used for incoming infrastructure messages
    :type in_ch_sdesc: B64
    :param log_ch_sdesc: Channel to be used for logging to the Frontend, defaults to None
    :type log_ch_sdesc: B64, optional
    :param args: Command args for starting Transport Agents, defaults to None in which
                 case the TCP transport agent is initiated as default
    :type args: list[str], optional
    :param env: Environment variables, defaults to None
    :type env: dict[str, str], optional
    :param gateway_channels: gateway channels to be set in transport agent's environment.
                             If not set, it is assumed they are already set in the parent
                             environment or the one input, defaults to None
    :type gateway_channels: Iterable[Channel], optional
    :param infrastructure: whether this agent will be used for infrastructure messaging, defaults to False
    :type infrastructure: bool, optional
    :return: Transport agent child process
    :rtype: subprocess.Popen
    """

    _env = ChainMap()
    if env is not None:
        _env.maps.append(env)
    _env.maps.append(os.environ)

    if infrastructure:
        service = dls.LA_BE
    else:
        service = dls.LS

    LOGGER = logging.getLogger(str(service)).getChild('start_transport_agent')
    for key, val in _env.items():
        LOGGER.debug(f'{key}: {val}')

    if gateway_channels is not None:
        for i, ch in enumerate(gateway_channels):
            _env[f'{facts.GW_ENV_PREFIX}{i+1}'] = str(B64(ch.serialize()))

    # TODO: This needs to be tightened up
    _args = resolve_args(args) + [
        str(node_index),
        '--ch-in-sdesc', str(in_ch_sdesc),
        '--log-sdesc', str(log_ch_sdesc)
    ]

    LOGGER.debug(f'Starting transport agent: {_args}')
    return subprocess.Popen(_args, bufsize=0, stdin=subprocess.DEVNULL, env=_env)

src/dragon/transport/__init__.py
DragonHPC/dragon
from __future__ import annotations

import ctypes
import ctypes.util
import enum
import os
import re
import socket
from typing import Any, Union


libc = ctypes.CDLL(ctypes.util.find_library("c"), use_errno=True)


class SIOCGIFFLAGS(enum.IntFlag):
    """Linux interface flags used by the SIOCGIFFLAGS ioctl(2) operation;
    see <linux/if.h> and <net/if.h>.
    """
    IFF_UP          = 1<<0      # Interface is up
    IFF_BROADCAST   = 1<<1      # Broadcast address valid
    IFF_DEBUG       = 1<<2      # Turn on debugging
    IFF_LOOPBACK    = 1<<3      # Is a loopback net
    IFF_POINTOPOINT = 1<<4      # Interface is point-to-point link
    IFF_NOTRAILERS  = 1<<5      # Avoid use of trailers
    IFF_RUNNING     = 1<<6      # Resources allocated
    IFF_NOARP       = 1<<7      # No address resolution protocol
    IFF_PROMISC     = 1<<8      # Receive all packets
    IFF_ALLMULTI    = 1<<9      # Receive all multicast packets
    IFF_MASTER      = 1<<10     # Master of a load balancer
    IFF_SLAVE       = 1<<11     # Slave of a load balancer
    IFF_MULTICAST   = 1<<12     # Supports multicast
    IFF_PORTSEL     = 1<<13     # Can set media type
    IFF_AUTOMEDIA   = 1<<14     # Auto media select active
    IFF_DYNAMIC     = 1<<15     # Dialup device with changing addresses
    IFF_LOWER_UP    = 1<<16
    IFF_DORMANT     = 1<<17
    IFF_ECHO        = 1<<18


class EthernetProtocol(enum.IntEnum):
    """Linux ethernet protocols; see <linux/if_ether.h>.

    The definitive list of ethernet protocols is maintained by the IEEE at
    https://standards-oui.ieee.org/ethertype/eth.txt.
    """
    ETH_P_LOOP      = 0x0060    # Ethernet Loopback packet
    ETH_P_PUP       = 0x0200    # Xerox PUP packet
    ETH_P_PUPAT     = 0x0201    # Xerox PUP Addr Trans packet
    ETH_P_TSN       = 0x22F0    # TSN (IEEE 1722) packet
    ETH_P_ERSPAN2   = 0x22EB    # ERSPAN version 2 (type III)
    ETH_P_IP        = 0x0800    # Internet Protocol packet
    ETH_P_X25       = 0x0805    # CCITT X.25
    ETH_P_ARP       = 0x0806    # Address Resolution packet
    ETH_P_BPQ       = 0x08FF    # G8BPQ AX.25 Ethernet Packet    [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_IEEEPUP   = 0x0a00    # Xerox IEEE802.3 PUP packet
    ETH_P_IEEEPUPAT = 0x0a01    # Xerox IEEE802.3 PUP Addr Trans packet
    ETH_P_BATMAN    = 0x4305    # B.A.T.M.A.N.-Advanced packet [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_DEC       = 0x6000    # DEC Assigned proto
    ETH_P_DNA_DL    = 0x6001    # DEC DNA Dump/Load
    ETH_P_DNA_RC    = 0x6002    # DEC DNA Remote Console
    ETH_P_DNA_RT    = 0x6003    # DEC DNA Routing
    ETH_P_LAT       = 0x6004    # DEC LAT
    ETH_P_DIAG      = 0x6005    # DEC Diagnostics
    ETH_P_CUST      = 0x6006    # DEC Customer use
    ETH_P_SCA       = 0x6007    # DEC Systems Comms Arch
    ETH_P_TEB       = 0x6558    # Trans Ether Bridging
    ETH_P_RARP      = 0x8035    # Reverse Addr Res packet
    ETH_P_ATALK     = 0x809B    # Appletalk DDP
    ETH_P_AARP      = 0x80F3    # Appletalk AARP
    ETH_P_8021Q     = 0x8100    # 802.1Q VLAN Extended Header
    ETH_P_ERSPAN    = 0x88BE    # ERSPAN type II
    ETH_P_IPX       = 0x8137    # IPX over DIX
    ETH_P_IPV6      = 0x86DD    # IPv6 over bluebook
    ETH_P_PAUSE     = 0x8808    # IEEE Pause frames. See 802.3 31B
    ETH_P_SLOW      = 0x8809    # Slow Protocol. See 802.3ad 43B
    ETH_P_WCCP      = 0x883E    # Web-cache coordination protocol defined in draft-wilson-wrec-wccp-v2-00.txt
    ETH_P_MPLS_UC   = 0x8847    # MPLS Unicast traffic
    ETH_P_MPLS_MC   = 0x8848    # MPLS Multicast traffic
    ETH_P_ATMMPOA   = 0x884c    # MultiProtocol Over ATM
    ETH_P_PPP_DISC  = 0x8863    # PPPoE discovery messages
    ETH_P_PPP_SES   = 0x8864    # PPPoE session messages
    ETH_P_LINK_CTL  = 0x886c    # HPNA, wlan link local tunnel
    ETH_P_ATMFATE   = 0x8884    # Frame-based ATM Transport over Ethernet
    ETH_P_PAE       = 0x888E    # Port Access Entity (IEEE 802.1X)
    ETH_P_PROFINET  = 0x8892    # PROFINET
    ETH_P_REALTEK   = 0x8899    # Multiple proprietary protocols
    ETH_P_AOE       = 0x88A2    # ATA over Ethernet
    ETH_P_ETHERCAT  = 0x88A4    # EtherCAT
    ETH_P_8021AD    = 0x88A8    # 802.1ad Service VLAN
    ETH_P_802_EX1   = 0x88B5    # 802.1 Local Experimental 1.
    ETH_P_PREAUTH   = 0x88C7    # 802.11 Preauthentication
    ETH_P_TIPC      = 0x88CA    # TIPC
    ETH_P_LLDP      = 0x88CC    # Link Layer Discovery Protocol
    ETH_P_MRP       = 0x88E3    # Media Redundancy Protocol
    ETH_P_MACSEC    = 0x88E5    # 802.1ae MACsec
    ETH_P_8021AH    = 0x88E7    # 802.1ah Backbone Service Tag
    ETH_P_MVRP      = 0x88F5    # 802.1Q MVRP
    ETH_P_1588      = 0x88F7    # IEEE 1588 Timesync
    ETH_P_NCSI      = 0x88F8    # NCSI protocol
    ETH_P_PRP       = 0x88FB    # IEC 62439-3 PRP/HSRv0
    ETH_P_CFM       = 0x8902    # Connectivity Fault Management
    ETH_P_FCOE      = 0x8906    # Fibre Channel over Ethernet
    ETH_P_IBOE      = 0x8915    # Infiniband over Ethernet
    ETH_P_TDLS      = 0x890D    # TDLS
    ETH_P_FIP       = 0x8914    # FCoE Initialization Protocol
    ETH_P_80221     = 0x8917    # IEEE 802.21 Media Independent Handover Protocol
    ETH_P_HSR       = 0x892F    # IEC 62439-3 HSRv1
    ETH_P_NSH       = 0x894F    # Network Service Header
    ETH_P_LOOPBACK  = 0x9000    # Ethernet loopback packet, per IEEE 802.3
    ETH_P_QINQ1     = 0x9100    # deprecated QinQ VLAN [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_QINQ2     = 0x9200    # deprecated QinQ VLAN [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_QINQ3     = 0x9300    # deprecated QinQ VLAN [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_EDSA      = 0xDADA    # Ethertype DSA [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_DSA_8021Q = 0xDADB    # Fake VLAN Header for DSA [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_DSA_A5PSW = 0xE001    # A5PSW Tag Value [ NOT AN OFFICIALLY REGISTERED ID ]
    ETH_P_IFE       = 0xED3E    # ForCES inter-FE LFB type
    ETH_P_AF_IUCV   = 0xFBFB    # IBM af_iucv [ NOT AN OFFICIALLY REGISTERED ID ]

    # If the value in the ethernet type is more than this value hen the frame is
    # Ethernet II. Else it is 802.3
    ETH_P_802_3_MIN = 0x0600

    # Non DIX types. Won't clash for 1500 types.

    ETH_P_802_3     = 0x0001    # Dummy type for 802.3 frames
    ETH_P_AX25      = 0x0002    # Dummy protocol id for AX.25
    ETH_P_ALL       = 0x0003    # Every packet (be careful!!!)
    ETH_P_802_2     = 0x0004    # 802.2 frames
    ETH_P_SNAP      = 0x0005    # Internal only
    ETH_P_DDCMP     = 0x0006    # DEC DDCMP: Internal only
    ETH_P_WAN_PPP   = 0x0007    # Dummy type for WAN PPP frames
    ETH_P_PPP_MP    = 0x0008    # Dummy type for PPP MP frames
    ETH_P_LOCALTALK = 0x0009    # Localtalk pseudo type
    ETH_P_CAN       = 0x000C    # CAN: Controller Area Network
    ETH_P_CANFD     = 0x000D    # CANFD: CAN flexible data rate
    ETH_P_CANXL     = 0x000E    # CANXL: eXtended frame Length
    ETH_P_PPPTALK   = 0x0010    # Dummy type for Atalk over PPP
    ETH_P_TR_802_2  = 0x0011    # 802.2 frames
    ETH_P_MOBITEX   = 0x0015    # Mobitex (kaz@cafe.net)
    ETH_P_CONTROL   = 0x0016    # Card specific control frames
    ETH_P_IRDA      = 0x0017    # Linux-IrDA
    ETH_P_ECONET    = 0x0018    # Acorn Econet
    ETH_P_HDLC      = 0x0019    # HDLC frames
    ETH_P_ARCNET    = 0x001A    # 1A for ArcNet :-)
    ETH_P_DSA       = 0x001B    # Distributed Switch Arch.
    ETH_P_TRAILER   = 0x001C    # Trailer switch tagging
    ETH_P_PHONET    = 0x00F5    # Nokia Phonet frames
    ETH_P_IEEE802154 = 0x00F6   # IEEE802.15.4 frame
    ETH_P_CAIF      = 0x00F7    # ST-Ericsson CAIF protocol
    ETH_P_XDSA      = 0x00F8    # Multiplexed DSA protocol
    ETH_P_MAP       = 0x00F9    # Qualcomm multiplexing and aggregation protocol
    ETH_P_MCTP      = 0x00FA    # Management component transport protocol packets


class ARPHardware(enum.IntEnum):
    ARPHRD_NETROM   = 0        # from KA9Q: NET/ROM pseudo
    ARPHRD_ETHER    = 1     # Ethernet 10Mbps
    ARPHRD_EETHER   = 2     # Experimental Ethernet
    ARPHRD_AX25     = 3     # AX.25 Level 2
    ARPHRD_PRONET   = 4     # PROnet token ring
    ARPHRD_CHAOS    = 5     # Chaosnet
    ARPHRD_IEEE802  = 6     # IEEE 802.2 Ethernet/TR/TB
    ARPHRD_ARCNET   = 7     # ARCnet
    ARPHRD_APPLETLK = 8     # APPLEtalk
    ARPHRD_DLCI     = 15    # Frame Relay DLCI
    ARPHRD_ATM      = 19    # ATM
    ARPHRD_METRICOM = 23    # Metricom STRIP (new IANA id)
    ARPHRD_IEEE1394 = 24    # IEEE 1394 IPv4 - RFC 2734
    ARPHRD_EUI64    = 27    # EUI-64
    ARPHRD_INFINIBAND = 32  # InfiniBand

    # Dummy types for non ARP hardware
    ARPHRD_SLIP     = 256
    ARPHRD_CSLIP    = 257
    ARPHRD_SLIP6    = 258
    ARPHRD_CSLIP6   = 259
    ARPHRD_RSRVD    = 260   # Notional KISS type
    ARPHRD_ADAPT    = 264
    ARPHRD_ROSE     = 270
    ARPHRD_X25      = 271   # CCITT X.25
    ARPHRD_HWX25    = 272   # Boards with X.25 in firmware
    ARPHRD_CAN      = 280   # Controller Area Network
    ARPHRD_MCTP     = 290
    ARPHRD_PPP      = 512
    ARPHRD_CISCO    = 513   # Cisco HDLC
    ARPHRD_HDLC        = ARPHRD_CISCO
    ARPHRD_LAPB     = 516   # LAPB
    ARPHRD_DDCMP    = 517   # Digital's DDCMP protocol
    ARPHRD_RAWHDLC  = 518   # Raw HDLC
    ARPHRD_RAWIP    = 519   # Raw IP

    ARPHRD_TUNNEL   = 768    # IPIP tunnel
    ARPHRD_TUNNEL6  = 769    # IP6IP6 tunnel
    ARPHRD_FRAD     = 770    # Frame Relay Access Device
    ARPHRD_SKIP     = 771    # SKIP vif
    ARPHRD_LOOPBACK = 772    # Loopback device
    ARPHRD_LOCALTLK = 773    # Localtalk device
    ARPHRD_FDDI     = 774    # Fiber Distributed Data Interface
    ARPHRD_BIF      = 775    # AP1000 BIF
    ARPHRD_SIT      = 776    # sit0 device - IPv6-in-IPv4
    ARPHRD_IPDDP    = 777    # IP over DDP tunneller
    ARPHRD_IPGRE    = 778    # GRE over IP
    ARPHRD_PIMREG   = 779    # PIMSM register interface
    ARPHRD_HIPPI    = 780    # High Performance Parallel Interface
    ARPHRD_ASH      = 781    # Nexus 64Mbps Ash
    ARPHRD_ECONET   = 782    # Acorn Econet
    ARPHRD_IRDA     = 783    # Linux-IrDA
    # ARP works differently on different FC media .. so
    ARPHRD_FCPP     = 784    # Point to point fibrechannel
    ARPHRD_FCAL     = 785    # Fibrechannel arbitrated loop
    ARPHRD_FCPL     = 786    # Fibrechannel public loop
    ARPHRD_FCFABRIC = 787    # Fibrechannel fabric
    # 787->799 reserved for fibrechannel media types
    ARPHRD_IEEE802_TR = 800  # Magic type ident for TR
    ARPHRD_IEEE80211 = 801   # IEEE 802.11
    ARPHRD_IEEE80211_PRISM = 802     # IEEE 802.11 + Prism2 header
    ARPHRD_IEEE80211_RADIOTAP = 803  # IEEE 802.11 + radiotap header
    ARPHRD_IEEE802154 = 804
    ARPHRD_IEEE802154_MONITOR = 805  # IEEE 802.15.4 network monitor

    ARPHRD_PHONET   = 820    # PhoNet media type
    ARPHRD_PHONET_PIPE = 821 # PhoNet pipe header
    ARPHRD_CAIF     = 822    # CAIF media type
    ARPHRD_IP6GRE   = 823    # GRE over IPv6
    ARPHRD_NETLINK  = 824    # Netlink header
    ARPHRD_6LOWPAN  = 825    # IPv6 over LoWPAN
    ARPHRD_VSOCKMON = 826    # Vsock monitor header

    ARPHRD_VOID     = 0xFFFF # Void type, nothing is known
    ARPHRD_NONE     = 0xFFFE # zero header length


class PacketType(enum.IntEnum):
    PACKET_HOST      = 0    # To us
    PACKET_BROADCAST = 1    # To all
    PACKET_MULTICAST = 2    # To group
    PACKET_OTHERHOST = 3    # To someone else
    PACKET_OUTGOING  = 4    # Outgoing of any type
    PACKET_LOOPBACK  = 5    # MC/BRD frame looped back
    PACKET_USER      = 6    # To user space
    PACKET_KERNEL    = 7    # To kernel space
    # Unused, PACKET_FASTROUTE and PACKET_LOOPBACK are invisible to user space
    PACKET_FASTROUTE = 6    # Fastrouted frame


class sockaddr_in(ctypes.Structure):
    _fields_ = [
        ('sin_family', ctypes.c_ushort),        # Always AF_INET
        ('sin_port',   ctypes.c_uint16),
        ('sin_addr',   ctypes.c_byte * 4),
        ('sin_zero',   ctypes.c_byte * 8),      # Ignored
    ]

    def asdict(self) -> dict[str, Any]:
        return {
            'family': socket._intenum_converter(self.sin_family, socket.AddressFamily),
            'port': int(self.sin_port),
            'addr': socket.inet_ntop(self.sin_family, self.sin_addr),
        }


class sockaddr_in6(ctypes.Structure):
    _fields_ = [
        ('sin6_family',   ctypes.c_ushort),     # Always AF_INET6
        ('sin6_port',     ctypes.c_uint16),
        ('sin6_flowinfo', ctypes.c_uint32),
        ('sin6_addr',     ctypes.c_byte * 16),
        ('sin6_scope_id', ctypes.c_uint32),
    ]

    def asdict(self) -> dict[str, Any]:
        return {
            'family': socket._intenum_converter(self.sin6_family, socket.AddressFamily),
            'port': int(self.sin6_port),
            'flowinfo': int(self.sin6_flowinfo),
            'addr': socket.inet_ntop(self.sin6_family, self.sin6_addr),
            'scope_id': int(self.sin6_scope_id),
        }


class sockaddr_ll(ctypes.Structure):
    _fields_ = [
        ('sll_family',   ctypes.c_ushort),      # Always AF_PACKET
        ('sll_protocol', ctypes.c_ushort),      # Physical-layer protocol
        ('sll_ifindex',  ctypes.c_int),         # Interface number
        ('sll_hatype',   ctypes.c_ushort),      # ARP hardware type
        ('sll_pkttype',  ctypes.c_ubyte),       # Packet type
        ('sll_halen',    ctypes.c_ubyte),       # Length of address
        ('sll_addr',     ctypes.c_ubyte * 8),   # Physical-layer address
    ]

    def asdict(self) -> dict[str, Any]:
        return {
            'family': socket._intenum_converter(self.sll_family, socket.AddressFamily),
            'protocol': socket._intenum_converter(socket.ntohs(self.sll_protocol), EthernetProtocol),
            'ifindex': int(self.sll_ifindex),
            'hatype': socket._intenum_converter(self.sll_hatype, ARPHardware),
            'pkttype': socket._intenum_converter(self.sll_pkttype, PacketType),
            'halen': int(self.sll_halen),
            'addr': bytes(self.sll_addr)[:self.sll_halen],
        }


class sockaddr(ctypes.Structure):
    _fields_ = [
        ('sa_family', ctypes.c_ushort),
        ('sa_data',   ctypes.c_byte * 14),
    ]

    af_type = {
        socket.AF_INET: ctypes.POINTER(sockaddr_in),
        socket.AF_INET6: ctypes.POINTER(sockaddr_in6),
        socket.AF_PACKET: ctypes.POINTER(sockaddr_ll),
    }

    def cast(self) -> Union[sockaddr_in, sockaddr_in6, sockaddr_ll]:
        af_type = self.af_type.get(self.sa_family)
        if af_type is None:
            return self
        return ctypes.cast(ctypes.addressof(self), af_type).contents

    def asdict(self) -> dict[str, Any]:
        return {
            'family': socket.AddressFamily(self.sa_family),
            'data': bytes(self.sa_data),
        }


class ifa_ifu(ctypes.Union):
    _fields_ = [
        ('ifu_broadaddr', ctypes.POINTER(sockaddr)), # broadcast address
        ('ifu_dstaddr',   ctypes.POINTER(sockaddr)), # other end of link
    ]


class ifaddrs(ctypes.Structure):
    _anonymous_ = ('ifa_ifu',)

    def __iter__(self):
        yield self
        p = self
        while p.ifa_next:
            p = p.ifa_next.contents
            yield p

    def asdict(self) -> dict[str, Any]:
        ifa = {
            'name': self.ifa_name.decode('utf-8'),
            'flags': SIOCGIFFLAGS(int(self.ifa_flags)),
        }

        addr = self.ifa_addr.contents if self.ifa_addr else None
        if addr is not None:
            ifa['addr'] = addr.cast().asdict()

        if self.ifa_netmask:
            netmask = self.ifa_netmask.contents
            if not netmask.sa_family and addr is not None:
                netmask.sa_family = addr.sa_family
            ifa['netmask'] = netmask.cast().asdict()

        if self.ifa_flags & SIOCGIFFLAGS.IFF_BROADCAST:
            ifa['broadaddr'] = None
            if self.ifu_broadaddr:
                broadaddr = self.ifu_broadaddr.contents
                if not broadaddr.sa_family and addr is not None:
                    broadaddr.sa_family = addr.sa_family
                ifa['broadaddr'] = broadaddr.cast().asdict()

        elif self.ifa_flags & SIOCGIFFLAGS.IFF_POINTOPOINT:
            ifa['dstaddr'] = None
            if self.ifu_dstaddr:
                dstaddr = self.ifu_dstaddr.contents
                if not dstaddr.sa_family and addr is not None:
                    dstaddr.sa_family = addr.sa_family
                ifa['dstaddr'] = dstaddr.cast().asdict()

        # TODO Support data field?
        return ifa

ifaddrs._fields_ = [
    ('ifa_next',    ctypes.POINTER(ifaddrs)),  # Pointer to next struct
    ('ifa_name',    ctypes.c_char_p),          # Interface name
    ('ifa_flags',   ctypes.c_uint),            # Interface flags
    ('ifa_addr',    ctypes.POINTER(sockaddr)), # Interface address
    ('ifa_netmask', ctypes.POINTER(sockaddr)), # Intrfacee netmask
    ('ifa_ifu',     ifa_ifu),
    ('ifa_data',    ctypes.c_void_p),          # Address specific data
]


def getifaddrs() -> list[dict]:
    """Returns a list of dicts that mirror an ifaddrs structure.

    NOTE: Dict items corresponding to sockaddr pointer types in the ifaddrs
    structure (i.e., addr, netmask, broadaddr, dstaddr) will only be present if
    they were set by getifaddrs(3). Do not assume they exist! Nested sockaddr
    mappings contain items specific family, though values have been coerced
    into more appropriate types (e.g., a packed IPv4 address is returned as a
    string in dotted notation). Also, address specific data (i.e., ifa_data) is
    not currently exposed.
    """
    ifap = ctypes.POINTER(ifaddrs)()
    if libc.getifaddrs(ctypes.byref(ifap)) < 0:
        errno = ctypes.get_errno()
        raise OSError(errno, os.strerror(errno))
    try:
        # TODO Convert to ipaddress.IPv4Interface or ipaddress.IPv6Interface
        return [ifa.asdict() for ifa in iter(ifap.contents)]
    finally:
        libc.freeifaddrs(ifap)


class InterfaceAddressFilter:

    def __init__(self, *filters):
        self.filters = list(filters)

    def clear(self):
        self.filters.clear()

    def __call__(self, ifa: dict[str, Any]) -> bool:
        return all(f(ifa) for f in self.filters)

    def require_flags(self, flags: SIOCGIFFLAGS, default: SIOCGIFFLAGS = SIOCGIFFLAGS(0)) -> InterfaceAddressFilter:
        """Include interfaces with all specified flags set."""
        def require_flags_filter(ifa):
            return ifa.get('flags', default) & flags == flags
        self.filters.append(require_flags_filter)
        return self

    def any_flags(self, flags: SIOCGIFFLAGS, default: SIOCGIFFLAGS = SIOCGIFFLAGS(0)) -> InterfaceAddressFilter:
        """Include interfaces with any specified flags set."""
        def any_flags_filter(ifa):
            return ifa.get('flags', default) & flags != 0
        self.filters.append(any_flags_filter)
        return self

    def exclude_flags(self, flags: SIOCGIFFLAGS, default: SIOCGIFFLAGS = SIOCGIFFLAGS(0)) -> InterfaceAddressFilter:
        """Exclude interfaces with any specified flags set."""
        def exclude_flags_filter(ifa):
            return ifa.get('flags', default) & flags == 0
        self.filters.append(exclude_flags_filter)
        return self

    def family(self, *args: socket.AddressFamily) -> InterfaceAddressFilter:
        """Include interfaces with any of the specified address families."""
        if not args:
            return self
        family = set(args)
        def family_filter(ifa):
            addr = ifa.get('addr')
            if addr is None:
                return False
            return addr.get('family') in family
        self.filters.append(family_filter)
        return self

    def name_re(self, *patterns: re.Pattern) -> InterfaceAddressFilter:
        """Include interfaces with a name that matches any of the specified
        regex patterns.
        """
        if not patterns:
            return self
        def name_re_filter(ifa):
            name = ifa.get('name', '')
            return any(p.match(name) is not None for p in patterns)
        self.filters.append(name_re_filter)
        return self

    def name_prefix(self, *prefixes: str) -> InterfaceAddressFilter:
        """Include interfaces with a names that starts with any of the
        specified prefixes.
        """
        if not prefixes:
            return self
        def name_prefix_filter(ifa):
            name = ifa.get('name', '')
            return any(name.startswith(p) for p in prefixes)
        self.filters.append(name_prefix_filter)
        return self

    def up_and_running(self) -> InterfaceAddressFilter:
        return self.require_flags(SIOCGIFFLAGS.IFF_UP|SIOCGIFFLAGS.IFF_RUNNING)

    def no_loopback(self) -> InterfaceAddressFilter:
        return self.exclude_flags(SIOCGIFFLAGS.IFF_LOOPBACK)

    def af_inet(self, inet6: bool = True, loopback: bool = False) -> InterfaceAddressFilter:
        family = {socket.AddressFamily.AF_INET}
        if inet6:
            family.add(socket.AddressFamily.AF_INET6)
        f = self.family(*family)
        if not loopback:
            f = f.no_loopback()
        return f

    def high_speed_if_names(self) -> InterfaceAddressFilter:
        return self.name_prefix('hsn', 'ipogif', 'ib')


def _json_preprocess(obj):
    if isinstance(obj, (list, tuple)):
        return [_json_preprocess(item) for item in obj]
    if isinstance(obj, dict):
        return {_json_preprocess(k): _json_preprocess(v) for k, v in obj.items()}
    if isinstance(obj, (bytes, bytearray)):
        return ':'.join(f'{b:02x}' for b in obj)
    if isinstance(obj, enum.IntFlag):
        return str(obj).removeprefix(obj.__class__.__name__).lstrip('.').split('|')
    if isinstance(obj, enum.IntEnum):
        return obj.name
    return obj


def main(args=None):
    import argparse
    from functools import reduce
    import json
    import operator

    parser = argparse.ArgumentParser(
        prog='ifaddrs',
        description="Get network interface addresses (Linux only).",
        epilog="""
            Output is a JSON encoded list of objects corresponding to an
            ifaddrs structure; see getifaddr(3) for more information.
        """,
    )
    parser.add_argument(
        'prefix',
        metavar='PREFIX',
        nargs='*',
        help="show interfaces with names matching prefix",
    )
    parser.add_argument(
        '-n', '--name',
        metavar='REGEX',
        action='append',
        default=[],
        help="show interfaces with names matching regex pattern",
    )

    addr_group = parser.add_argument_group('filter on address')
    addr_group.add_argument(
        '-f', '--family',
        metavar='FAMILY',
        choices=set(socket.AddressFamily.__members__),
        action='append',
        default=[],
        help="socket address family: %(choices)s",
    )
    addr_group.add_argument(
        '--ip',
        action='store_true',
        default=False,
        help="show interfaces with IPv4 addresses, same as -f AF_INET",
    )
    addr_group.add_argument(
        '--ip6',
        action='store_true',
        default=False,
        help="show interfaces with IPv6 addresses, same as -f AF_INET6",
    )

    _flags = set(SIOCGIFFLAGS.__members__)
    flags_group = parser.add_argument_group(
        'filter on flags',
        description=f"Available SIOCGIFFLAGS flags: {', '.join(_flags)}",
    )
    flags_group.add_argument(
        '-r', '--require-flags',
        metavar='FLAG',
        choices=_flags,
        action='append',
        default=[],
        help="all flags required to be set",
    )
    flags_group.add_argument(
        '-a', '--any-flags',
        metavar='FLAG',
        choices=_flags,
        action='append',
        default=[],
        help="any flags which may be set",
    )
    flags_group.add_argument(
        '-x', '--exclude-flags',
        metavar='FLAG',
        choices=_flags,
        action='append',
        default=[],
        help="all flags that must not be set",
    )
    flags_group.add_argument(
        '--up',
        action='store_true',
        default=False,
        help="show up interfaces, same as -r IFF_UP",
    )
    flags_group.add_argument(
        '--running',
        action='store_true',
        default=False,
        help="show running interfaces, same as -r IFF_RUNNING",
    )
    flags_group.add_argument(
        '--no-loopback',
        dest='loopback',
        action='store_false',
        default=True,
        help="exclude loopback interfaces, same as -x IFF_LOOPBACK",
    )

    args = parser.parse_args(args)

    args.family = {getattr(socket.AddressFamily, af) for af in args.family}
    if args.ip:
        args.family.add(socket.AddressFamily.AF_INET)
    if args.ip6:
        args.family.add(socket.AddressFamily.AF_INET6)

    args.require_flags = reduce(
        operator.or_,
        {getattr(SIOCGIFFLAGS, f) for f in args.require_flags},
        SIOCGIFFLAGS(0),
    )
    args.any_flags = reduce(
        operator.or_,
        {getattr(SIOCGIFFLAGS, f) for f in args.any_flags},
        SIOCGIFFLAGS(0),
    )
    args.exclude_flags = reduce(
        operator.or_,
        {getattr(SIOCGIFFLAGS, f) for f in args.exclude_flags},
        SIOCGIFFLAGS(0),
    )
    if args.up:
        args.require_flags |= SIOCGIFFLAGS.IFF_UP
    if args.running:
        args.require_flags |= SIOCGIFFLAGS.IFF_RUNNING
    if not args.loopback:
        args.exclude_flags |= SIOCGIFFLAGS.IFF_LOOPBACK

    ifaddr_filter = InterfaceAddressFilter()
    if args.family:
        ifaddr_filter.family(*args.family)
    if args.require_flags:
        ifaddr_filter.require_flags(args.require_flags)
    if args.any_flags:
        ifaddr_filter.any_flags(args.any_flags)
    if args.exclude_flags:
        ifaddr_filter.exclude_flags(args.exclude_flags)
    if args.name:
        ifaddr_filter.name_re(*args.name)
    if args.prefix:
        ifaddr_filter.name_prefix(*args.prefix)

    ifaddrs = list(filter(ifaddr_filter, getifaddrs()))
    print(json.dumps(_json_preprocess(ifaddrs)))


if __name__ == '__main__':
    main()

src/dragon/transport/ifaddrs.py
DragonHPC/dragon
import subprocess
import logging
import shutil
import shlex

from ...infrastructure.facts import TransportAgentOptions, DEFAULT_OVERLAY_NETWORK_PORT, TRANSPORT_AGENT_ALIASES
from ...utils import B64
from ...dlogging.util import DragonLoggingServices as dls


def start_overlay_network(ch_in_sdesc: B64,
                          ch_out_sdesc: B64,
                          log_sdesc: B64,
                          host_ids: list[str],
                          ip_addrs: list[str],
                          frontend: bool = False,
                          env: dict = None):
    """Entry point to start the launcher overlay network

    :param ch_in_sdesc: Channel to be used for infrastructure messages incoming
                        from launcher frontend/backend. One component of "CTRL ch"
    :type ch_in_sdesc: B64
    :param ch_out_sdesc: Channel to be used for infrastructure messages outgoing
                         to launcher frontend/backend. One component of "CTRL ch"
    :type ch_out_sdesc: B64
    :param log_sdesc: Channel to be used for sending logging messages to the Frontend
    :type log_sdesc: B64
    :param host_ids: List of host IDs for all nodes that will be communicated with
    :type host_ids: list[str]
    :param ip_addrs: List of IP address for all nodes that will be communicated with
    :type ip_addrs: list[str]
    :param frontend: Is parent process the launcher frontned defaults to False
    :type frontend: bool, optional
    :param env: environment variables. Should container gateway channels, defaults to None
    :type env: dict, optional
    :return: Popen object for the overlay network process
    :rtype: subprocess.Popen
    """

    log = logging.getLogger(dls.ON).getChild('start_overlay_network')

    # If no transport agent settings are set in launch parameters
    args = shlex.split(str(TransportAgentOptions.TCP))
    alias = TRANSPORT_AGENT_ALIASES.get(args[0].lower())
    if alias:
        args = alias + args[1:]
    # Resolve executable
    cmd = shutil.which(args[0])
    if cmd:  # Only replace if resolved
        args[0] = cmd

    args = args \
        + ['0', '--infrastructure',
           '--ip-addrs'] + ip_addrs \
        + ['--host-ids'] + host_ids \
        + ['--ch-in-sdesc', str(ch_in_sdesc),
           '--ch-out-sdesc', str(ch_out_sdesc),
           '--port', str(DEFAULT_OVERLAY_NETWORK_PORT),
           '--log-sdesc', str(log_sdesc)]

    if frontend:
        args.append('--frontend')

    log.info(f'overlay args: {args}')
    # start_new_session to detach this process and avoid SIGINT going to it
    return subprocess.Popen(args,
                            bufsize=0,
                            env=env,
                            start_new_session=True)

src/dragon/transport/overlay/__init__.py
DragonHPC/dragon
"""Dragon TCP transport agent.

.. note::

    This library uses `asyncio` to provide concurrent I/O, which means an event
    loop is required. Calls to functions which lead to blocking I/O should be
    run via a thread-pool executor (e.g., via `asyncio.to_thread`) to avoid
    blocking the event loop.  By default, this package uses the default executor
    which may be customized by `asyncio.loop.set_default_executor`.

The Dragon TCP transport agent is comprised of four primary components:

*   `StreamTransport` -- Used by `Server` and `Client` instances to reliably
    communicate with peers. Responsible for establishing and managing
    connections as well as outstanding requests.

*   `Server` -- Handles `Request` messages from clients and replies with
    `Response` messages as appropriate.

*   `Client` -- Processes `GatewayMessages` to fulfill *send*, *receive*, and
    *poll* operations on remote channels.

*   *(TODO)* `control.Manager` -- Manages the agents main thread of execution
    and handles control commands from local services.

.. rubric:: Local Transport Agent

For testing purposes it may be convenient to implement a transport agent
that only sends messages to itself. The key difference is to use `Transport`
instead of `StreamTransport`::

    import asyncio

    from dragon.transport.tcp import cancel_all_tasks, Client, Server, Transport, LOOPBACK_ADDRESS_IPv4

    async def agent(channels):
        # Create a Transport instance to handle Request/Response I/O.
        # By default, Transport._handle_send() implements a loopback
        # strategy by calling Tranport._handle_recv().
        transport = Transport(LOOPBACK_ADDRESS_IPv4)

        # Create a Server instance to process Requests (i.e., remote
        # channel operations).
        server = Server(transport)

        # Using the default transport, the node mapping should
        # ensure every address is the transport's address, for
        # consistency.
        nodes = defaultdict(lambda: transport.addr)

        # Create Client instances for each gateway channel to send
        # Requests to the Server and process Responses in order to
        # complete gateway messages. This example assumes the gateway channels
        # have been created and the corresponding serialized channel
        # descriptors are passed in via the `channels` list.
        clients = [Client(ch_sd, transport, nodes) for ch_sd in channels]

        # NOTE: Transport is not a TaskMixin as there is nothing that requires a
        # separate thread of execution, so there is no need to call
        # `transport.start()` or `await transport.wait_started()` as is the case
        # with StreamTransport.

        # Start the server
        server.start(name=f'Transport-server')

        # Start each client
        for idx, c in enumerate(clients):
            c.start(name=f'Transport-client-{idx}')

        # At this point, the server and client tasks are running and you
        # can do whatever, but if you need a control loop...
        try:
            while True:
                await asyncio.sleep(5.0)
        finally:
            # Shut down clients and server
            await asyncio.gather(*[c.stop() for c in clients], return_exceptions=True)
            try:
                await server.stop()
            except:
                pass
            # Cancel all remaining tasks for good measure
            await cancel_all_tasks()

    if __name__ == '__main__':
        SERIALIZED_CHANNEL_DESCRIPTORS = [...]
        asyncio.run(agent(SERIALIZED_CHANNEL_DESCRIPTORS))
"""

from .transport import Address, StreamTransport, Transport, \
                       LOOPBACK_ADDRESS_IPv4, LOOPBACK_ADDRESS_IPv6
from .client import Client
from .server import Server
from .task import cancel_all_tasks

src/dragon/transport/tcp/__init__.py
DragonHPC/dragon
import asyncio
import logging
import os
import ssl
from typing import Union

from ...channels import Channel
from ...infrastructure import messages as dmsg
from ...infrastructure.connection import Connection, ConnectionOptions
from ...infrastructure.facts import GW_ENV_PREFIX, DEFAULT_TRANSPORT_PORT, FRONTEND_HOSTID
from ...infrastructure.facts import DEFAULT_OVERLAY_NETWORK_PORT, DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE
from ...launcher.util import next_tag
from ...utils import B64
from ...utils import host_id as get_host_id, set_host_id
from ...dtypes import DEFAULT_WAIT_MODE, IDLE_WAIT

from .agent import Agent
from .task import cancel_all_tasks
from .transport import Address, StreamTransport

FRONTEND_HELP = """FRONTEND_IP specifies the single IP address for the frontend of
the dragon runtime"""

BACKEND_HELP = """BACKEND_IP specifies the list of backend compute node IP addresses
for execution of the dragon runtime backend"""

PORT_HELP = f'''PORT specifies the port to be used for infrastructure communication. By
default, {DEFAULT_OVERLAY_NETWORK_PORT} is used for runtime infrastructure
or {DEFAULT_TRANSPORT_PORT} for backend communiction'''

LOGGER = logging.getLogger('dragon.transport.tcp.__main__')


def single_recv(ch_sdesc: bytes) -> Union[tuple(dmsg.all_message_classes)]:
    ch = Channel.attach(ch_sdesc)
    try:
        with Connection(inbound_initializer=ch) as conn:
            return dmsg.parse(conn.recv())
    finally:
        ch.detach()


def open_connection(ch_in_sdesc: bytes = None, ch_out_sdesc: bytes = None) -> Connection:
    if ch_in_sdesc is None and ch_out_sdesc is None:
        raise ValueError('Requires at least one serialized channel descriptor')
    ch_in = ch_out = conn = None
    if ch_in_sdesc is not None:
        ch_in = Channel.attach(ch_in_sdesc)
    if ch_out_sdesc is not None:
        ch_out = Channel.attach(ch_out_sdesc)
    conn = Connection(
        inbound_initializer=ch_in,
        outbound_initializer=ch_out,
        options=ConnectionOptions(creation_policy=ConnectionOptions.CreationPolicy.EXTERNALLY_MANAGED),
    )
    conn.ghost = True
    conn.open()
    return conn


def close_connection(conn: Connection) -> None:
    conn.close()
    if conn.inbound_chan is not None:
        conn.inbound_chan.detach()
    if conn.outbound_chan is not None:
        conn.outbound_chan.detach()


async def tcp_transport_agent(node_index: str = None,
                              ch_in_sdesc: B64 = None,
                              ch_out_sdesc: B64 = None,
                              default_port: int = DEFAULT_TRANSPORT_PORT,
                              tls_enabled: bool = True,
                              cafile: str = None,
                              certfile: str = None,
                              keyfile: str = None,
                              tls_verify: bool = True,
                              max_threads: int = None,
                              host_ids: list[str] = None,
                              ip_addrs: list[str] = None,
                              infrastructure: bool = False,
                              frontend: bool = False) -> None:

    from ...dlogging.util import DragonLoggingServices
    if infrastructure:
        LOGGER = logging.getLogger(DragonLoggingServices.ON).getChild('transport.tcp.__main__')
        if frontend:
            up_msg = dmsg.OverlayPingLA(next_tag())
            halt_msg = dmsg.LAHaltOverlay
        else:
            up_msg = dmsg.OverlayPingBE(next_tag())
            halt_msg = dmsg.BEHaltOverlay

    else:
        LOGGER = logging.getLogger(DragonLoggingServices.TA).getChild('transport.tcp.__main__')
        up_msg = dmsg.TAPingSH(next_tag())
        halt_msg = dmsg.SHHaltTA

    if max_threads is not None:
        # Set a new default executor that limits the maximum number of worker
        # threads to use.
        from concurrent.futures import ThreadPoolExecutor
        executor = ThreadPoolExecutor(
            max_workers=int(max_threads),
            thread_name_prefix='asyncio',
        )
        loop = asyncio.get_running_loop()
        loop.set_default_executor(executor)

    # Initial receive from input channel to get LAChannelsInfo message
    if not infrastructure:
        la_channels_info = await asyncio.to_thread(single_recv, ch_in_sdesc.decode())
        assert isinstance(la_channels_info, dmsg.LAChannelsInfo), "Did not receive LAChannelsInfo from local services"

        LOGGER.debug(f'Node {node_index}: LAChannelsInfo: {la_channels_info.get_sdict()}')

        try:  # Validate LAChannelsInfo nodes_desc structure
            # Verify given node index in LAChannelInfo node descriptors
            assert node_index in la_channels_info.nodes_desc, f'Invalid node index: {node_index}'
        except AssertionError:
            LOGGER.critical(f'Invalid LAChannelsInfo={la_channels_info.get_sdict()}')
            raise

        # Get my node descriptor
        node_desc = la_channels_info.nodes_desc[node_index]

        # TODO Add support for multiple addresses per host ID
        # Create host ID => address mapping from LAChannelsInfo node descriptors
        nodes = {}
        for x in la_channels_info.nodes_desc.values():
            try:
                # XXX Use the first IP address until multiple addresses per host ID
                # XXX are supported.
                addr = Address.from_netloc(str(x.ip_addrs[0]))
                nodes[int(x.host_id)] = Address(addr.host, addr.port or default_port)
            except Exception:
                LOGGER.critical(f'Failed to build node-address mapping from LAChannelsInfo.nodes_desc: {la_channels_info.nodes_desc}')
                raise

    # If infrastructure, create same data structures as above, but use information input
    # since there's no LAChannelsInfo for us
    else:
        nodes = {}
        for ip_addr, host_id in zip(ip_addrs, host_ids):
            try:
                # XXX Use the first IP address until multiple addresses per host ID
                # XXX are supported.
                addr = Address.from_netloc(str(ip_addr))
                nodes[int(host_id)] = Address(addr.host, addr.port or default_port)

            except Exception:
                LOGGER.critical('Failed to build node-address mapping')
                raise

    try:
        # Establish connection for command-and-control
        if not infrastructure:
            ch_out_sdesc = B64.from_str(node_desc.shep_cd)
        control = await asyncio.to_thread(open_connection, ch_in_sdesc.decode(), ch_out_sdesc.decode())
    except Exception:
        if not infrastructure:
            LOGGER.critical(f'Failed to initialize the control connection: Requires valid input and output channel descriptors for node index {node_index}: input_channel={ch_in_sdesc}, output_channel={ch_out_sdesc}, LAChannelsInfo={la_channels_info.get_sdict()}')
        else:
            LOGGER.critical(f'Failed to initialize the control connection: Requires valid input and output channel descriptors: input_channel={ch_in_sdesc}, output_channel={ch_out_sdesc}')
        raise

    try:
        # Create transport
        if not infrastructure:
            transport = StreamTransport(nodes[int(node_desc.host_id)])
            wait_mode = DEFAULT_WAIT_MODE
        else:
            transport = StreamTransport(nodes[int(get_host_id())])
            wait_mode = IDLE_WAIT
        LOGGER.info(f'Created transport: {transport.addr} with wait mode {wait_mode}')

        if tls_enabled:
            client_ssl_ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
            server_ssl_ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
            client_ssl_ctx.minimum_version = ssl.TLSVersion.TLSv1_3
            server_ssl_ctx.minimum_version = ssl.TLSVersion.TLSv1_3
            if cafile is not None:
                server_ssl_ctx.load_verify_locations(cafile)
                client_ssl_ctx.load_verify_locations(cafile)
            if certfile is not None or keyfile is not None:
                server_ssl_ctx.load_cert_chain(certfile, keyfile)
                client_ssl_ctx.load_cert_chain(certfile, keyfile)
            if not tls_verify:
                client_ssl_ctx.check_hostname = False
                client_ssl_ctx.verify_mode = ssl.CERT_NONE
                server_ssl_ctx.verify_mode = ssl.CERT_NONE
            transport.server_options.update(ssl=server_ssl_ctx)
            transport.default_connection_options.update(ssl=client_ssl_ctx)

        # Create agent
        async with Agent(transport, nodes, wait_mode=wait_mode) as agent:
            LOGGER.info('Created agent')

            # Create clients for each gateway channel
            if not infrastructure:
                n_gw = la_channels_info.num_gw_channels
            else:
                n_gw = DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE

            for i in range(n_gw):
                encoded_channel_sdesc = os.environ[GW_ENV_PREFIX + str(i+1)]
                channel_sdesc = B64.from_str(encoded_channel_sdesc).decode()
                agent.new_client(channel_sdesc)
                LOGGER.debug(f'Created client for gateway channel {i}')

            # Send TAPingSH to local services to acknowledge transport is active
            # XXX What is tag?
            await asyncio.to_thread(control.send, up_msg.serialize())
            LOGGER.info(f'Sent {type(up_msg)} reply')

            while agent.is_running():
                LOGGER.debug('Agent is running, polling control')
                # XXX Should we sleep in between polls to avoid having a
                # XXX thread constantly blocking waiting for control input?
                #await asyncio.sleep(5.0)

                # Wait for control message
                ready = await asyncio.to_thread(control.poll, timeout=5.0)
                if ready:
                    # Receive control message
                    data = await asyncio.to_thread(control.recv)
                    msg = dmsg.parse(data)
                    # Process control message
                    if isinstance(msg, halt_msg):
                        LOGGER.info(f'Received {type(msg)}')
                        break
                    LOGGER.warning(f'Received unsupported control message: {msg}')

            LOGGER.debug('Agent is not running, terminating')
            await cancel_all_tasks()

    except Exception as e:
        LOGGER.exception(f'Unable to run TCP agent: {e}')
        await asyncio.to_thread(close_connection, control)
        raise e

    # Hand back our connection to local services so we can signal our exit
    return control


def main(args=None):
    import argparse
    from distutils.util import strtobool

    from ...dlogging.util import DragonLoggingServices, setup_BE_logging
    from ...infrastructure.facts import PROCNAME_OVERLAY_TA, PROCNAME_TCP_TA
    from ...infrastructure.util import range_expr
    from ...utils import set_procname

    parser = argparse.ArgumentParser(description='Runs Dragon TCP socket transport agent (TSTA)')

    parser.add_argument('--infrastructure', action='store_true',
                        help="Set to use for runtime infrastructure rather than backend communiction")
    parser.add_argument('--ip-addrs', metavar='FRONTEND_IP', dest='ip_addrs', nargs='+',
                        type=str, help=FRONTEND_HELP)
    parser.add_argument('--ch-in-sdesc',
                        type=B64.from_str,
                        help="Base64 encoded serialized input channel descriptor")
    parser.add_argument('--ch-out-sdesc',
                        type=B64.from_str,
                        help="Base64 encoded serialized output channel descriptor")
    parser.add_argument('--host-ids', dest='host_ids', type=str, nargs='+')

    parser.add_argument('-p', '--port', type=int, help="Listening port (default: %(default)s)")

    parser.add_argument('--tls', dest='tls_enabled', type=lambda x: bool(strtobool(x)), help=argparse.SUPPRESS)
    parser.add_argument('--no-tls', dest='tls_enabled', action='store_false', help="Disable TLS")
    parser.add_argument('--cafile', metavar='FILE', help="CA certificate for verifying clients")
    parser.add_argument('--certfile', metavar='FILE', help="X509 certificate for server and client auth")
    parser.add_argument('--keyfile', metavar='FILE', help="Private key corresponding to --certfile")
    parser.add_argument('--tls-verify', dest='tls_verify', type=lambda x: bool(strtobool(x)), help=argparse.SUPPRESS)
    parser.add_argument('--no-tls-verify', dest='tls_verify', action='store_false', help="Disable TLS verficiation")

    parser.add_argument('--max-threads', type=int, help="Maximum number of threads to use")
    parser.add_argument('--frontend', action='store_true', help="Whether this is being started on the frontend")

    def get_log_level(level):
        x = str(level).upper()
        if x not in logging._nameToLevel:
            raise argparse.ArgumentTypeError(f'Unknown log level: {level}')
        return x

    parser.add_argument(
        '--log-level',
        choices=logging._nameToLevel.keys(),
        type=get_log_level,
        help="Set log level (default: %(default)s)",
    )
    parser.add_argument(
        '--dragon-logging',
        type=lambda x: bool(strtobool(x)),
        help=argparse.SUPPRESS,
    )
    parser.add_argument(
        '--no-dragon-logging',
        dest='dragon_logging',
        action='store_false',
        help="Disable Dragon logging",
    )
    parser.add_argument(
        '--log-sdesc',
        type=B64.from_str,
        help="Serialized descriptor for Dragon logs",
    )

    if hasattr(os, 'sched_setaffinity'):
        max_cpu = os.cpu_count() - 1
        parser.add_argument(
            '--cpu-list',
            type=lambda s: set(range_expr(s)),
            default=f'0-{max_cpu}' if max_cpu > 0 else '0',
            help="List of processors separated by commas and may include ranges (default: %(default)s)",
        )

    parser.add_argument(
        'node_index',
        help="Node index in LAChannelsInfo message of this agent",
    )

    parser.set_defaults(
        port=os.environ.get('DRAGON_TRANSPORT_TCP_PORT', DEFAULT_TRANSPORT_PORT),

        tls_enabled=os.environ.get('DRAGON_TRANSPORT_TCP_TLS_ENABLED', True),
        cafile=os.environ.get('DRAGON_TRANSPORT_TCP_CAFILE'),
        certfile=os.environ.get('DRAGON_TRANSPORT_TCP_CERTFILE'),
        keyfile=os.environ.get('DRAGON_TRANSPORT_TCP_KEYFILE'),
        tls_verify=os.environ.get('DRAGON_TRANSPORT_TCP_TLS_VERIFY', True),

        max_threads=os.environ.get('DRAGON_TRANSPORT_TCP_MAX_THREADS'),

        log_level=os.environ.get('DRAGON_TRANSPORT_LOG_LEVEL', 'WARNING'),
        dragon_logging=True,
        log_sdesc=None
    )

    args = parser.parse_args()

    if args.frontend:
        set_host_id(FRONTEND_HOSTID)

    if not args.infrastructure:
        set_procname(PROCNAME_TCP_TA)
    else:
        set_procname(PROCNAME_OVERLAY_TA)

    if args.dragon_logging and args.log_sdesc is not None:
        if not args.infrastructure:
            service = DragonLoggingServices.TA
        else:
            service = DragonLoggingServices.ON
        log_level, _ = setup_BE_logging(service=service,
                                        logger_sdesc=args.log_sdesc)
        # We will ignore the argument level because we want a unified setting
        # defined by user arguments to the launcher
        LOGGER.setLevel(log_level)
    else:
        logging.basicConfig(
            format='%(asctime)s %(levelname)-9s %(process)d %(name)s:%(filename)s:%(lineno)d %(message)s',
            level=args.log_level,
        )

    LOGGER.info(f'Command line arguments: {args}')

    if hasattr(os, 'sched_setaffinity'):
        try:
            os.sched_setaffinity(os.getpid(), args.cpu_list)
        except Exception:
            LOGGER.exception('Failed to set CPU affinity')
            raise

    try:
        control = asyncio.run(tcp_transport_agent(
                                       args.node_index,
                                       args.ch_in_sdesc,
                                       args.ch_out_sdesc,
                                       args.port,
                                       args.tls_enabled,
                                       args.cafile,
                                       args.certfile,
                                       args.keyfile,
                                       args.tls_verify,
                                       args.max_threads,
                                       args.host_ids,
                                       args.ip_addrs,
                                       args.infrastructure,
                                       args.frontend
        ))
    except Exception:
        LOGGER.exception('')
        raise

    # Send exit control message
    if args.infrastructure:
        down_msg = dmsg.OverlayHalted(tag=next_tag())
    else:
        down_msg = dmsg.TAHalted(tag=next_tag())

    try:
        control.send(down_msg.serialize())
    except Exception:
        LOGGER.exception('Failed to send TAHalted')
        raise
    else:
        LOGGER.info('Sent TAHalted and exiting')
    finally:
        close_connection(control)


if __name__ == "__main__":
    main()

src/dragon/transport/tcp/__main__.py
DragonHPC/dragon
import asyncio
from collections import defaultdict
from collections.abc import Iterable
import logging
from typing import Optional

from .client import Client
from .server import Server
from .task import TaskMixin
from .transport import Address, LOOPBACK_ADDRESS_IPv4, Transport

from ...dtypes import DEFAULT_WAIT_MODE, WaitMode

LOGGER = logging.getLogger('dragon.transport.tcp.agent')


class Agent:
    """TCP Transport Agent.

    Requires `Transport` instance and a mapping of node indexes to addresses.
    The `Server` will be automatically created unless specified. Use
    `Agent.new_client` to create a `Client`, though an initial set may also be
    given. For example::

        nodes = {
            1: Address.from_netloc('127.0.0.1:8888'),
            2: Address.from_netloc('127.0.0.1:7777'),
        }

        transport = StreamTransport(nodes[1])

        try:
            with Agent(transport, nodes) as agent:
                # Customize control loop
                while True:
                    await asyncio.sleep(5.0)
        finally:
            await cancel_all_tasks()

    """
    _running = False

    def __init__(self,
                 transport: Transport,
                 nodes: dict[int, Address],
                 server: Optional[Server] = None,
                 clients: Optional[Iterable[Client]] = None,
                 wait_mode: WaitMode = DEFAULT_WAIT_MODE):
        if server is None:
            server = Server(transport, wait_mode=wait_mode)
        if clients is None:
            clients = set()
        else:
            clients = set(clients)
        self._transport = transport
        self._server = server
        self._clients = clients
        self._wait_mode = wait_mode
        self.nodes = nodes

    @classmethod
    def loopback(cls, addr: Address = LOOPBACK_ADDRESS_IPv4):
        """Create an `Agent` instance initialized with a loopback transport."""
        transport = Transport(addr)
        nodes = defaultdict(lambda: transport.addr)
        return cls(transport, nodes)

    async def __aenter__(self):
        if not self.is_running():
            await self.start()
        return self

    async def __aexit__(self, *exc):
        if self.is_running():
            await self.stop()
        return False

    def is_running(self):
        return getattr(self, '_running', False)

    async def start(self):
        LOGGER.debug('Agent starting')
        if isinstance(self._transport, TaskMixin):
            await self._start_transport()
        self._start_server()
        self._start_clients()
        self._running = True
        LOGGER.debug('Agent started')

    async def stop(self):
        LOGGER.debug('Agent stopping')
        del self._running
        await self._stop_clients()
        if self._server.is_running():
            await self._stop_server()
        if isinstance(self._transport, TaskMixin):
            if self._transport.is_running():
                await self._stop_transport()
        LOGGER.debug('Agent stopped')

    async def _start_transport(self):
        assert isinstance(self._transport, TaskMixin), "Transport is not a TaskMixin"
        self._transport.start(name=f'Transport-{self._transport.addr}')
        await self._transport.wait_serving()
        LOGGER.debug('Transport started')

    async def _stop_transport(self):
        assert isinstance(self._transport, TaskMixin), "Transport is not a TaskMixin"
        assert self._transport.is_running(), "Transport not running"
        await self._transport.stop()
        LOGGER.debug('Transport stopped')

    def _start_server(self):
        self._server.start(name='Server')
        LOGGER.debug('Server started')

    async def _stop_server(self):
        assert self._server.is_running(), "Server not running"
        await self._server.stop()
        LOGGER.debug('Server stopped')

    def _start_clients(self):
        for c in self._clients:
            if not c.is_running():
                self._start_client(c)
        LOGGER.debug('Clients started')

    async def _stop_clients(self):
        await asyncio.gather(*[self._stop_client(c) for c in self._clients if c.is_running()], return_exceptions=True)
        LOGGER.debug('Clients stopped')

    @staticmethod
    def _start_client(client):
        name = f'Client-{id(client)}'
        client.start(name)
        LOGGER.debug(f'{name} started')

    @staticmethod
    async def _stop_client(client):
        name = f'Client-{id(client)}'
        assert client.is_running(), f'{name} not running'
        try:
            await client.stop()
        except:
            LOGGER.exception(f'{name} errored while stopping')
            raise
        LOGGER.debug(f'{name} stopped')

    def new_client(self, channel_sdesc: bytes):
        """Create a `Client` and start it if the agent is already running.

        :param channel_sdesc: Serialized channel descriptor
        :return: Created `Client`
        """
        client = Client(channel_sdesc, self._transport, self.nodes, wait_mode=self._wait_mode)
        self.add_client(client)
        return client

    def add_client(self, client: Client):
        """Add `Client` and start it if the agent is already running.

        :param Client: Client instance to add
        """
        self._clients.add(client)
        if self.is_running():
            self._start_client(client)


if __name__ == '__main__':
    from .task import cancel_all_tasks

    async def main():
        try:
            async with Agent.loopback() as agent:
                # run control loop
                while agent.is_running():
                    await asyncio.sleep(5.0)
        finally:
            await cancel_all_tasks()

    asyncio.run(main())

src/dragon/transport/tcp/agent.py
DragonHPC/dragon
import asyncio
from functools import partial, singledispatchmethod
import logging
from typing import Optional

from ...channels import GatewayMessage, Channel, ChannelEmpty, ChannelRecvTimeout
from ...dtypes import WaitMode, DEFAULT_WAIT_MODE
from .errno import get_errno, DRAGON_TIMEOUT
from .io import UUIDBytesIO
from .messages import ErrorResponse, EventRequest, EventResponse, \
    RecvRequest, RecvResponse, Request, Response, SendMemoryRequest, \
    SendRequest, SendResponse, SendReturnMode
from .task import TaskMixin, run_forever
from .transport import Address, Transport
from .util import create_msg


LOGGER = logging.getLogger('dragon.transport.tcp.client')


class Client(TaskMixin):
    """Attaches to channel, receives gateway messages, and processes requests
    with the corresponding transport server.
    """

    def __init__(self,
                 channel_sdesc: bytes,
                 transport: Transport,
                 nodes: Optional[dict[int, Address]] = None,
                 wait_mode: WaitMode = DEFAULT_WAIT_MODE):
        self.channel_sdesc = channel_sdesc
        self.transport = transport
        if nodes is None:
            nodes = {}
        self.nodes = nodes
        self._channel = None
        self._recv_handle = None
        self._wait_mode = wait_mode
        self._background_tasks = set()

    @run_forever
    async def run(self) -> None:
        await asyncio.to_thread(self.open)
        try:
            while True:
                msg = await self.recv()
                LOGGER.debug(f'Received gateway message: {msg}')
                try:
                    task = self.process(msg)
                except BaseException as e:
                    LOGGER.exception(f'Error processing gateway message: {msg}')

                    # Try to complete the gateway message with the error
                    try:
                        msg.complete_error(get_errno(e))
                    except:
                        LOGGER.exception('Failed to complete gateway message with error')

                    # Ensure gateway message is destroyed
                    try:
                        msg.destroy()
                    except:
                        LOGGER.exception('Failed to destroy gateway message')

                    # Finally, ignore Exceptions, but not BaseExceptions (e.g., KeyboardInterrupt)
                    try:
                        raise e
                    except Exception:
                        continue
                else:
                    self._background_tasks.add(task)
                    task.add_done_callback(self._background_tasks.discard)
                    LOGGER.debug(f'Processed gateway message: {msg}')

        finally:
            await asyncio.to_thread(self.close)

    def open(self) -> None:
        self._channel = Channel.attach(self.channel_sdesc)
        try:
            self._recv_handle = self._channel.recvh(wait_mode=self._wait_mode)
            self._recv_handle.open()
        except:
            self._channel.detach()
            raise

    def close(self) -> None:
        self._recv_handle.close()
        self._channel.detach()

    async def recv(self, interval: float = 0.1) -> GatewayMessage:
        # TODO Consider revisiting this recv loop
        try:
            msg = self._recv_handle.recv(blocking=False)
        except (ChannelEmpty, ChannelRecvTimeout):
            while True:
                try:
                    msg = await asyncio.to_thread(self._recv_handle.recv, timeout=interval)
                    break
                except (ChannelRecvTimeout, ChannelEmpty):
                    continue
                except asyncio.CancelledError:
                    await asyncio.sleep(interval)
                    raise
        try:
            return GatewayMessage.from_message(msg)
        finally:
            msg.destroy()

    def process(self, msg: GatewayMessage) -> asyncio.Task:
        # Look up destination node address
        try:
            to_addr = self.nodes[msg.target_hostid]
        except KeyError:
            raise ValueError(f'Unknown target host ID: {msg.target_hostid}')
        # Create Request from gateway message
        req = create_request(msg)
        # Send request
        fut = self.transport.write_request(req, to_addr)
        # Handle responses asynchronously. Ordering of responses is
        # intentionally not guaranteed since the remote server is expected to
        # process requests in the order they are received.
        return asyncio.create_task(self.wait_for_response(fut, msg, req))

    async def wait_for_response(self, fut: asyncio.Future, msg: GatewayMessage, req: Request) -> None:
        try:
            # SendResponses to SendRequests with return modes IMMEDIATELY or
            # WHEN_BUFFERED may be received before the request has been sent
            # (if remote) or processed (if local). Waiting for the request's
            # I/O event is required for WHEN_BUFFERED return mode to guarantee
            # local requests are properly processed.
            if isinstance(req, SendRequest) and req.return_mode == SendReturnMode.WHEN_BUFFERED:
                await req._io_event.wait()
            # Wait for response
            try:
                resp, addr = await fut
            except BaseException as e:
                LOGGER.exception(f'Error awaiting response to gateway message: {msg}')
                # Create local ErrorResponse response; seqno=None implies it
                # was never actually sent/received.
                resp = ErrorResponse(None, get_errno(e), f'Error awaiting response to gateway message: {msg}')
                # Inidicate local transport address (i.e., myself)
                addr = self.transport.addr
            # Handle response and complete the gateway message
            try:
                await self.handle_response(resp, addr, msg)
            except BaseException:
                LOGGER.exception(f'Error handling response to gateway message: {msg}')
            finally:
                # Set the I/O event on the response. Critical for properly
                # handling RecvResponse messages when the transport does NOT use
                # write_message() and read_message() to handle message I/O,
                # e.g., Transport.
                resp._io_event.set()
        finally:
            # Destroy the gateway message
            msg.destroy()

    @singledispatchmethod
    async def handle_response(self, resp: Response, addr: Address, msg: GatewayMessage) -> None:
        raise NotImplementedError(f'Unsupported response type: {type(resp)}')

    @handle_response.register
    async def _(self, resp: ErrorResponse, addr: Address, msg: GatewayMessage) -> None:
        """Handle ErrorResponse messages."""
        # Log errors except time outs on events
        if not (msg.is_event_kind and resp.errno == DRAGON_TIMEOUT):
            LOGGER.error(f'Request {resp.seqno} errored: {resp.errno}: {resp.text}')
        # Need to complete the message in order to pass back the response's
        # error code.
        msg.complete_error(resp.errno)

    @handle_response.register
    async def _(self, resp: SendResponse, addr: Address, msg: GatewayMessage) -> None:
        """Handle SendResponse messages."""
        assert msg.is_send_kind
        msg.send_complete()

    @handle_response.register
    async def _(self, resp: RecvResponse, addr: Address, msg: GatewayMessage) -> None:
        """Handle RecvResponse messages."""
        assert msg.is_get_kind
        try:
            msg_recv = await asyncio.to_thread(create_msg, resp.payload, self._channel, msg.deadline, msg.get_dest_mem_descr_ser)
        except BaseException as e:
            try:
                msg.complete_error(get_errno(e))
            except:
                LOGGER.exception('Failed to complete gateway message with error')
            raise e
        msg.get_complete(msg_recv)

    @handle_response.register
    async def _(self, resp: EventResponse, addr: Address, msg: GatewayMessage) -> None:
        """Handle EventResponse messages."""
        assert msg.is_event_kind
        msg.event_complete(resp.result, resp.errno)


def create_request(msg: GatewayMessage) -> Request:
    if msg.is_send_kind:
        if msg.is_send_return_immediately:
            send_return_mode = SendReturnMode.IMMEDIATELY
        elif msg.is_send_return_when_buffered:
            send_return_mode = SendReturnMode.WHEN_BUFFERED
        elif msg.is_send_return_when_deposited:
            send_return_mode = SendReturnMode.WHEN_DEPOSITED
        elif msg.is_send_return_when_received:
            send_return_mode = SendReturnMode.WHEN_RECEIVED
        else:
            raise ValueError('Unsupported send return mode')
        sendhid = UUIDBytesIO.decode(msg.send_payload_message_attr_sendhid)
        payload = msg.send_payload_message
        mem_sd = msg.send_dest_mem_descr_ser
        if mem_sd is None:
            cls = partial(SendRequest, return_mode=send_return_mode, sendhid=sendhid, payload=payload)
        else:
            cls = partial(SendMemoryRequest, return_mode=send_return_mode, sendhid=sendhid, payload=payload, mem_sd=mem_sd)
    elif msg.is_get_kind:
        cls = RecvRequest
    elif msg.is_event_kind:
        cls = partial(EventRequest, mask=msg.event_mask)
    else:
        raise ValueError('Unknown kind of gateway message')
    req = cls(seqno=None, timeout=0.0, channel_sd=msg.target_ch_ser)
    req.deadline = msg.deadline
    return req

src/dragon/transport/tcp/client.py
DragonHPC/dragon
"""Module contains dragonError_t values, for convenience."""

from functools import singledispatch
from dragon.rc import DragonError


DRAGON_SUCCESS = DragonError.SUCCESS
DRAGON_TIMEOUT = DragonError.TIMEOUT
DRAGON_FAILURE = DragonError.FAILURE
DRAGON_INVALID_ARGUMENT = DragonError.INVALID_ARGUMENT
DRAGON_NOT_IMPLEMENTED = DragonError.NOT_IMPLEMENTED
DRAGON_INTERNAL_MALLOC_FAIL = DragonError.INTERNAL_MALLOC_FAIL


def get_errno(e):
    return getattr(e, 'lib_err', _errno(e))


@singledispatch
def _errno(e):
    return DRAGON_FAILURE

@_errno.register
def _(e: ValueError):
    return DRAGON_INVALID_ARGUMENT

@_errno.register
def _(e: NotImplementedError):
    return DRAGON_NOT_IMPLEMENTED

@_errno.register
def _(e: TimeoutError):
    return DRAGON_TIMEOUT

src/dragon/transport/tcp/errno.py
DragonHPC/dragon
"""Read and write objects from `asyncio-streams`."""

import asyncio
from enum import Enum
from ipaddress import ip_address, IPv4Address, IPv6Address
from struct import Struct
from typing import Any, Callable, Optional, Union
from uuid import UUID, uuid4


class FixedBytesIO:
    """Base class that reads and writes fixed size bytes from/to streams."""

    size: int
    """The number of bytes to read or write."""

    async def read(self, reader: asyncio.StreamReader) -> bytes:
        """Read bytes from stream.

        :param reader: Stream to read
        :return: Exactly `.size` bytes
        :raises: See `asyncio.StreamReader.readexactly` for possible exceptions
        """
        data = await reader.readexactly(self.size)
        return data

    def write(self, writer: asyncio.StreamWriter, data: bytes) -> None:
        """Write bytes to stream.

        :param writer: Stream to write
        :param data: Bytes to write
        :raises AssertionError: If `data` is not `.size` bytes
        """
        assert len(data) == self.size
        writer.write(data)


class CodableIO(FixedBytesIO):
    """Encode and decode fixed-bytes I/O."""

    encode: Callable[..., Any]
    """Encode bytes before writing."""

    decode: Callable[[Any], Any]
    """Decode bytes after reading."""

    async def read(self, reader: asyncio.StreamReader) -> Any:
        """Read decoded object(s) from stream.

        :param reader: Stream to read
        :return: Decoded object
        """
        data = await super().read(reader)
        return self.decode(data)

    def write(self, writer: asyncio.StreamWriter, *args: Any) -> None:
        """Write encoded object(s) to stream.

        :param writer: Stream to write
        :param args: Objects to encode
        """
        data = self.encode(*args)
        super().write(writer, data)


class StructIO(Struct, FixedBytesIO):
    """Unpack and pack `Struct` objects from/to streams.

    :param format: `Struct` format string
    """

    async def read(self, reader: asyncio.StreamReader) -> Any:
        """Read bytes from stream and unpack as `Struct`.

        :param reader: Stream to read
        :return: Unpacked object
        """
        packed_data = await super().read(reader)
        data = self.unpack(packed_data)
        return data if len(data) > 1 else data[0]

    def write(self, writer: asyncio.StreamWriter, *args: Any) -> None:
        """Pack objects as `Struct` and write bytes to stream.

        :param writer: Stream to write
        :param args: Objects to pack
        """
        packed_data = self.pack(*args)
        super().write(writer, packed_data)


class VariableBytesIO(StructIO):
    """Read and write a variable number of bytes from/to streams.

    Uses a simple length-based encoding scheme where the total number of bytes
    to read/write is first read/written from/to a stream. Instances must
    provide the `Struct` format string for encoding/decoding the number of
    bytes.

    :param format: `Struct` format string; restricts the maximum number of bytes which may be read/written
    """

    async def read(self, reader: asyncio.StreamReader) -> bytes:
        """Read variable number of bytes from stream.

        :param reader: Stream to read
        :return: Bytes read
        :raises: See `asyncio.StreamReader.readexactly` for possible exceptions
        """
        num_bytes = await super().read(reader)
        data = await reader.readexactly(num_bytes)
        return data

    def write(self, writer: asyncio.StreamWriter, data: bytes) -> None:
        """Write variable number of bytes to stream.

        :param writer: Stream to write
        :param data: Bytes to write
        """
        super().write(writer, len(data))
        writer.write(data)


class VariableTextIO(CodableIO, VariableBytesIO):
    """Encode and decode variable length text from/to streams.

    :param format: `Struct` format string; restricts the maximum number of bytes which may be read/written before/after decoding/encoding.
    """

    codec = 'utf-8'
    """Codec for encoding/decoding."""

    @classmethod
    def encode(cls, s: str) -> bytes:
        """Encode string to bytes."""
        return s.encode(cls.codec)

    @classmethod
    def decode(cls, s: bytes) -> str:
        """Decode bytes as string."""
        return s.decode(cls.codec)


class UUIDBytesIO(CodableIO):
    """Encode and decode UUIDs from/to streams."""

    # UUID is 16 bytes
    size = len(uuid4().bytes)

    @classmethod
    def encode(cls, u: UUID) -> bytes:
        """Encode UUID to bytes."""
        return u.bytes

    @classmethod
    def decode(cls, s: bytes) -> UUID:
        return UUID(bytes=s)


class IPAddressIO(CodableIO):
    """Base class to read and write IP addresses."""

    @staticmethod
    def encode(ip: Union[IPv4Address, IPv6Address]) -> bytes:
        """Encode an `ipaddress.IPv4Address` or `ipaddress.IPv6Address` as *packed* bytes."""
        return ip.packed

    decode = staticmethod(ip_address)
    """Decode bytes as an `ipaddress.IPv4Address` or `ipaddress.IPv6Address`."""


class IPv4AddressIO(IPAddressIO):
    """Read and write `ipaddress.IPv4Address` instances."""
    size = len(ip_address('127.0.0.1').packed)


class IPv6AddressIO(IPAddressIO):
    """Read and write `ipaddress.IPv6Address` instances."""
    size = len(ip_address('::1').packed)


class EnumIO(CodableIO, StructIO):
    """Read and write an `enum.Enum` packed into a minimum-sized integer value
    based on the maximum `enum.Enum` value.

    :param enum_cls: `enum.Enum` subclass to read and write
    :param format: If not specified, the smallest possible `Struct` format is automatically selected.
    """

    @staticmethod
    def encode(e: Enum) -> int:
        """Encode an `enum.Enum` value as an integer."""
        return int(e.value)

    def __init__(self, enum_cls: Enum, format: Optional[str] = None):
        if format is None:
            max_value = max(e.value for e in enum_cls)
            if max_value < 1<<8:
                format = '!B'
            elif max_value < 1<<16:
                format = '!H'
            elif max_value < 1<<32:
                format = '!L'
            elif max_value < 1<<64:
                format = '!Q'
            else:
                raise ValueError(f'Maximum value of {enum_cls} is out of range')
        super().__init__(format)
        self.decode = enum_cls
        """Decode `int` value as `enum.Enum` instance."""

src/dragon/transport/tcp/io.py
DragonHPC/dragon
import asyncio
from dataclasses import dataclass, field
from enum import Enum
from functools import cache
from inspect import isclass
from ipaddress import IPv4Address, IPv6Address
import logging
import math
import time
from typing import Annotated, get_args, get_origin, get_type_hints, Optional
from uuid import UUID

from .io import *
from .util import seconds_remaining


LOGGER = logging.getLogger('dragon.transport.tcp.messages')


uint8 = Annotated[int, StructIO('!B')]
""":class:`Int <int>` object packed into an unsigned 8-bit integer value."""

uint16 = Annotated[int, StructIO('!H')]
""":class:`Int <int>` object packed into an unsigned 16-bit integer value."""

uint64 = Annotated[int, StructIO('!Q')]
""":class:`Int <int>` object packed into an unsigned 64-bit integer value."""

float32 = Annotated[float, StructIO('!f')]
""":class:`Float <float>` object packed into a signed 32-bit floating point value."""

uuid_bytes = Annotated[UUID, UUIDBytesIO()]
""":class:`UUID` object packed into a 16 byte value."""

varbytes_short = Annotated[bytes, VariableBytesIO('!H')]
""":class:`Bytes <bytes>` object packed into at most (2**16 - 1) variable length bytes."""

varbytes = Annotated[bytes, VariableBytesIO('!Q')]
""":class:`Bytes <bytes>` object packed into at most (2**64 - 1) variable length bytes."""

vartext = Annotated[str, VariableTextIO('!Q')]
""":class:`String <str>` object encoded in utf-8 and packed into at most (2**64 - 1) variable length bytes."""

ipv4addr = Annotated[IPv4Address, IPv4AddressIO()]
""":class:`IPv4Address <ipaddress.IPv4Address>` object packed into 4 bytes."""

ipv6addr = Annotated[IPv6Address, IPv6AddressIO()]
""":class:`IPv6Address <ipaddress.IPv6Address>` object packed into 16 bytes."""


#SendReturnMode = Enum('SendReturnMode', 'IMMEDIATELY WHEN_BUFFERED WHEN_DEPOSITED WHEN_RECEIVED')
class SendReturnMode(Enum):
    """Return behavior of `SendRequest` messages."""

    IMMEDIATELY = 1
    """Return as soon as possible.  The source buffer can not be reused safely
    until some other user-coded sychronization ensures the message was received
    by another process. The source buffer can never be reused by the sender when
    transfer-of-ownership is specified on the send.
    """

    WHEN_BUFFERED = 2
    """Return as soon as the source buffer can be reused by the sender, except
    when transfer-of-ownership is specified on the send.
    """

    WHEN_DEPOSITED = 3
    """Return as soon as the message is deposited into the destination channel.
    The source buffer can be reused by the sender except when
    transfer-of-ownership is specified on the send.
    """

    WHEN_RECEIVED = 4
    """Return as soon as the message is received by some process out of the
    target channel.  The source buffer can be reused by the sender except when
    transfer-of-ownership is specified on the send.
    """


send_return_mode = Annotated[SendReturnMode, EnumIO(SendReturnMode)]


@dataclass
class TransportMessage:
    """Base class for messages to be sent or received using a `Transport`.

    .. rubric:: Type ID

    Subclasses must define a unique ``typeid`` parameter in order to be
    transmittable. For example, `Hello` is defined with a ``typeid`` of
    ``\\x40``. The ``typeid`` is the first byte of all transmitted messages and
    informs the receiver about what type of message to read.

    .. rubric:: Serialization

    The `read_message` and `write_message` functions (de)serialize a
    `TransportMessage` according to its ``__annotations__``, which is guaranteed
    to be an ordered mapping, in class declaration order.  The
    `dataclasses.dataclass` decorator makes it convenient to add common methods
    to a class in coordination with its ``__annotations__`` as dataclass fields.
    Transmittable fields must be annotated with an instance of `FixedBytesIO`,
    which provides the ``read`` and ``write`` methods for serializing the
    annotated type from an `asyncio.StreamReader` or to an
    `asyncio.StreamWriter`, respectively.

    :example: To define a transmittable integer field, serialized as an unsigned
    64-bit integer::

            from dataclasses import dataclass
            from typing import Annotated

            from dragon.transport.tcp.io import StructIO

            @dataclass
            class CustomMessage(TransportMessage):
                custom_field: Annotated[int, StructIO('!Q')]

    """

    _types = {}
    """Maps type ID to subclass.

    This dictionary is automatically populated when subclasses are defined with
    the ``typeid`` parameter. Subclasses of `TransportMessage` that do **not**
    define ``typeid``, and hence are not in this mapping, are not transmittable.
    """

    def __init_subclass__(cls, /, *, typeid: Optional[bytes] = None, **kwargs):
        super().__init_subclass__(**kwargs)
        if typeid is not None:
            assert isinstance(typeid, bytes)
            assert len(typeid) == 1
            assert typeid not in cls._types
            cls._typeid = typeid
            cls._types[typeid] = cls

    @property
    def _io_event(self):
        try:
            return self.__io_event
        except AttributeError:
            self.__io_event = asyncio.Event()
        return self.__io_event


@dataclass
class Hello(TransportMessage, typeid=b'\x40'):
    """Handshake messasge to communicate corresponding IPv4 server address."""

    host: ipv4addr
    """Host `ipaddress.IPv4Address`"""

    port: uint16
    """Port number"""


@dataclass
class Hello6(Hello, typeid=b'\x60'):
    """Handshake message to communicate corresponding IPv6 server address."""

    host: ipv6addr
    """Host `ipaddress.IPv6Address`"""


@dataclass
class SequencedTransportMessage(TransportMessage):
    """Base class for messages with sequnce numbers."""

    seqno: uint64
    """Sequence number"""


@dataclass
class Request(SequencedTransportMessage):
    """Base class for requests."""

    # Concerning absolute deadlines versus relative timeout.
    #
    # A GatewayMessage includes a deadline, which is based on the timeout
    # argument passed to ChannelSendH.send() and ChannelRecvH.recv() on the
    # client side. Currently, this is the only information the transport has
    # available concerning how to set the timeout argument to
    # ChannelSendH.send() or ChannelRecvH.recv() on the server side.
    #
    # The transport has two basic options, neither ideal:
    #
    #   1. A Request includes the deadline as-is and the server side computes
    #      the timeout argument as "seconods remaining" before the deadline
    #      expires. This assumes clocks are synced and any skew will impact the
    #      computed timeout.
    #
    #   2. A Request includes the timeout, which is computed on the client side
    #      prior to send. The fidelity of the deadline computed at the server is
    #      impacted by network and queuing (on both the client and server)
    #      latency.
    #
    # So, pick your poison. The current implementation is (2). The baseline
    # Request.timestamp is not sent to the server. Instead servers will evaluate
    # the deadline against their own monotonic clock.

    timestamp: float = field(default_factory=time.monotonic, init=False, repr=False, compare=False)
    """Base timestamp from which to evaluate `timeout`."""

    timeout: float32
    """Request timeout relative to Request.timestamp.

    Use math.inf to indicate the request never times out. This is consistent
    with the sentinel value used in a GatewayMessage which has no deadline.

    .. warning::
        Timeouts are approximate since messages incur network and queuing
        latency during tranmission; however, they should be relatively
        unaffected by clock skew between peers.

    .. note::
        A request provides the same semantics as ChannelSendH.send() and
        ChannelRecvH.recv(). Request timeout impacts `blocking` and `timeout`
        parameters to those functions as follows:

        - abs(Request.timeout) == math.inf  => blocking=True, timeout=None
        - Request.timeout == 0 => blocking=True, timeout=0 (same as blocking=False)
        - Request.timeout > 0  => blocking=True, timeout=Request.timeout

        Recall that blocking=True and timeout<0 will raise an exception and
        therefore isn't a combination the transport needs to worry about when
        interpretting the deadline of a GatewayMessage.

    """

    channel_sd: varbytes_short
    """Serialized target channel descriptor suitable for passing to
    `dragon.channels.Channel.attach`.
    """

    @property
    def deadline(self):
        """Request deadline based on Request.timeout relative to Request.timestamp."""
        return self.timestamp + self.timeout

    @deadline.setter
    def deadline(self, value):
        """Set the Request's timeout and timestamp based on the specified
        deadline and the current monotonic clock.

        .. note::
            `req.deadline = req.deadline` effectively normalizes the existing
            timeout against the current clock.
        """
        self.timeout, self.timestamp = seconds_remaining(value, _inf=math.inf)


@dataclass
class SendRequest(Request, typeid=b'\x01'):
    """Request corresponding to a gateway message with kind *send*."""

    return_mode: send_return_mode
    """Send return mode."""

    sendhid: uuid_bytes
    """The ``sendhid`` attribute on a gateway message; a 16 byte UUID.  Send
    concurrency is limited based on target channel and the send handle ID.
    """

    payload: varbytes
    """Message payload."""


@dataclass
class SendMemoryRequest(SendRequest, typeid=b'\x02'):
    """Request corresponding to a gateway message with kind
    ``DRAGON_GATEWAY_MESSAGE_SEND`` where a serialized memory descriptor on the
    remote peer for the payload is available.
    """

    mem_sd: varbytes_short
    """Serialized memory descriptor on the remote peer."""


class RecvRequest(Request, typeid=b'\x03'):
    """Request corresponding to a gateway message with kind
    ``DRAGON_GATEWAY_MESSAGE_GET``.
    """
    pass


@dataclass
class EventRequest(Request, typeid=b'\x04'):
    """Request corresponding to a gateway message with kind
    ```DRAGON_GATEWAY_MESSAGE_EVENT``.
    """

    mask: uint16
    """Event mask argument to `dragon.channels.Channel.poll`."""


class Response(SequencedTransportMessage):
    """Base class for responses."""
    pass


@dataclass
class ErrorResponse(Response, typeid=b'\xff'):
    """Response when the server experiences an error processing a `Request`."""

    errno: uint16
    """Dragon error number."""

    text: vartext
    """Error string."""


class SendResponse(Response, typeid=b'\xfe'):
    """Response to a `SendRequest`.

    .. seealso::
        Depending on the `SendReturnMode`, a corresponding `SendResponse` may
        be sent by the client prior to the `SendRequest` actually being sent
        or processed at the target server.
    """
    pass


@dataclass
class RecvResponse(Response, typeid=b'\xfc'):
    """Response to a `RecvRequest`."""

    payload: varbytes
    """Message payload."""


@dataclass
class EventResponse(Response, typeid=b'\xfb'):
    """Response to an `EventRequest`."""

    errno: uint16
    """Dragon error number."""

    result: uint64
    """Return value from `dragon.channels.Channel.poll`."""


@cache
def _get_io_annotations(cls):
    """Returns dictionary of object (e.g., class) ``__annotations__`` which
    include a `FixedBytesIO` instance.

    :param cls: Object to be inspected for *IO* ``__annotations__``
    """
    D = {}
    for name, hints in get_type_hints(cls, include_extras=True).items():
        origin = get_origin(hints)
        if not isclass(origin):
            continue
        if not issubclass(origin, Annotated):
            continue
        _, *annotations = get_args(hints)
        for obj in annotations:
            if isinstance(obj, FixedBytesIO):
                D[name] = obj
    return D


async def read_message(reader: asyncio.StreamReader) -> TransportMessage:
    """Read a message from the specified stream.

    See `TransportMessage` for more information about how messages are
    serialized based on annotations.

    :param reader: Stream to read
    :return: Message read
    """
    typeid = await reader.readexactly(1)
    try:
        cls = TransportMessage._types[typeid]
    except KeyError:
        raise NotImplementedError(f'Unknown message type identifier: {typeid}')
    args = []
    for name, io in _get_io_annotations(cls).items():
        try:
            item = await io.read(reader)
        except:
            LOGGER.exception(f'Error reading message attribute: {cls.__name__}.{name}')
            raise
        else:
            args.append(item)
    msg = cls(*args)
    # Set the I/O event since the message was read from a stream
    msg._io_event.set()
    return msg


async def write_message(writer: asyncio.StreamWriter, msg: TransportMessage) -> None:
    """Write message to specified stream.

    See `TransportMessage` for more information about how messages are
    deserialized based on annotations.

    :param writer: Stream to write
    :param msg: Message to write
    """
    cls = type(msg)
    try:
        typeid = cls._typeid
    except AttributeError:
        raise NotImplementedError(f'TransportMessage type not writable: {cls}')
    writer.write(typeid)
    for name, io in _get_io_annotations(cls).items():
        try:
            io.write(writer, getattr(msg, name))
        except:
            LOGGER.exception(f'Error writing message attribute: {cls.__name__}.{name}')
            raise
    await writer.drain()
    # Set the I/O event since the message has been written to a stream
    msg._io_event.set()

src/dragon/transport/tcp/messages.py
DragonHPC/dragon
import asyncio
from collections import defaultdict
from functools import singledispatchmethod
import logging
import sys
import traceback

from .errno import get_errno, DRAGON_TIMEOUT
from .messages import ErrorResponse, \
                      EventRequest, EventResponse, \
                      RecvRequest, RecvResponse, \
                      SendRequest, SendResponse, SendReturnMode
from .task import TaskMixin, run_forever
from .transport import Address, Transport
from .util import poll_channel, recv_msg, send_msg

from ...dtypes import DEFAULT_WAIT_MODE, WaitMode

LOGGER = logging.getLogger('dragon.transport.tcp.server')


class Server(TaskMixin):
    """Process `Request` messages from clients and reply with `Response`
    messages as appropriate.
    """

    IDLE_PROCESS_TIMEOUT = 60.0

    def __init__(self,
                 transport: Transport,
                 wait_mode: WaitMode = DEFAULT_WAIT_MODE):
        self.transport = transport
        self._process_tasks = {}
        self._requests = defaultdict(asyncio.Queue)
        self._background_tasks = set()
        self._wait_mode = wait_mode

    @run_forever
    async def run(self):
        while True:
            req, addr = await self.transport.read_request()
            if isinstance(req, SendRequest):
                key = (req.channel_sd, req.sendhid)
            elif isinstance(req, RecvRequest):
                key = (req.channel_sd, RecvRequest)
            else:
                key = (req.channel_sd, EventRequest)

            self._requests[key].put_nowait((req, addr))
            self._ensure_process_task(key)

    def _ensure_process_task(self, key):
        # Check if send task is still running
        try:
            task = self._process_tasks[key]
        except KeyError:
            task = None
        else:
            if task.done():
                # Remove reference to task
                del self._process_tasks[key]
                task = None
        if task is None:
            # (Re-)Start send task
            self._process_tasks[key] = asyncio.create_task(self.process(key))

    @run_forever
    async def process(self, key):
        while True:
            try:
                req, addr = await asyncio.wait_for(self._requests[key].get(), self.IDLE_PROCESS_TIMEOUT)
            except asyncio.TimeoutError:
                if self._requests[key].empty():
                    # Idle timeout and still no requests to process, shutdown
                    del self._requests[key]
                    # Comment left here because these are helpful in debugging deadlock situations.
                    # LOGGER.debug(f'Removing queue for key={key}')
                    # LOGGER.debug(f'Exiting Task {self._process_tasks[key]}')
                    # LOGGER.debug(f'Here are the remaining Queue sizes')
                    # LOGGER.debug(f'{[self._requests[key].qsize() for key in self._requests]}')
                    # LOGGER.debug(f'Here are the remaining Tasks')
                    del self._process_tasks[key]
                    # LOGGER.debug(f'{[self._process_tasks[key] for key in self._process_tasks]}')
                    break
                continue

            # Dispatch request to corresponding handler
            try:
                await self.handle_request(req, addr)
            except BaseException as e:
                # Send ErrorResponse
                ex_type, ex_value, ex_tb = sys.exc_info()
                tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
                resp = ErrorResponse(req.seqno, get_errno(e), f'Error while handling request: {tb_str}')
                self.transport.write_response(resp, addr)

                # Ignore Exceptions but not BaseExceptions
                try:
                    raise
                except Exception:
                    pass
            finally:
                # Set the I/O event on the request. Critical for properly
                # handling SendRequest messages when the transport does NOT use
                # write_message() and read_message() to handle message I/O,
                # e.g., Transport. See Client.wait_for_response().
                req._io_event.set()

    @singledispatchmethod
    async def handle_request(self, req, addr: Address) -> None:
        raise NotImplementedError(f'Unsupported message type: {type(req)}')

    @handle_request.register
    async def _(self, req: SendRequest, addr: Address) -> None:
        await asyncio.to_thread(
            send_msg,
            req.channel_sd,
            req.payload,
            req.deadline,
            getattr(req, 'mem_sd', None),
            copy_on_send=not req._io_event.is_set(),
            wait_mode=self._wait_mode
        )
        if req.return_mode in (SendReturnMode.WHEN_DEPOSITED, SendReturnMode.WHEN_RECEIVED):
            resp = SendResponse(req.seqno)
            self.transport.write_response(resp, addr)

    @handle_request.register
    async def _(self, req: RecvRequest, addr: Address) -> None:
        msg = await asyncio.to_thread(recv_msg, req.channel_sd, req.deadline, wait_mode=self._wait_mode)
        task = None  # Ensure task is defined for use in exception handler
        destroy_callback = lambda t: msg.destroy()
        try:
            # Create the response. Note the use of a memory view into the
            # message payload! This means we need to take special care to not
            # destroy the message until AFTER the response has been sent.
            resp = RecvResponse(req.seqno, msg.bytes_memview())
            # Create a background task to destroy the message after the
            # response is sent.
            task = asyncio.create_task(resp._io_event.wait())
            self._background_tasks.add(task)
            task.add_done_callback(self._background_tasks.discard)
            # Add a callback to destroy the received message after it has been
            # sent.
            task.add_done_callback(destroy_callback)
            # Write the response
            self.transport.write_response(resp, addr)
        except:
            # Remove the destroy_callback since we're going to immediately
            # destroy it.
            try:
                if isinstance(task, asyncio.Task):
                    task.remove_done_callback(destroy_callback)
            except:
                pass
            msg.destroy()
            raise

    @handle_request.register
    async def _(self, req: EventRequest, addr: Address) -> None:
        ok, result = await asyncio.to_thread(poll_channel, req.channel_sd, req.mask, req.deadline)
        if ok:
            # Although the result should always be an int when ok is True,
            # explicitly return 0 when result is None just to be safe.
            resp = EventResponse(req.seqno, result.rc, result.value)
        else:
            resp = ErrorResponse(req.seqno, DRAGON_TIMEOUT, 'Poll timed out')
        self.transport.write_response(resp, addr)

src/dragon/transport/tcp/server.py
DragonHPC/dragon
import asyncio
from functools import partial, wraps
import logging
from typing import Optional


LOGGER = logging.getLogger('dragon.transport.tcp.task')


def run_forever(func=None, /, *, ignore=(Exception,)):
    """Ignores exceptions and continues to yield.

    Any unignored exception will terminate the generator and log the
    exception. Note that asyncio.exceptions.CanncelledError exceptions are
    not logged.

    If no exception is raised, it is assumed the calling context intended
    to terminate.
    """

    if func is None:
        # We're being called as @run_forever() with parens
        return partial(run_forever, ignore=ignore)

    assert asyncio.iscoroutinefunction(func)

    @wraps(func)
    async def wrapper(*args, **kwds):
        while True:
            try:
                await func(*args, **kwds)
            except ignore:
                LOGGER.exception(f'Uncaught exception, continuing task: {asyncio.current_task()}')
                pass
            except asyncio.exceptions.CancelledError:
                LOGGER.debug(f'Task cancelled: {asyncio.current_task()}')
                return
            except:
                LOGGER.exception(f'Uncaught exception, terminating task: {asyncio.current_task()}')
                return
            else:
                LOGGER.debug(f'Successfully terminating: {asyncio.current_task()}')
                return

    return wrapper


async def cancel_all_tasks():
    """Cancels all running tasks, except the current one."""
    tasks = asyncio.all_tasks()
    tasks.remove(asyncio.current_task())
    for t in tasks:
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)


class TaskMixin:
    """Mixin class that provides a conventional interface for managing a task
    corresponding to the ``run()`` method.
    """

    _task = None

    def start(self, name: Optional[str] = None) -> None:
        assert self._task is None
        self._task = asyncio.create_task(self.run(), name=name)

    async def stop(self) -> None:
        assert self._task is not None
        self._task.cancel()
        try:
            await self._task
        except asyncio.CancelledError:
            pass
        finally:
            del self._task

    def is_running(self):
        return self._task is not None and not self._task.done()

    @run_forever
    async def run(self) -> None:
        raise NotImplementedError

src/dragon/transport/tcp/task.py
DragonHPC/dragon
from __future__ import annotations

import asyncio
from collections import ChainMap, defaultdict
from dataclasses import dataclass
from functools import singledispatchmethod
from ipaddress import ip_address, IPv4Address, IPv6Address
from itertools import chain, count
import logging
import os
from typing import IO, Union
from urllib.parse import urlsplit
from warnings import warn
from weakref import WeakSet, WeakValueDictionary

from .errno import *
from .messages import *
from .task import TaskMixin, run_forever
from .util import unget_nowait


LOGGER = logging.getLogger('dragon.transport.tcp.transport')


@dataclass(frozen=True)
class Address:
    """Node address"""

    host: Union[IPv4Address, IPv6Address]
    """Host IP address"""

    port: int = None
    """Port number"""

    @classmethod
    def from_netloc(cls, netloc: str) -> Address:
        """Returns an `Address` from a *netloc* string.

        Uses `urllib.parse.urlsplit` to extract host and port information from
        the given string. The extracted hostname is expected to be a valid IPv4
        or IPv6 address and will be validated.

        :param netloc: String of the form ``host:port``
        :raise ValueError:
            If host and port cannot be extracted from the specified string or
            cannot be coerced into appropriate types, IP address and integer
            respectively.

        :example:
            >>> Address.from_netloc('127.0.0.1')
            >>> Address.from_netloc('[127.0.0.1]')
            >>> Address.from_netloc('[::1]')
            >>> Address.from_netloc('127.0.0.1:8888')
            >>> Address.from_netloc('[127.0.0.1]:8888')
            >>> Address.from_netloc('[::1]:8888')

        .. seealso:: :rfc:`2732` *Format for Literal IPv6 Addresses in URL's*

        """
        if not netloc:
            raise ValueError('A valid hostname or IPv4 or IPv6 address is required.')
        info = urlsplit('tcp://' + str(netloc))
        host = ip_address(info.hostname)
        return cls(host, info.port)

    def __str__(self):
        """Returns netloc string."""
        if isinstance(self.host, IPv6Address):
            host = f'[{self.host}]'
        else:
            host = str(self.host)
        if self.port is None:
            return host
        return f'{host}:{self.port}'

    def hello(self) -> Union[Hello, Hello6]:
        """Returns corresponding Hello message for this address."""
        assert self.port is not None
        cls = Hello6 if isinstance(self.host, IPv6Address) else Hello
        return cls(self.host, self.port)

    @classmethod
    def from_hello(cls, hello: Union[Hello, Hello6]) -> Address:
        """Returns Address from `Hello` message."""
        return cls(ip_address(hello.host), int(hello.port))

    async def do_handshake(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> Address:
        """Exchange Hello messages with peer.

        :param reader: Stream for messages from peer
        :param writer: Stream for messages to peer
        :return: Peer address
        """
        await write_message(writer, self.hello())
        resp = await read_message(reader)
        assert isinstance(resp, Hello)
        return type(self).from_hello(resp)


LOOPBACK_ADDRESS_IPv4 = Address.from_netloc('127.0.0.1')
LOOPBACK_ADDRESS_IPv6 = Address.from_netloc('[::1]')


class Transport:
    """Base class for TCP transports; provides loopback."""

    def __init__(self, addr: Address):
        """
        :param addr: Server address to accept peer connections.
        """
        self.addr = addr
        """Server address to accept peer connections."""
        self._seqno = count(start=1)
        self._responses = {}
        self._last_request = defaultdict(int)
        self._requests = asyncio.Queue()

    def write_request(self, req: Request, addr: Address, /) -> asyncio.Future:
        """Send a `Request`.

        :param req: `Request` message to send
        :param addr: Recipient `Address`
        :return:
            `asyncio.Future` which will be set to (`Response`, `Address`) pair
            when the corresponding response is available.

        """
        LOGGER.debug(f'Writing request to {addr}: {req}')
        # Set seqno before queuing or handling request
        req.seqno = next(self._seqno)
        # Create future that will be completed when the response is received
        fut = self._responses[req.seqno] = asyncio.Future()
        try:
            self._handle_send(req, addr)
        except:
            # Clean up the future response
            try:
                del self._responses[req.seqno]
            except KeyError:
                pass
            raise
        return fut

    def write_response(self, resp: Response, addr: Address, /) -> None:
        """Send a `Response`.

        :param resp: `Response` message to send
        :param addr: Recipient `Address`
        """
        LOGGER.debug(f'Writing response to {addr}: {resp}')
        self._handle_send(resp, addr)

    def _handle_send(self, msg: TransportMessage, addr: Address, /) -> None:
        # Reminder: Subclasses (e.g., StreamTransport) may rely on this
        # implementation to deliver messages to itself. As a result, any
        # additional overhead introduced here may impact subclasses.
        assert addr == self.addr

        # Even though the message is not actually sent anywhere; but rather,
        # passed immediately to the receive handler, it still needs to process
        # the SendResponse for SendRequests with IMMEDIATELY or WHEN_BUFFERED
        # return modes.
        if isinstance(msg, SendRequest) and msg.return_mode in (SendReturnMode.IMMEDIATELY, SendReturnMode.WHEN_BUFFERED):
            resp = SendResponse(msg.seqno)
            self._handle_recv(resp, self.addr)

        # Important: Do NOT call msg._io_event.set() here! The message is never
        # actually written to and read from a stream, so Server and Client are
        # responsible for properly setting the I/O event when the message has
        # been handled.  Otherwise the sender may misbehave, e.g., prematurely
        # destroy the corresponding Dragon message!

        # Immediately handle receiving the message.
        self._handle_recv(msg, self.addr)

    @singledispatchmethod
    def _handle_recv(self, msg: TransportMessage, addr: Address, /) -> None:
        """Handle a message received from the specified address."""
        # Default handler immediately sends error response
        LOGGER.error(f'Received unsupported type of message from {addr}: {type(msg)}')
        # An unsupported message most likely does not have a sequence number,
        # so default to seqno=0 in that case, which should NOT be used
        # by any valid transport since sequence numbers start at 1.
        resp = ErrorResponse(getattr(msg, 'seqno', 0), DRAGON_NOT_IMPLEMENTED, f'Unsupported type of message: {type(msg)}')
        self.write_response(resp, addr)
        # This handler is a dead end; no client or server will ever process
        # this transport message. As a result, we need to ensure the I/O event
        # is set just in case something else is waiting on it.
        msg._io_event.set()

    @_handle_recv.register
    def _(self, resp: Response, addr: Address, /) -> None:
        """Handle a `Response`."""
        LOGGER.debug(f"Received response from {addr}: {resp}")
        try:
            fut = self._responses.pop(resp.seqno)
        except KeyError:
            LOGGER.warning(f'Received unexpected response: {resp}')
            # May include additional text (e.g., ErrorResponse)
            text = getattr(resp, 'text', '')
            if text:
                LOGGER.warning(resp.text)
        else:
            fut.set_result((resp, addr))

    @_handle_recv.register
    def _(self, req: Request, addr: Address, /) -> None:
        """"Handle a `Request."""
        LOGGER.debug(f"Received request from {addr}: {req}")
        # Protect against inadvertant replays by checking the last seqno from
        # each address.
        if req.seqno <= self._last_request[addr]:
            LOGGER.error('Received duplicate request from {addr}: {req}')
            return
        self._last_request[addr] = req.seqno
        # Queue request so it can be read
        self._requests.put_nowait((req, addr))

    async def read_request(self) -> tuple[Request, Address]:
        """:return: Next `Request` received, when available."""
        req, addr = await self._requests.get()
        LOGGER.debug(f'Read request from {addr}: {req}')
        return req, addr

    async def read_response(self, seqno) -> tuple[Response, Address]:
        try:
            fut = self._responses[seqno]
        except KeyError:
            raise ValueError(f'Invalid sequence number: {seqno}')
        resp, addr = await fut
        return resp, addr


class StreamTransport(Transport, TaskMixin):
    """Asynchronous message transport using `asyncio-streams`.

    `StreamTransport` instances must be given a valid `Address` which indicates
    the host and port on which to accept peer connections::

        addr = Address.from_netloc('127.0.0.1:8888')
        transport = StreamTransport(addr)

    Additional server options to be passed to `asyncio.start_server` may be set
    by updating the `StreamTransport.server_options` dictionary. For example,
    to protect the transport listening socket using TLS::

        import ssl

        # Create server SSLContext
        ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
        # Disable client verification (e.g., enable privacy without authentication)
        ctx.verify_mode = ssl.CERT_NONE

        transport.server_options.update(ssl=ctx)

    Similarly, default options for all transport client connections, which are
    passed to `asyncio.open_connection` may be set by updating the
    `StreamTransport.default_connection_options` dictionary::

        # Create SSLContext for client connections
        ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
        # Disable server verification (e.g., enable privacy without authentication)
        ctx.check_hostname = False
        ctx.verify_mode = ssl.CERT_NONE

        transport.default_connection_options.update(ssl=ctx)

    Connection options for specific peers may be set by updating the
    `StreamTransport.connection_options` dictionary using the corresponding
    `Address` as the key::

        # Disable SSL for connections to specific peer
        peer_addr = Address.from_netloc('127.0.0.1:9999')

        transport.connection_options[peer_addr].update(ssl=None)

    After desired server and connection options have been configured, the
    transport must be started to run in a separate task::

        # Run transport in a separate task with a specific name
        transport.start(name=f'Transport-{transport.addr}')

        # Wait for the transport to start the server and be accepting
        # connections.
        await transport.wait_serving()

    When running, a `StreamTransport` will automatically accept and open
    connections from/to peers as necessary. See `StreamTransport.connect` to
    explicitly create a connection to a peer; although, this is not typically
    needed. Immediately after a connection is accepted or opened, a lightweight
    handshake is required to identify the client's corresponding server address
    before the connection may be considered
    *established* and messages freely exchanged. For more details, see
    `StreamTransport._do_handshake`.

    Once a connection is established, the transport adds the corresponding
    `asyncio.StreamWriter` to a list of available writers and creates a task to
    run a receive-loop for the corresponding `asyncio.StreamReader`, see
    `StreamTransport._do_recv`. A receive-loop task does not terminate unless
    its `asyncio.StreamReader` is closed or encounters an unrecoverable error.
    When a message is queued to be sent (e.g., `StreamTransport.write_request`,
    `StreamTransport.write_response`), the transport ensures a corresponding
    send-loop task exists or creates one using using the most recent
    `asyncio.StreamWriter` available for the recipient, see
    `StreamTransport._do_send`.

    Messages are not tied to underlying connections and do not need to be sent
    and received on corresponding `asyncio.StreamReader` and
    `asyncio.StreamWriter` pairs.  This enables the transport to mitigate
    intermittent connection failures and provide resilient message delivery.  As
    a result, a `Response` is not guaranteed to be received on the
    `asyncio.StreamReader` corresponding to the `asyncio.StreamWriter` on which
    the `Request` was sent.  Conversely, a `Response` is not guaranteed to be
    sent on the `asyncio.StreamWriter` corresponding to the
    `asyncio.StreamReader` on which the `Request` was received.

    Use `StreamTransport.write_request` to queue a `Request` to be sent to a
    peer. It assigns a sequence number to the `Request` and returns an
    `asyncio.Future`, which will be completed when the corresponding `Response`
    is received. It is the responsibility of the caller to wait for the returned
    `asyncio.Future` instance to be completed in order to actually process the
    `Response`. For example::

        from . import messages

        # Create request
        req = messages.SendRequest(...)

        # Send request to recipient
        fut = transport.write_request(req, peer_addr)

        # Wait for response
        resp, resp_addr = await fut

        # Check that the responder was the request recipient
        assert resp_addr == peer_addr

    Use `StreamTransport.read_request` to get the *next* `Request` to process
    and `StreamTransport.write_response` to send a corresponding `Response`. For
    example::

        # Receive request
        req, addr = await transport.read_request()

        # Process request and create response
        resp = messages.SendResponse(req.seqno)

        # Send response
        transport.write_response(resp, addr)

    To stop the transport task and wait for it to complete::

        await transport.stop()

    """

    IDLE_SEND_TIMEOUT = 60.0

    def __init__(self, addr: Address):
        """
        :param addr: Server address to accept peer connections.
        """
        super().__init__(addr)
        self.server_options = {}
        """Server connection options."""
        self.default_connection_options = ChainMap()
        """Default connection options for all client connections."""
        self.connection_options = defaultdict(dict)
        """Connection options for specific Addresses applied on top of the
        default connection options.
        """
        self._server = None
        self._connection_locks = WeakValueDictionary()
        self._writers = defaultdict(list)
        self._recv_tasks = defaultdict(WeakSet)
        self._mailboxes = defaultdict(asyncio.Queue)
        self._send_tasks = WeakValueDictionary()

    @run_forever
    async def run(self):
        """Run the transport server.

        Typically ran in a separate `asyncio.Task` via `StreamTransport.start`.
        """
        self._server = await asyncio.start_server(self.accept_connection, str(self.addr.host), int(self.addr.port), **self.server_options)
        try:
            async with self._server:
                await self._server.serve_forever()
        finally:
            self._server = None

    async def stop(self) -> None:
        await super().stop()
        await self.close()

    async def close(self) -> None:
        # Cancel and clean up recv tasks
        recv_tasks = list(chain.from_iterable(self._recv_tasks.values()))
        if recv_tasks:
            for t in recv_tasks:
                t.cancel()
            await asyncio.gather(*recv_tasks, return_exceptions=True)
        self._recv_tasks.clear()
        # Cancel and clean up send tasks
        send_tasks = list(self._send_tasks.values())
        if send_tasks:
            for t in send_tasks:
                t.cancel()
            await asyncio.gather(*send_tasks, return_exceptions=True)
        self._send_tasks.clear()
        # Close and clean up writers
        closing_writers = [close_writer(w) for w in chain.from_iterable(self._writers.values())]
        if closing_writers:
            await asyncio.gather(*closing_writers, return_exceptions=True)
        self._writers.clear()

    def is_serving(self):
        return self._server is not None and self._server.is_serving()

    async def wait_serving(self, interval: float = 0.1):
        """Wait for the transport server to be running and serving clients.

        :param interval: Sleep interval between checking server status.

        :example:
            Intended to be called after `StreamTransport.start` to ensure the
            server is running before proceeding::

                transport = StreamTransport(Address.from_netloc('127.0.0.1:8888'))
                transport.start()
                await transport.wait_serving()

        """
        while not self.is_serving():
            await asyncio.sleep(interval)

    async def wait_started(self, *args, **kwds):
        warn('Use wait_serving() instead of wait_started()', DeprecationWarning)
        return self.wait_serving(*args, **kwds)

    async def accept_connection(self, reader:asyncio.StreamReader, writer: asyncio.StreamWriter):
        """Accept a client connection.

        Completes handshake with client before starting the receive-loop with
        `reader` and adding the `writer` to the available pool.

        .. note::
            Called by the transport server when accepting new connections. Not
            intended to be called directly.
        """
        addr = await self.addr.do_handshake(reader, writer)
        # XXX Verifying the connection in this way prevents test code from
        # XXX forging addresses, but also impacts legitimate use cases, e.g.,
        # XXX proxies. The only resolution is to real client authentication,
        # XXX e.g., verify the advertised address is a valid SAN in an X.509
        # XXX certificate.
        ## Verify connection is to the advertised host; port is not applicable
        ## from the server side.
        #host, _ = writer.get_extra_info('peername')
        #assert addr.host == ip_address(host)
        self.add_connection(addr, reader, writer)

    async def _open_connection(self, addr: Address, /) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
        """Open a connection to a peer transport server.

        :param addr: Peer address
        :return: `asyncio.StreamReader` and `asyncio.StreamWriter` pair
        """
        opts = self.default_connection_options.new_child(self.connection_options[addr])
        reader, writer = await asyncio.open_connection(str(addr.host), int(addr.port), **opts)
        addr = await self.addr.do_handshake(reader, writer)
        # XXX See comment above in accept_connection() on why this is not a
        # XXX tenable workaround for actual server authentication.
        ## Verify connection is to the advertised address
        #host, port = writer.get_extra_info('peername')
        #assert addr == Address(ip_address(host), int(port))
        self.add_connection(addr, reader, writer)
        return reader, writer

    def add_connection(self, addr: Address, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        self._writers[addr].append(writer)
        task = asyncio.create_task(self._do_recv(addr, reader), name=f'{type(self).__name__}-Receiver-{addr}')
        self._recv_tasks[addr].add(task)
        LOGGER.debug(f'Added connection to {addr}')

    async def connect(self, addr: Address) ->  asyncio.StreamWriter:
        """Returns an open `asyncio.StreamWriter` connected to the specified
        address, opening a new connection if necessary.

        A task will be created to run the receive-loop for the corresponding
        `asyncio.StreamReader` as appropriate.
        """
        # Do not allow self-connections, since messages to oneself are
        # short-circuited and never actually written to a connection.
        assert addr != self.addr
        # XXX Is there a race condition here? If peers attempt to open
        # XXX connections simultaneously, it is not possible to guarantee
        # XXX the resulting StreamWriters are ordered the same at both, i.e.,
        # XXX they may each have different views on which is the "latest". As a
        # XXX result two sockets may still be used.
        try:
            lock = self._connection_locks[addr]
        except KeyError:
            lock = self._connection_locks[addr] = asyncio.Lock()
        async with lock:
            # Search for open writers, opening a new connection as appropriate.
            while True:
                open_writers = [w for w in self._writers[addr] if not w.is_closing()]
                # If there are open writers, then this loop breaks and we only
                # want to keep one, so clearing self._writers[addr] here is
                # fine. Otherwise, if there are no open writers, then clearing
                # self._writers[addr] before waiting for an open connection
                # simply gets rid of closed writers, if any.
                self._writers[addr].clear()
                if open_writers:
                    break
                # Loop after _open_connection() completes to select an open
                # writer instead of choosing the writer returned in case any
                # connections were accepted while waiting.
                try:
                    await self._open_connection(addr)
                except ConnectionRefusedError:
                    LOGGER.exception(f'Failed to open connection to {addr}')
                    continue
            # TODO I now think sorting the open_writers is pointless, remove?
            # Sort open writers to increase the chance that both peers choose
            # to keep the same  attempt to reduce the number of open sockets
            open_writers.sort(key=writer_addrs)
            writer = open_writers.pop()
            self._writers[addr].append(writer)
        # Close remaining open writers.
        await asyncio.gather(*[close_writer(w) for w in open_writers], return_exceptions=True)
        return writer

    @run_forever
    async def _do_recv(self, addr: Address, reader: asyncio.StreamReader):
        while not reader.at_eof():
            try:
                msg = await read_message(reader)
            except asyncio.IncompleteReadError:
                # NOTE: IncompleteReadError.partial has the portion of bytes
                # read; however, since the stream will have gotten EOF there
                # isn't much to do with an incomplete message.
                continue
            LOGGER.debug(f"Received from {addr}: {msg}")
            try:
                self._handle_recv(msg, addr)
            except Exception as e:
                if isinstance(msg, Request):
                    resp = ErrorResponse(msg.seqno, DRAGON_FAILURE, 'Uncaught exception handling message')
                    self.write_response(resp, addr)
                raise

    def _handle_send(self, msg: TransportMessage, addr: Address, /) -> None:
        # Short-circuit delivery when sending to myself
        if addr == self.addr:
            super()._handle_send(msg, addr)
            return
        if isinstance(msg, SendRequest) and msg.return_mode == SendReturnMode.IMMEDIATELY:
            resp = SendResponse(msg.seqno)
            self._handle_recv(resp, self.addr)
        # Enqueue outgoing message
        self._mailboxes[addr].put_nowait(msg)
        # Ensure send task is running
        self._ensure_send_task(addr)

    def _ensure_send_task(self, addr: Address):
        # Check if send task is still running
        try:
            task = self._send_tasks[addr]
        except KeyError:
            task = None
        else:
            if task.done():
                # Remove reference to task
                del self._send_tasks[addr]
                task = None
        if task is None:
            # (Re-)Start send task
            LOGGER.debug(f'Starting sender: {type(self).__name__}-Sender-{addr}')
            self._send_tasks[addr] = asyncio.create_task(self._do_send(addr), name=f'{type(self).__name__}-Sender-{addr}')

    @run_forever
    async def _do_send(self, addr: Address) -> None:
        writer = await self.connect(addr)
        while True:
            try:
                msg = await asyncio.wait_for(self._mailboxes[addr].get(), self.IDLE_SEND_TIMEOUT)
            except asyncio.TimeoutError:
                if self._mailboxes[addr].empty():
                    LOGGER.debug('Cleaning up idle mailbox')
                    # Idle timeout and still no outgoing messages, shutdown
                    del self._mailboxes[addr]
                    break
                continue

            LOGGER.debug(f"Sending to {addr}: {msg}")
            if isinstance(msg, Request):
                # Normalizes Request timeout based on the current monotonic
                # clock. May scale timeout down to 0 if deadline exceeded, but
                # does not preemptively raise DRAGON_TIMEOUT in order to satisfy
                # "try once" semantics.
                msg.deadline = msg.deadline
            try:
                await write_message(writer, msg)
            except ConnectionError:
                LOGGER.exception(f'Connection error while writing to {addr}: {msg}')
                unget_nowait(self._mailboxes[addr], msg)
                # Close existing writer, then get a new one and retry
                await close_writer(writer)
                writer = await self.connect(addr)
            except:
                LOGGER.exception(f'Uncaught error while writing to {addr}: {msg}')
                unget_nowait(self._mailboxes[addr], msg)
                raise
            else:
                self._mailboxes[addr].task_done()
                if isinstance(msg, SendRequest) and msg.return_mode == SendReturnMode.WHEN_BUFFERED:
                    resp = SendResponse(msg.seqno)
                    self._handle_recv(resp, self.addr)


def writer_addrs(writer):
    addrs = [writer.get_extra_info('sockname'), writer.get_extra_info('peername')]
    addrs.sort()
    return addrs


async def close_writer(writer: asyncio.StreamWriter) -> None:
    """Write EOF and close an `asyncio.StreamWriter`."""
    if writer.is_closing():
        return
    if writer.can_write_eof():
        writer.write_eof()
    writer.close()
    try:
        await writer.wait_closed()
    except NotImplementedError:
        pass


async def create_streams(rfile: Optional[IO], wfile: Optional[IO]) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
    """Creates `asyncio.StreamReader` and `asyncio.StreamWriter` instances that
    wrap the specified file-like objects.

    :param rfile: Read file
    :param wfile: Write file

    :example:
        From stdin and stdout::

            import sys

            reader, writer = await create_streams(sys.stdin, sys.stdout)

    """
    if rfile is None and wfile is None:
        raise ValueError('Requires at least one file object')
    loop = asyncio.get_event_loop()
    if rfile is not None:
        reader = asyncio.StreamReader(loop=loop)
        await loop.connect_read_pipe(lambda: asyncio.StreamReaderProtocol(reader, loop=loop), rfile)
    else:
        reader = None
    if wfile is not None:
        w_transport, w_protocol = await loop.connect_write_pipe(lambda: asyncio.streams.FlowControlMixin(loop=loop), wfile)
        writer = asyncio.StreamWriter(w_transport, w_protocol, reader, loop)
    else:
        writer = None
    return reader, writer


async def create_pipe_streams() -> tuple[asyncio.StreamReader, asyncio.StreamWriter]:
    """Creates `asyncio.StreamReader` and `asyncio.StreamWriter` instaces from
    an `os.pipe`.

    :example:
        Useful for testing `read_message` and `write_message` round trip::

            from .messages import SendRequest, read_message, write_message

            reader, writer = create_pipe_streams()

            src_req = SendRequest(...)

            await write_message(writer, src_req)
            await close_writer(writer)

            dest_req = await read_message(reader)

            assert src_req == dest_req

    """
    r, w = os.pipe()
    reader, _ = await create_streams(os.fdopen(r, 'rb'), None)
    _, writer = await create_streams(None, os.fdopen(w, 'wb'))
    return reader, writer


async def create_pipe_connections() -> tuple[tuple[asyncio.StreamReader, asyncio.StreamWriter], tuple[asyncio.StreamReader, asyncio.StreamWriter]]:
    """Creates a connection pair (i.e., two pair of `asyncio.StreamReader` and
    `asyncio.StreamWriter` instances) from `os.pipe` file descriptors.

    :example:
        To simulate actors *Alice* and *Bob* communicating via
        `StreamTransports` without running transport servers or making actual
        network connections::

            from .messages import SendRequest, SendResponse

            # Create transports for Alice and Bob
            alice = StreamsTransport(Address.from_netloc('127.0.0.1:8888))
            bob = StreamTransport(Address.from_netloc('127.0.0.1:9999'))

            # Create connections
            alice_conn, bob_conn = create_pipe_connections()

            # Manually add connections to one another
            alice.add_connection(bob.addr, *alice_conn)
            bob.add_connection(alice.addr, *bob_conn)

            # Alice sends a request to Bob
            fut = alice.write_request(SendRequest(...), bob.addr)

            req, addr = await bob.read_request()

            # Verify request is from Alice
            assert addr == alice.addr

            # Bob processes the request and sends a response
            bob.write_response(SendResponse(req.seqno), addr)

            # Alice waits for the response
            resp, addr = wait fut

            # Verify the response is from Bob
            assert addr == bob.addr

    """
    r0, w0 = os.pipe()
    r1, w1 = os.pipe()
    connA = await create_streams(os.fdopen(r0, 'rb'), os.fdopen(w1, 'wb'))
    connB = await create_streams(os.fdopen(r1, 'rb'), os.fdopen(w0, 'wb'))
    return connA, connB


if __name__ == "__main__":
    import logging
    from random import choice
    from uuid import uuid4
    from .task import cancel_all_tasks


    logging.basicConfig(
        format='%(asctime)-15s %(levelname)-8s %(taskName)-15s %(message)s',
        level=logging.DEBUG,
    )

    def asyncio_task_log_filter(record):
        try:
            task = asyncio.current_task()
        except RuntimeError:
            record.task = None
            record.taskName = None
        else:
            record.task = task
            record.taskName = task.get_name()
        return True

    for h in logging.getLogger().handlers:
        h.addFilter(asyncio_task_log_filter)

    localhost = ip_address('127.0.0.1')

    async def sender(from_addr, to_addr):
        transport = StreamTransport(from_addr)
        transport.start(name=f'Sender-transport-{transport.addr}')
        await transport.wait_serving()
        while True:
            req = SendMemoryRequest(None, timeout=0.5, channel_sd=b'channel desc', return_mode=choice(list(SendReturnMode)), sendhid=uuid4().bytes, payload=b'payload', mem_sd=b'memory desc')
            fut = transport.write_request(req, to_addr)
            LOGGER.info(f"Waiting for response")
            resp, resp_addr = await fut
            LOGGER.info(f"Got response")
            if req.return_mode in (SendReturnMode.IMMEDIATELY, SendReturnMode.WHEN_BUFFERED):
                # Check that the response came from ourselves since the return
                # mode is WHEN_BUFFERED.
                assert resp_addr == from_addr
            else:
                assert resp_addr == to_addr
            assert req.seqno == resp.seqno
            await asyncio.sleep(0.1)
            LOGGER.info(f'Send tasks: {len(transport._send_tasks)}')

    async def receiver(addr):
        transport = StreamTransport(addr)
        transport.start(name=f'Receiver-transport-{transport.addr}')
        await transport.wait_serving()
        while True:
            req, from_addr = await transport.read_request()
            if req.return_mode in (SendReturnMode.WHEN_DEPOSITED, SendReturnMode.WHEN_RECEIVED):
                transport.write_response(SendResponse(req.seqno), from_addr)
            LOGGER.info(f'Receiver send tasks: {len(transport._send_tasks)}')

    async def main():
        to_addr = Address(localhost, 8888)
        for port in range(to_addr.port + 1, to_addr.port + 1 + 1):
            asyncio.create_task(sender(Address(localhost, port), to_addr), name=f'Sender-{port}')
        try:
            await receiver(to_addr)
        except KeyboardInterrupt:
            pass
        finally:
            LOGGER.critical('Shutting down...')
            await cancel_all_tasks()

    asyncio.run(main())

src/dragon/transport/tcp/transport.py
DragonHPC/dragon
from asyncio import Queue, QueueFull
from collections.abc import Generator
from contextlib import contextmanager
import math
import time
from typing import Any, Optional, Union

from ...channels import Channel, ChannelRecvH, ChannelSendH, Message, PollResult
from ...managed_memory import MemoryAlloc, MemoryPool
from ...dtypes import WaitMode, DEFAULT_WAIT_MODE


def unget_nowait(self: Queue, item: Any) -> None:
    """Re-queues an item at the front of a queue, essentially performing the
    opposite of get_nowait().

    Differs from put_nowait() in that the item is put on the front of the queue
    as opposed to the end, and the number of unfinished tasks is not
    incremented.  As a result, do NOT call task_done() for a requeued item.

    Raises QueueFull if the queue is at max size and the item cannot be
    requeued.
    """
    if self.full():
        raise QueueFull
    self._queue.appendleft(item)
    self._wakeup_next(self._getters)


def seconds_remaining(deadline: float, _inf: Optional[float] = None) -> (Optional[float], float):
    """Returns `(N, since)` where `N` is the number of seconds remaining until
    deadline is exceeded, or 0 if the deadline is exceeded, and `since` is the
    baseline timestamp, i.e., `time.monotonic()`.

    `N` is suitable for use as the `timeout` argument to `ChannelSendH.send()`
    and `ChannelRecvH.recv()`, where `blocking=True` (the default). For example,
    when

    - `abs(deadline) == math.inf`, then `N` will be None` and `blocking=True, timeout=None`
    - seconds remaining <= 0, then `N == 0`, which is equivalent to `blocking=False` or `blocking=True, timeout=0`
    - seconds remaining > 0, then `blocking=True, timeout=N`

    .. note::
        When called successively, `N` will decrease to, but not go below, `0`.
        As a result, it supports graceful degradation of the `timeout` argument
        to `ChannelSendH.send()` and `ChannelRecvH.recv()`.

    """
    since = time.monotonic()
    if abs(deadline) == math.inf:
        return _inf, since
    remaining = max(deadline - since, 0.0)
    return remaining, since


def mem_descr_msg(sdesc: bytes, data: bytes) -> Message:
    """Attaches to memory allocation given a serialized memory descriptor and
    writes the specified data. Yields a Dragon message created from the
    corresponding memory allocation.
    """
    # XXX Neither MemoryAlloc.attach() nor Message.create_from_mem() support
    # XXX a timeout or deadline.
    mem = MemoryAlloc.attach(sdesc)
    try:
        if mem.size < len(data):
            raise ValueError(f'Memory allocation too small: {mem}')
        v = mem.get_memview()
        v[:] = data
        return Message.create_from_mem(mem)
    except:
        mem.detach()
        raise


def mem_pool_msg(pool: MemoryPool, data: bytes, deadline: Optional[float] = None) -> Message:
    """Creates Dragon message from data in specific pool's memory allocation.

    Deadline is measured relative to time.monotonic().
    """
    # XXX ChannelSendH.send_bytes() does not call Message.create_alloc() with
    # XXX a timeout or deadline; hence it is re-implemented here.
    timeout, _ = seconds_remaining(deadline)
    msg = Message.create_alloc(pool, len(data), timeout=timeout)
    try:
        v = msg.bytes_memview()
        v[:] = data
    except:
        msg.destroy(True)
        raise
    return msg


def create_msg(data: bytes, channel: Optional[Channel] = None, deadline: Optional[float] = None, sdesc: Optional[bytes] = None) -> Message:
    """Creates a Dragon message from data.

    If a channel is specified, the message will be created from a memory
    allocation in the channel's memory pool with the given deadline. The
    deadline is measured relative to time.monotonic().

    If a serialized memory descriptor is specified, channel and deadline
    arguments are ignored, and the message will be created from the
    corresponding memory allocation.
    """
    if sdesc is None:
        assert channel is not None
        return mem_pool_msg(channel.get_pool(), data, deadline)
    else:
        return mem_descr_msg(sdesc, data)


@contextmanager
def attach_channel(sdesc: bytes) -> Generator[Channel, None, None]:
    ch = Channel.attach(sdesc)
    try:
        yield ch
    finally:
        ch.detach()


@contextmanager
def open_handle(h: Union[ChannelRecvH, ChannelSendH]) -> Generator[Union[ChannelRecvH, ChannelSendH], None, None]:
    h.open()
    try:
        yield h
    finally:
        h.close()


def send_msg(channel_sdesc: bytes,
             payload: bytes,
             deadline: float,
             mem_sd: Optional[bytes] = None,
             copy_on_send: bool = False,
             wait_mode: WaitMode = DEFAULT_WAIT_MODE) -> None:
    # TODO Should we cache attached channels and open handles?
    with attach_channel(channel_sdesc) as ch, open_handle(ch.sendh(wait_mode=wait_mode)) as h:
        msg = create_msg(payload, ch, deadline, mem_sd)
        timeout, _ = seconds_remaining(deadline)

        h.send(
            msg,
            ownership=h.copy_on_send if copy_on_send else h.transfer_ownership_on_send,
            timeout=timeout,
        )


def recv_msg(channel_sdesc: bytes,
             deadline: float,
             wait_mode: WaitMode = DEFAULT_WAIT_MODE) -> Message:
    # TODO Should we cache attached channels and open handles?
    with attach_channel(channel_sdesc) as ch, open_handle(ch.recvh(wait_mode=wait_mode)) as h:
        timeout, _ = seconds_remaining(deadline)
        return h.recv(timeout=timeout)


def poll_channel(channel_sdesc: bytes, event_mask: int, deadline: float) -> (bool, int):
    # TODO Should we cache attached channels and open handles?
    with attach_channel(channel_sdesc) as ch:
        result = PollResult()
        timeout, _ = seconds_remaining(deadline)
        ok = ch.poll(timeout=timeout, event_mask=event_mask, poll_result=result)
        return ok, result

src/dragon/transport/tcp/util.py
DragonHPC/dragon
from cryptography import x509
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.x509.oid import ExtendedKeyUsageOID, ExtensionOID, NameOID
import datetime
import ipaddress


def generate_private_key(key_size=4096):
    return rsa.generate_private_key(public_exponent=65537, key_size=key_size)


def pem_encode_pkcs8(private_key, password=None):
    if password is None:
        enc_alg = serialization.NoEncryption()
    else:
        if not isinstance(password, bytes):
            password = password.encode()
        enc_alg = serialization.BestAvailableEncryption(password)
    return private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=enc_alg,
    )


def certificate_builder(issuer_name, subject_name, public_key, *, serial_number=None, not_valid_before=None, not_valid_after=None):
    if not_valid_before is None:
        not_valid_before = datetime.datetime.utcnow()
    # Do not generate certs valid for more than a year
    max_not_valid_after = not_valid_before.replace(year=not_valid_before.year + 1)
    if not_valid_after is None:
        not_valid_after = max_not_valid_after
    elif not_valid_after > max_not_valid_after:
        raise ValueError(f'Cannot be valid after {max_not_valid_after}')
    if not_valid_before >= not_valid_after:
        raise ValueError(f'Invalid validity period: {not_valid_before} to {not_valid_after}')
    if serial_number is None:
        serial_number = x509.random_serial_number()
    return x509.CertificateBuilder(
        issuer_name=issuer_name,
        subject_name=subject_name,
        public_key=public_key,
        serial_number=serial_number,
        not_valid_before=not_valid_before,
        not_valid_after=not_valid_after,
    )


def add_ca_key_usage_extensions(builder):
    return builder.add_extension(
        x509.KeyUsage(
            digital_signature=True,
            content_commitment=True,
            key_encipherment=True,
            data_encipherment=True,
            key_agreement=True,
            key_cert_sign=True,
            crl_sign=True,
            encipher_only=False,
            decipher_only=False,
        ), critical=True,
    ).add_extension(
        x509.BasicConstraints(ca=True, path_length=0), critical=True,
    )


def add_server_key_usage_extensions(builder):
    return builder.add_extension(
        x509.KeyUsage(
            digital_signature=True,
            content_commitment=False,
            key_encipherment=False,
            data_encipherment=False,
            key_agreement=False,
            key_cert_sign=False,
            crl_sign=False,
            encipher_only=False,
            decipher_only=False,
        ), critical=True,
    ).add_extension(
        x509.BasicConstraints(ca=False, path_length=None), critical=True,
    ).add_extension(
        x509.ExtendedKeyUsage([
            ExtendedKeyUsageOID.SERVER_AUTH,
            ExtendedKeyUsageOID.CLIENT_AUTH,
        ]), critical=False,
    )


def sign(builder, private_key, algorithm=None):
    if algorithm is None:
        algorithm = hashes.SHA512()
    return builder.sign(private_key, algorithm)


def pem_encode(x509_obj):
    return x509_obj.public_bytes(
        encoding=serialization.Encoding.PEM,
    )


def get_common_name(name):
    """Returns common name string from x509.Name"""
    # Get subject common name
    common_name = name.get_attributes_for_oid(NameOID.COMMON_NAME)
    if not common_name:
        raise ValueError('Missing common name attribute')
    if len(common_name) > 1:
        raise ValueError('Contains more than one common name attribute')
    return str(common_name[0].value)


def check_subject_alternative_name(csr):
    san_ext = csr.extensions.get_extension_for_oid(ExtensionOID.SUBJECT_ALTERNATIVE_NAME)
    if san_ext is None:
        raise ValueError('CSR missing Subject Alternative Name extension')
    san = san_ext.value
    # Require the common name to be specified as SAN
    common_name = get_common_name(csr.subject)
    if common_name not in set(map(str, san.get_values_for_type(x509.GeneralName))):
        raise ValueError('Subject common name attribute is not a subject alternative name')
    # TODO: Validate other SAN entries?
    return san


class CertificateAuthority:

    def __init__(self, private_key, certificate):
        self.private_key = private_key
        self.certificate = certificate

    @classmethod
    def generate(cls, name=None):
        if name is None:
            name = x509.Name([
                x509.NameAttribute(NameOID.ORGANIZATION_NAME, u'Hewlett Packard Enterprise Development LP'),
                x509.NameAttribute(NameOID.ORGANIZATIONAL_UNIT_NAME, u'Dragon'),
                x509.NameAttribute(NameOID.COMMON_NAME, u'Dragon Root CA'),
            ])
        private_key = generate_private_key()
        public_key = private_key.public_key()
        builder =  certificate_builder(
            issuer_name=name,
            subject_name=name,
            public_key=public_key,
        )
        builder = add_ca_key_usage_extensions(builder).add_extension(
            x509.SubjectKeyIdentifier.from_public_key(public_key),
            critical=False
        ).add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(public_key),
            critical=False
        )
        cert = sign(builder, private_key)
        return cls(private_key, cert)

    def issue_server_certificate(self, csr, *, valid_for: datetime.timedelta = None):
        not_valid_before = datetime.datetime.utcnow()
        if valid_for is None:
            # one year
            not_valid_after = not_valid_before.replace(year=not_valid_before.year + 1)
        else:
            not_valid_after = not_valid_before + valid_for
        if not csr.is_signature_valid:
            raise ValueError(f'CSR signature invalid {csr}')
        san = check_subject_alternative_name(csr)
        public_key = csr.public_key()
        builder = certificate_builder(
            issuer_name=self.certificate.subject,
            subject_name=csr.subject,
            public_key=public_key,
            not_valid_before=not_valid_before,
            not_valid_after=not_valid_after,
        )
        builder = add_server_key_usage_extensions(builder).add_extension(
            x509.SubjectKeyIdentifier.from_public_key(public_key),
            critical=False
        ).add_extension(
            x509.AuthorityKeyIdentifier.from_issuer_public_key(self.certificate.public_key()),
            critical=False
        ).add_extension(san, critical=False)
        return sign(builder, self.private_key)


def server_csr_builder(subject_alternative_name):
    # Get subject from the first SAN
    subject = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, str(subject_alternative_name[0].value)),
    ])
    return x509.CertificateSigningRequestBuilder().subject_name(subject).add_extension(
        subject_alternative_name, critical=False,
    )


def generate_server_csr(ipaddr, *altnames):
    names = [x509.IPAddress(ipaddress.ip_address(ipaddr))]
    if altnames:
        names.extend(altnames)
    private_key = generate_private_key()
    builder = server_csr_builder(x509.SubjectAlternativeName(names))
    csr = sign(builder, private_key)
    return private_key, csr


def generate_server_self_signed_cert(ipaddr, *altnames, valid_for: datetime.timedelta = None):
    names = [x509.IPAddress(ipaddress.ip_address(ipaddr))]
    if altnames:
        names.extend(altnames)
    not_valid_before = datetime.datetime.utcnow()
    if valid_for is None:
        # one year
        not_valid_after = not_valid_before.replace(year=not_valid_before.year + 1)
    else:
        not_valid_after = not_valid_before + valid_for
    private_key = generate_private_key()
    public_key = private_key.public_key()
    subject = x509.Name([
        x509.NameAttribute(NameOID.COMMON_NAME, str(names[0].value)),
    ])
    builder = certificate_builder(
        issuer_name=subject,
        subject_name=subject,
        public_key=public_key,
        not_valid_before=not_valid_before,
        not_valid_after=not_valid_after,
    )
    builder = add_server_key_usage_extensions(builder).add_extension(
        x509.SubjectKeyIdentifier.from_public_key(public_key),
        critical=False
    ).add_extension(
        x509.AuthorityKeyIdentifier.from_issuer_public_key(public_key),
        critical=False
    ).add_extension(x509.SubjectAlternativeName(names), critical=False)
    cert = sign(builder, private_key)
    return private_key, cert


def dragon_uri(rank, host, port):
    return x509.UniformResourceIdentifier(f'urn:dragon:{rank}:{host}:{port}')


if __name__ == "__main__":
    from tempfile import TemporaryDirectory
    import os

    with TemporaryDirectory() as workdir:

        def writefile(path, contents):
            with open(os.path.join(workdir, path), 'wb') as f:
                f.write(contents)

        def openssl_x509_text(path):
            os.system(f'openssl x509 -noout -text -in {workdir}/{path}')

        # Generate CA
        ca = CertificateAuthority.generate()

        writefile('ca-key.pem', pem_encode_pkcs8(ca.private_key))
        writefile('ca-cert.pem', pem_encode(ca.certificate))

        openssl_x509_text('ca-cert.pem')

        cluster_nodes = [
            ('127.0.0.1', 7575),
            ('127.0.0.1', 7576),
        ]

        for rank, (host, port) in enumerate(cluster_nodes):
            uri = dragon_uri(rank, host, port)
            key, csr = generate_server_csr(host, uri)
            cert = ca.issue_server_certificate(csr)

            writefile(f'key-{rank}.pem', pem_encode_pkcs8(key))
            writefile(f'cert-{rank}.pem', pem_encode(cert))

            openssl_x509_text(f'cert-{rank}.pem')

src/dragon/transport/x509.py
DragonHPC/dragon
#!/usr/bin/env python3
def main():
    infile = open('./return_codes.h','r')
    outfile = open('./return_codes_map.h', 'w')
    outfile.write('/*\n')
    outfile.write('Copyright 2020, 2022 Hewlett Packard Enterprise Development LP\n\n')
    outfile.write('This file was auto-generated by return_codes_map_maker.py. Do Not Edit. */\n\n')
    outfile.write('#ifndef DRAGON_RETURN_CODES_MAP\n')
    outfile.write('#define DRAGON_RETURN_CODES_MAP\n\n')
    outfile.write('/* This defines a static array to map\n')
    outfile.write(' * return values in return_codes.h\n')
    outfile.write(' * to a string representation. This\n')
    outfile.write(' * source file was auto-generated by\n')
    outfile.write(' * the tool return_codes_map_maker.py\n')
    outfile.write('*/\n\n')
    outfile.write('/* This file is not to be included by other C code.\n')
    outfile.write(' * Instead, use the dragon_get_rc_string API call\n')
    outfile.write(' * to translate a return code to its string representation.\n')
    outfile.write('*/\n\n')
    outfile.write('#include <dragon/return_codes.h>\n\n')
    outfile.write('static const char* dragon_rc_map[] = {\n')

    outfile2 = open('../../dragon/return_codes.pxd', 'w')
    outfile2.write('# This file was auto-generated by return_codes_map_maker.py. Do Not Edit.\n')
    outfile2.write('cdef extern from "<dragon/return_codes.h>":\n')
    outfile2.write('\n')
    outfile2.write('    ctypedef enum dragonError_t:\n')

    outfile3 = open('../../dragon/pydragon_rc.pyx', 'w')
    outfile3.write('# This file was auto-generated by return_codes_map_maker.py. Do Not Edit.\n')
    outfile3.write('from dragon.return_codes cimport *\n\n')
    outfile3.write('cpdef enum DragonError:\n')

    mapping = False
    for line in infile:
        if 'dragonError_t' in line:
            mapping = False

        if len(line.strip()) != 0 and mapping:
            tmp = line.strip().rstrip(',')
            if tmp.startswith('DRAGON_'):
                lst = tmp.split()
                val = lst[0]
                out_line = '    "'+lst[0]+'",'
                dragonError_t_val = lst[0] # No commas, no quotes in the return_codes.pxd file.
                enum_val = dragonError_t_val[7:]
                if val == 'DRAGON_BAD_RETURN_CODE':
                    # for last return code, don't end line with a comma.
                    out_line = out_line[:-1]
                outfile.write(out_line+'\n')
                outfile2.write('        '+dragonError_t_val+'\n')
                outfile3.write('    '+enum_val+' = '+dragonError_t_val+'\n')

        if 'dragonError_st' in line:
            mapping = True

    outfile.write('};\n\n')
    outfile.write(f'static const dragonError_t dragon_max_rc_value = DRAGON_BAD_RETURN_CODE;\n')
    outfile.write('\n#endif\n')

    outfile.close()
    outfile2.close()
    outfile3.close()


if __name__ == "__main__":
    main()

src/include/dragon/return_codes_map_maker.py
DragonHPC/dragon
import os
from distutils import log
from distutils.command.clean import clean as _clean
from distutils.dir_util import remove_tree
from functools import partial
from pathlib import Path
from sysconfig import get_config_var

from setuptools import Extension, find_packages, setup
from setuptools.command.build import build as _build
from setuptools.command.build_py import build_py as _build_py
from Cython.Build import cythonize

from dragon.cli import entry_points


ROOTDIR = str(Path(__file__).parent)


DragonExtension = partial(Extension,
    include_dirs=[f'{ROOTDIR}/lib', f'{ROOTDIR}/include'],
    library_dirs=[f'{ROOTDIR}/lib'],
    libraries=['dragon', 'rt'],
)


extensions = [
    DragonExtension("dragon.utils", ["dragon/pydragon_utils.pyx"]),
    DragonExtension("dragon.rc", ["dragon/pydragon_rc.pyx"]),
    DragonExtension("dragon.dtypes", ["dragon/pydragon_dtypes.pyx"]),
    DragonExtension("dragon.pheap", ["dragon/pydragon_heap.pyx"]),
    DragonExtension("dragon.locks", ["dragon/pydragon_lock.pyx"]),
    DragonExtension("dragon.channels", ["dragon/pydragon_channels.pyx"]),
    DragonExtension("dragon.heapmanager", ["dragon/pydragon_heapmanager.pyx"]),
    DragonExtension("dragon.managed_memory", ["dragon/pydragon_managed_memory.pyx"]),
    DragonExtension("dragon.dlogging.logger", ["dragon/dlogging/pydragon_logging.pyx"]),
    DragonExtension("dragon.pmod", ["dragon/pydragon_pmod.pyx"]),
    DragonExtension("dragon.perf", ["dragon/pydragon_perf.pyx"]),
    Extension(
        "dragon.launcher.pmsgqueue",
        ["dragon/launcher/pydragon_pmsgqueue.pyx"],
        include_dirs=[f'{ROOTDIR}/include', f'{ROOTDIR}/dragon/launcher/include'],
        library_dirs=[f'{ROOTDIR}/lib'],
        libraries=["pmsgqueue", "rt"],
    ),
]

class build(_build):

    user_options = _build.user_options + [
        ('cythonize', None, 'cythonize packages and modules'),
    ]

    boolean_options = _build.boolean_options + ['cythonize']

    def initialize_options(self):
        super().initialize_options()
        self.cythonize = 0

    def run(self):
        _cythonize = partial(cythonize,
            nthreads=int(os.environ.get('DRAGON_BUILD_NTHREADS', os.cpu_count())),
            show_all_warnings=True,
            compiler_directives={
                "language_level": 3,
            },
            build_dir=self.build_temp,
        )

        # Cythonize core extensions
        self.distribution.ext_modules = _cythonize(
            self.distribution.ext_modules,
            force=bool(self.cythonize),
        )

        if self.cythonize:
            # Cythonize everything except dragon.cli and dragon.__main__
            self.distribution.packages = ['dragon.cli']
            self.distribution.py_modules = ['dragon.__main__']
            self.distribution.ext_modules.extend(_cythonize(
                'dragon/**/*.py',
                exclude=['dragon/**/setup.py', 'dragon/cli/*.py', 'dragon/__main__.py'],
                compiler_directives={
                    "language_level": 3,
                    # Disable annotation typing when compiling arbitrary
                    # Python 3.x code.
                    "annotation_typing": False,
                },
                force=True,
            ))
            #build_py_options = self.distribution.command_options.setdefault('build_py', {})
            #build_py_options['compile'] = ('setup script', 1)
            build_ext_options = self.distribution.command_options.setdefault('build_ext', {})
            build_ext_options['inplace'] = ('setup script', 0)

        super().run()


class build_py(_build_py):

    # The find_modules() implementation in setuptools (really distutils)
    # automatically adds __init__.py for any non-root package module listed in
    # py_modules. As a result, our implementation is a simplified variant of
    # the original that does not automatically add __init__.py.

    def find_modules(self):
        """Finds individually-specified Python modules, ie. those listed by
        module name in 'self.py_modules'.  Returns a list of tuples (package,
        module_base, filename): 'package' is a tuple of the path through
        package-space to the module; 'module_base' is the bare (no
        packages, no dots) module name, and 'filename' is the path to the
        ".py" file (relative to the distribution root) that implements the
        module.
        """
        # Maps package names to the corresponding directory
        packages = {}

        # List of (package, module, filename) tuples to return
        modules = []

        for module in self.py_modules:
            path = module.split('.')
            package = '.'.join(path[0:-1])
            module_base = path[-1]

            try:
                package_dir = packages[package]
            except KeyError:
                package_dir = self.get_package_dir(package)
                # We explicitly do not what to automatically add __init__.py,
                # hence we pass None as the first argument, but we still want
                # to check the package_dir in order to properly resolve the
                # module_file below.
                self.check_package(None, package_dir)
                packages[package] = package_dir

            module_file = os.path.join(package_dir, module_base + ".py")
            if not self.check_module(module, module_file):
                continue

            modules.append((package, module_base, module_file))

        return modules


class clean(_clean):

    def run(self):
        super().run()

        # Clean up any in-place extensions
        ext_suffix = get_config_var('EXT_SUFFIX')
        for ext_file in Path('dragon').glob(f'**/*{ext_suffix}'):
            log.info(f"removing '{ext_file}'")
            if not self.dry_run:
                ext_file.unlink()

        # Clean __pycache__ directories
        for directory in Path('dragon').glob('**/__pycache__'):
            remove_tree(str(directory), dry_run=self.dry_run)


setup(
    cmdclass={"build": build, "build_py": build_py, "clean": clean},
    options = {
        "build_py": {"compile": 1},
        "build_ext": {"inplace": 1},
    },
    name="dragon",
    version=os.environ.get('DRAGON_VERSION', 'latest'),
    description="Python multiprocessing over the Dragon distributed runtime",
    packages=find_packages(),
    ext_modules = extensions,
    entry_points=entry_points,
    python_requires=">=3.9",
    install_requires=['cloudpickle', 'numpy', ],

)

src/setup.py
DragonHPC/dragon
from dragon.channels import Channel, Message, ChannelSendH, ChannelRecvH

# create a Channel
ch = Channel(1)

# serialze the Channel and write that to a file
ser = ch.serialize()
ser_file = open("channel.dat", "wb")
ser_file.write(ser)
ser_file.close()

# create and open a send handle to Channel 1 and a recv handle to Channel 2
sh = ch.sendh()
sh.open()

# create a message we will keep sending over and over again
sm=Message().create_alloc(512)
mb=sm.bytes_memview()
mb[1] = b"K"
mb[511] = b"b"

# run a loop where we will keep sending the message
for i in range(100):
    sh.send(sm)

# destroy the message and by default the underlying managed memory allocation
sm.destroy()

# wait for user input
input("Enter to continue: ")

# cleanup the channel
ch.destroy()

test/channels_subtests/ch1.py
DragonHPC/dragon
from dragon.channels import Channel, Message, ChannelSendH, ChannelRecvH

# read the serialized channel from a file
ser_file = open("channel.dat", "rb")
ser = ser_file.read()
ser_file.close()

# attach to a Channel
ch = Channel.attach(ser)

# create and open a recv handle to Channel
rh = ch.recvh()
rh.open()

# run a loop where we will keep receiving the message
for i in range(100):
    msg = rh.recv()
    rb = msg.bytes_memview()
    print(f"Value at 1 is {rb[1]}")
    msg.destroy()

# wait for user input
input("Enter to continue: ")

test/channels_subtests/ch2.py
DragonHPC/dragon
import multiprocessing as mp
from multiprocessing import Process
from dragon.channels import Channel, Message, ChannelSendH, ChannelRecvH
from dragon.managed_memory import MemoryPool

def recv_worker():
    # read the serialized channel from a file
    ser_file = open("channel.dat", "rb")
    ser = ser_file.read()
    ser_file.close()

    # attach to a Channel
    ch = Channel.attach(ser)

    # create and open a recv handle to Channel
    rh = ch.recvh()
    rh.open()

    # run a loop where we will keep receiving the message
    for i in range(26):
        msg = rh.recv()
        rb = msg.bytes_memview()
        rb = rb.cast('c') # Cast to char to spit out letters instead of ascii values
        print(f"Value at 1 is {rb[1]}")
        msg.destroy()

    rh.close()
    ch.detach()

if __name__ == '__main__':
    mp.set_start_method('spawn', force=True)

    pool_name = "pydragon_channel_test"
    pool_size = 1073741824 # 1GB
    pool_uid = 1
    mpool = MemoryPool(pool_size, pool_name, pool_uid)

    # create a Channel
    ch = Channel(mpool, 1)

    # serialze the Channel and write that to a file
    ser = ch.serialize()
    ser_file = open("channel.dat", "wb")
    ser_file.write(ser)
    ser_file.close()

    # create and open a send handle to Channel 1 and a recv handle to Channel 2
    sh = ch.sendh()
    sh.open()

    # create a message we will keep sending over and over again
    sm = Message.create_alloc(mpool, 512)
    mb = sm.bytes_memview()
    mb = mb.cast('c') # Cast to char to directly assign letters in

    alpha = "abcdefghijklmnopqrstuvwxyz"
    # run a loop where we will keep sending the message
    for i in range(26):
        # Because messages are buffered by default we can reuse the same message allocation multiple times
        mb[1] = alpha[i].encode()
        sh.send(sm)

    # destroy the message and by default the underlying managed memory allocation
    # This does not release all buffered messages
    sm.destroy()
    # Close the send handle
    sh.close()

    worker = mp.Process(target=recv_worker, args=())
    worker.start()
    worker.join()

    # wait for user input
    input("Enter to continue: ")

    # cleanup the channel
    ch.destroy()
    mpool.destroy()

test/channels_subtests/ch_ex.py
DragonHPC/dragon
#!/usr/bin/env python3

import time
import sys
from dragon.channels import Channel, Message, ChannelSendH, ChannelRecvH
from dragon.managed_memory import MemoryPool

POOL_NAME = "pydragon_channel_test"
POOL_SIZE = 1073741824 # 1GB
POOL_UID = 0

N_REPEATS = 10000
N_MSGS = 100
MSG_WORDS = 4096
C_UID = 0
B_PER_MBLOCK = 1024
BILL = 1000000000
MSG_SIZE = 4 * MSG_WORDS # memoryviews in integer format take 4 byte ints

mpool = MemoryPool(POOL_SIZE, POOL_NAME, POOL_UID)
ch = Channel(mpool, C_UID)

ch_ser = ch.serialize()

send_hdl = ch.sendh()
send_hdl.open()

recv_hdl = ch.recvh()
recv_hdl.open()

msg = Message().create_alloc(mpool, MSG_SIZE)
msg_view = msg.bytes_memview()
msg_view = msg_view.cast('i')

for i in range(0, MSG_WORDS):
    msg_view[i] = i

print(f"Measuring performance for {MSG_SIZE} byte messages")

start_time = time.monotonic_ns()

for j in range(N_REPEATS):
    for i in range(N_MSGS):
        send_hdl.send(msg)

    for i in range(N_MSGS):
        recv_msg = recv_hdl.recv()

        if i == (N_MSGS - 1):
            recv_view = recv_msg.bytes_memview()
            recv_view = recv_view.cast('i')
            for k in range(0, MSG_WORDS):
                if recv_view[k] != k:
                    print(f"Expected {k} at index {k} but got {recv_view[k]}")
                    sys.exit(0)

        recv_msg.destroy()

end_time = time.monotonic_ns()
total = 1.0 * (end_time - start_time) * 1e-9
mrate = N_REPEATS * N_MSGS / total
uni_lat = 5e5 * total / (N_REPEATS * N_MSGS)
bandwidth = MSG_SIZE * N_REPEATS * N_MSGS / (1024 * 1024 * total)
print("Perf: pt2pt message rate = {:0.3f} msgs/sec, uni-directional latency = {:0.3f} usec, BW = {:0.3f} MiB/s".format(mrate, uni_lat, bandwidth))

recv_hdl.close()
send_hdl.close()
ch.destroy()
mpool.destroy()

test/channels_subtests/perf_channels.py
DragonHPC/dragon
#!/usr/bin/env python3

import time
import unittest
import os
import multiprocessing as mp
from dragon.channels import Channel, Message, ChannelSendH, ChannelRecvH, ChannelError, ChannelSendError, \
    ChannelRecvError, OwnershipOnSend, LockType, FlowControl, ChannelFull, ChannelEmpty, ChannelHandleNotOpenError, \
    ChannelRecvTimeout, ChannelSendTimeout, EventType, ChannelBarrierBroken, ChannelBarrierReady, MASQUERADE_AS_REMOTE, \
    ChannelSet, ChannelSetTimeout, POLLIN, SPIN_WAIT, ADAPTIVE_WAIT
from dragon.managed_memory import MemoryPool, DragonMemoryError
import sys

BARRIER_CHANNEL_CAPACITY = 5
MAX_SPINNERS = 5

def worker_attach_detach(ch_ser, pool_ser):
    ch = Channel.attach(ch_ser)
    sendh = ch.sendh()
    sendh.open()

    mpool = MemoryPool.attach(pool_ser)
    msg = Message.create_alloc(mpool, 512)
    mview = msg.bytes_memview()
    mview[0:5] = b"Hello"
    sendh.send(msg)

    sendh.close()
    msg.destroy()
    mpool.detach()
    ch.detach()

def worker_pickled_attach(ch, mpool):
    sendh = ch.sendh()
    sendh.open()

    msg = Message.create_alloc(mpool, 512)
    mview = msg.bytes_memview()
    mview[:5] = b"Howdy"
    sendh.send(msg)

    sendh.close()
    msg.destroy()
    mpool.detach()
    ch.detach()

def worker_send_recv(id, ch_ser, ch2_ser, pool_ser):
    try:
        ch = Channel.attach(ch_ser)
        ch2 = Channel.attach(ch2_ser)
        mpool = MemoryPool.attach(pool_ser)

        sendh = ch.sendh(wait_mode=SPIN_WAIT)
        sendh.open()

        msg = Message.create_alloc(mpool, 32)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        sendh.send(msg)

        recvh = ch2.recvh(wait_mode=ADAPTIVE_WAIT)
        recvh.open()

        msg = recvh.recv()
        msg.destroy()
        recvh.close()
        ch.detach()
        sys.exit(0)

    except Exception as ex:
        print(repr(ex))
        sys.exit(1)

def worker_fill_poll_empty(ch_ser, pool_ser):
    try:
        ch = Channel.attach(ch_ser)
        sendh = ch.sendh()
        sendh.open()

        mpool = MemoryPool.attach(pool_ser)

        for j in range(3):
            for i in range(BARRIER_CHANNEL_CAPACITY):
                msg = Message.create_alloc(mpool, 512)
                mview = msg.bytes_memview()
                mview[0:5] = b"Hello"
                sendh.send(msg)

            if not ch.poll(event_mask=EventType.POLLEMPTY, timeout=5):
                print("Error in polling for empty")
                sys.exit(1)


    except Exception as ex:
        print(repr(ex))
        sys.exit(1)

def worker_empty_poll_full(ch_ser):
    try:
        ch = Channel.attach(ch_ser)
        recvh = ch.recvh()
        recvh.open()

        for j in range(3):
            if not ch.poll(event_mask=EventType.POLLFULL, timeout=5):
                print("Error in polling for full")
                sys.exit(1)

            for i in range(BARRIER_CHANNEL_CAPACITY):
                msg = recvh.recv()
                mview = msg.bytes_memview()
                if mview[0:5] != b"Hello":
                    print("Error in received message")
                    sys.exit(1)

    except Exception as ex:
        print(repr(ex))
        sys.exit(1)


def worker_barrier_wait(ch_ser):
    try:
        ch = Channel.attach(ch_ser)
        recvh = ch.recvh()
        recvh.open()

        try:
            if not ch.poll(event_mask=EventType.POLLBARRIER, timeout=0):
                print('Error calling Barrier Poll', flush=True)
                sys.exit(1)
        except ChannelBarrierReady:
            ch.poll(event_mask=EventType.POLLBARRIER_RELEASE, timeout=0)

        msg = recvh.recv(timeout=5)
        x = int.from_bytes(msg.bytes_memview(), byteorder=sys.byteorder, signed=True)
        if x < -1 or x > BARRIER_CHANNEL_CAPACITY:
            print(f'Error: Got Barrier Value {x}', flush=True)
            sys.exit(1)

        if x == -1:
            # 2 is used to communicate an abort on the barrier back to caller.
            sys.exit(2)

        sys.exit(0)

    except ChannelBarrierBroken as ex:
        # return a 2 on the exit to indicate a barrier broken exception.
        sys.exit(2)

    except Exception as ex:
        #print("There was an exception while running the barrier wait process.", flush=True)
        #print(repr(ex), flush=True)
        sys.exit(1)


class ChannelCreateTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.pool_name = f"pydragon_channel_test_{os.getpid()}"
        pool_size = 1073741824  # 1GB
        pool_uid = 1
        cls.mpool = MemoryPool(pool_size, cls.pool_name, pool_uid)

    @classmethod
    def tearDownClass(cls):
        cls.mpool.destroy()

    def test_pool_create_destroy(self):
        pool = pool_name = "pydragon_pool_test"
        pool_size = 1073741824  # 1GB
        pool_uid = 42
        mpool = MemoryPool(pool_size, pool_name, pool_uid)
        mpool.destroy()
        mpool = MemoryPool(pool_size, pool_name, pool_uid)
        mpool.destroy()

    def test_create_destroy(self):
        ch = Channel(self.mpool, 1)
        ch.destroy()
        with self.assertRaises(ChannelError):
            ch.serialize()

    def test_with_stmt(self):
        msg = Message.create_alloc(self.mpool, 1)
        mview = msg.bytes_memview()
        mview[0] = 1

        with Channel(self.mpool, 1) as ch:
            with ch.recvh() as recvh:
                with self.assertRaises(TimeoutError):
                    recvh.recv(timeout = 1.5)

            with ch.sendh() as sendh:
                sendh.send(msg)

            with ch.recvh() as recvh:
                msg2 = recvh.recv(timeout = 1.5)
                mview2 = msg2.bytes_memview()
                self.assertEqual(1, mview2[0])

    def test_with_stmt_on_attach(self):
        msg = Message.create_alloc(self.mpool, 1)
        mview = msg.bytes_memview()
        mview[0] = 1

        # Make sure we're getting back a copy of the serializer, not a reference
        ch = Channel(self.mpool, 1)
        ch_ser = ch.serialize()  # Will return a pre-loaded copy of the serializer based on the above detach call
        with ch.sendh() as sendh:
            sendh.send(msg)

        with Channel.attach(ch_ser) as ch2:
            with ch2.recvh() as recvh:
                msg2 = recvh.recv(timeout = 1.5)
                mview2 = msg2.bytes_memview()
                self.assertEqual(1, mview2[0])

        ch = Channel.attach(ch_ser)
        ch.destroy()

    def test_create_destroy_repeat(self):
        for i in range(10):
            ch = Channel(self.mpool, 1)
            ch.destroy()
            with self.assertRaises(ChannelError):
                ch.serialize()

    def test_negative_uid(self):
        # Causes overflow error since the parameter type is uint64_t, can't cast from negative
        with self.assertRaises(OverflowError):
            Channel(self.mpool, -1)

    def test_zero_uid(self):
        ch = Channel(self.mpool, 0)
        ch.destroy()

    def test_string_uid(self):
        with self.assertRaises(TypeError):
            Channel(self.mpool, "0")

    def test_small_message(self):
        ch = Channel(self.mpool, 1)
        sendh = ch.sendh()
        recvh = ch.recvh()
        msg = Message.create_alloc(self.mpool, 1)
        mview = msg.bytes_memview()
        mview[0] = 1

        sendh.open()
        sendh.send(msg)
        sendh.close()
        msg.destroy()

        recvh.open()
        msg = recvh.recv()
        mview = msg.bytes_memview()
        self.assertEqual(1, mview[0])
        recvh.close()
        msg.destroy()
        ch.destroy()

    def test_attr_channel(self):
        with self.assertRaises(ChannelError):
            ch = Channel(self.mpool, 1, block_size=10)

        with self.assertRaises(ChannelError):
            ch = Channel(self.mpool, 1, block_size=256, capacity=0)

        with self.assertRaises(TypeError):
            ch = Channel()
            ch.recvh()

        ch = Channel(self.mpool, 1, block_size=256, capacity=10)
        sendh = ch.sendh()
        recvh = ch.recvh()
        msg = Message.create_alloc(self.mpool, 1)
        mview = msg.bytes_memview()
        mview[0] = 1

        sendh.open()
        sendh.send(msg)
        sendh.close()
        msg.destroy()

        recvh.open()
        msg = recvh.recv()
        mview = msg.bytes_memview()
        self.assertEqual(1, mview[0])
        recvh.close()

        self.assertEqual(ch.capacity, 10)
        self.assertEqual(ch.block_size, 256)

        msg.destroy()
        ch.destroy()

    def test_channel_full(self):
        ch = Channel(self.mpool, 1, block_size=256, capacity=1)
        sendh = ch.sendh()
        recvh = ch.recvh()
        msg = Message.create_alloc(self.mpool, 1)
        mview = msg.bytes_memview()
        mview[0] = 1
        sendh.open()
        recvh.open()
        sendh.send(msg)

        msg2 = Message.create_alloc(self.mpool, 1)
        mview2 = msg2.bytes_memview()
        mview2[0] = 1

        with self.assertRaises(TimeoutError):
            sendh.send(msg2, blocking = True, timeout=1)

        msg = recvh.recv()
        msg.destroy()
        msg2.destroy()
        sendh.close()
        recvh.close()
        ch.destroy()


    def test_attr_channel2(self):
        ch = Channel(self.mpool, 1, fc_type=FlowControl.RESOURCES_FLOW_CONTROL, lock_type=LockType.GREEDY)
        sendh = ch.sendh()
        recvh = ch.recvh()
        msg = Message.create_alloc(self.mpool, 1)
        mview = msg.bytes_memview()
        mview[0] = 1

        sendh.open()
        sendh.send(msg)
        sendh.close()
        msg.destroy()

        recvh.open()
        msg = recvh.recv()
        mview = msg.bytes_memview()
        self.assertEqual(1, mview[0])
        recvh.close()

        msg.destroy()
        ch.destroy()


    def test_multiple_channels(self):
        channels = []
        for i in range(0, 5):
            ch = Channel(self.mpool, i)
            channels.append(ch)
            sendh = ch.sendh()
            msg = Message.create_alloc(self.mpool, 512)
            mview = msg.bytes_memview()
            mview[1] = i
            sendh.open()
            sendh.send(msg)
            sendh.close()
            msg.destroy()

        for i in range(0, 5):
            ch = channels[i]
            recvh = ch.recvh()
            recvh.open()
            msg = recvh.recv()
            mview = msg.bytes_memview()
            self.assertEqual(i, mview[1])
            recvh.close()
            msg.destroy()
            ch.destroy()

    def test_uid_in_use(self):
        # @MCB TODO: Review the error messages on this.
        # It fails, but in an unexpected way.  The error
        # comes back through the managed memory library in a weird way.
        ch1 = Channel(self.mpool, 1)
        with self.assertRaises(ChannelError):
            Channel(self.mpool, 1)
        ch1.destroy()

    @unittest.skip("Cython auto-truncates floats to integer types where possible.  Need to determine behavior.")
    def test_uid_float(self):
        with self.assertRaises(TypeError):
            Channel(self.mpool, 1.5)

    def test_serialize_after_detach(self):
        # Make sure we can store a serializer after detaching
        ch = Channel(self.mpool, 1)
        ch.detach(serialize=True)
        ch_ser = ch.serialize()
        ch.attach(ch_ser)
        ch.destroy()

    @unittest.skip("Not valid yet. Need to do some investigation as to similar problem.")
    def test_destroy_after_detach(self):
        # Make sure we can store a serializer after detaching
        ch = Channel(self.mpool, 1, flags=MASQUERADE_AS_REMOTE)
        ch.detach()
        ch.destroy()

    def test_serialize_lifecycle(self):
        # Make sure we're getting back a copy of the serializer, not a reference
        ch = Channel(self.mpool, 1)
        ch.detach(serialize=True)  # Stores a serializer in the object and sets a flag indicating such
        ch_ser = ch.serialize()  # Will return a pre-loaded copy of the serializer based on the above detach call
        del ch  # the __del__ call will check a flag to see if serialize was ever called, and free if necessary
        ch = Channel.attach(ch_ser)
        ch.destroy()

    def test_get_pool(self):
        # Simple test to see if we get a usable memorypool object
        ch = Channel(self.mpool, 1)
        mpool2 = ch.get_pool()
        mem = mpool2.alloc(512)
        mem.free()
        ch.destroy()

    def test_message_count(self):
        ch = Channel(self.mpool, 1)
        sh = ch.sendh()
        rh = ch.recvh()
        sh.open()
        rh.open()

        x = bytearray("Hello", 'utf-8')
        sh.send_bytes(x, len(x))

        self.assertEqual(1, ch.num_msgs)

        rmsg = rh.recv()
        self.assertEqual(0, ch.num_msgs)
        rmsg.destroy()

        ch.destroy()

    def test_serialized_info_no_attach(self):
        cuid = 90210
        ch = Channel(self.mpool, cuid)
        ch_ser = ch.serialize()

        (s_cuid, typeid) = Channel.serialized_uid_type(ch_ser)
        self.assertEqual(s_cuid, cuid, "ID's do not match")
        self.assertEqual(typeid, 1, "Expected enum type 1")

        ch.destroy()

    def test_serialized_pool_info_no_attach(self):
        cuid = 90210
        ch = Channel(self.mpool, cuid)
        ch_ser = ch.serialize()

        (pool_uid, pool_fname) = Channel.serialized_pool_uid_fname(ch_ser)
        self.assertEqual(pool_uid, 1, "Pool IDs don't match, expected 1")
        self.assertEqual(pool_fname, "/_dragon_" + self.pool_name + "_manifest")

        ch.destroy()


class ChannelTests(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        pool_name = f"pydragon_channel_test_{os.getpid}"
        pool_size = 1073741824  # 1GB
        pool_uid = 1
        cls.mpool = MemoryPool(pool_size, pool_name, pool_uid)

    @classmethod
    def tearDownClass(cls):
        cls.mpool.destroy()

    def setUp(self):
        self.ch = Channel(self.mpool, c_uid=1, max_spinners=MAX_SPINNERS)
        self.ch2 = Channel(self.mpool, c_uid=2, capacity=BARRIER_CHANNEL_CAPACITY, max_spinners=MAX_SPINNERS)

    def tearDown(self):
        self.ch.destroy()
        self.ch2.destroy()

    def test_poll_timeout(self):
        timeout = 0.4
        start = time.monotonic()
        result = self.ch.poll(timeout=timeout)
        delta = time.monotonic() - start

        self.assertFalse(result)
        self.assertGreater(delta, timeout, 'delta should be bigger than timeout')

    def test_unopened_recv_hdl(self):
        recvh = self.ch.recvh()
        with self.assertRaises(ChannelHandleNotOpenError):
            recvh.recv()

    def test_poll_time(self):
        # Without a message it times out
        self.assertEqual(self.ch.poll(timeout=1.1),False)

    def test_poll(self):
        recvh = self.ch.recvh()
        recvh.open()

        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLIN), False)
        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLIN|EventType.POLLOUT), True)

        sendh = self.ch.sendh()
        sendh.open()
        msg = Message.create_alloc(self.mpool, 32)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        sendh.send(msg, timeout=0)

        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLIN), True)

        # There is a message
        self.assertEqual(self.ch.poll(), True)
        recvh.recv()

        # Without a message it times out
        self.assertEqual(self.ch.poll(timeout=1.5), False)

        msg.destroy()
        sendh.close()
        recvh.close()

    def test_poll_empty_full(self):
        ch_ser = self.ch2.serialize()
        pool_ser = self.mpool.serialize()

        proc1 = mp.Process(target=worker_fill_poll_empty, args=(ch_ser, pool_ser))
        proc1.start()

        proc2 = mp.Process(target=worker_empty_poll_full, args=(ch_ser,))
        proc2.start()

        proc1.join()
        proc2.join()
        self.assertTrue(proc1.exitcode == 0, 'The fill and poll for empty process did not complete successfully')
        self.assertTrue(proc2.exitcode == 0, 'The empty and poll for full process did not complete successfully')

    def test_poll_barrier(self):
        ch_ser = self.ch2.serialize()
        proc_list = []

        for i in range(BARRIER_CHANNEL_CAPACITY):
            proc = mp.Process(target=worker_barrier_wait, args=(ch_ser, ))
            proc.start()
            proc_list.append(proc)

        for i in range(BARRIER_CHANNEL_CAPACITY):
            proc_list[i].join()
            self.assertEqual(proc_list[i].exitcode, 0, 'Non-zero exitcode from barrier proc')

    def test_poll_barrier2(self):
        ch_ser = self.ch2.serialize()
        try:
            for j in range(5):
                proc_list = []

                for i in range(BARRIER_CHANNEL_CAPACITY*2):
                    proc = mp.Process(target=worker_barrier_wait, args=(ch_ser, ))
                    proc.start()
                    proc_list.append(proc)

                for i in range(BARRIER_CHANNEL_CAPACITY*2):
                    proc_list[i].join()
                    self.assertEqual(proc_list[i].exitcode, 0, 'Non-zero exitcode from barrier proc')
        except:
            print('**** Error in test_poll_barrier2')
            print(f'*** Receivers={self.ch2.blocked_receivers}')
            print(f'*** Number of messages={self.ch2.num_msgs}')
            print(f'*** Barrier Count is {self.ch2.barrier_count}')

    def test_poll_barrier3_with_abort(self):
        ch_ser = self.ch2.serialize()
        success = 0
        aborted = 0
        try:

            proc_list = []

            for i in range(BARRIER_CHANNEL_CAPACITY):
                proc = mp.Process(target=worker_barrier_wait, args=(ch_ser, ))
                proc.start()
                proc_list.append(proc)

            for i in range(BARRIER_CHANNEL_CAPACITY):
                proc_list[i].join()
                ec = proc_list[i].exitcode
                if ec == 0:
                    success+=1
                elif ec == 2:
                    aborted+=1

            proc_list = []

            for i in range(BARRIER_CHANNEL_CAPACITY*2):
                proc = mp.Process(target=worker_barrier_wait, args=(ch_ser, ))
                if i==4:
                    self.ch2.poll(event_mask=EventType.POLLBARRIER_ABORT)
                proc.start()
                proc_list.append(proc)

            for i in range(BARRIER_CHANNEL_CAPACITY*2):
                proc_list[i].join()
                ec = proc_list[i].exitcode
                if ec == 0:
                    success+=1
                elif ec == 2:
                    aborted+=1

            self.assertEqual(success, 5, f'Five should have succeeded. Instead there were {success}')
            self.assertEqual(aborted, 10, f'Ten should have aborted, instead there were {aborted}')
            self.ch2.poll(event_mask=EventType.POLLRESET)

        except Exception as ex:
            print(repr(ex))
            print('**** Error in test_poll_barrier3_with_abort')
            print(f'*** Receivers={self.ch2.blocked_receivers}')
            print(f'*** Number of messages={self.ch2.num_msgs}')
            print(f'*** Barrier Count is {self.ch2.barrier_count}')

    def test_poll_barrier_abort(self):
        ch_ser = self.ch2.serialize()
        proc_list = []

        for i in range(BARRIER_CHANNEL_CAPACITY-1):
            proc = mp.Process(target=worker_barrier_wait, args=(ch_ser, ))
            proc.start()
            proc_list.append(proc)

        time.sleep(1)
        self.ch2.poll(event_mask=EventType.POLLBARRIER_ABORT, timeout=5)

        for i in range(BARRIER_CHANNEL_CAPACITY-1):
            proc_list[i].join()
            self.assertEqual(proc_list[i].exitcode, 2, 'Non-zero exitcode from barrier proc')

        self.ch2.poll(event_mask=EventType.POLLRESET, timeout=5)

        proc_list = []

        for i in range(BARRIER_CHANNEL_CAPACITY):
            proc = mp.Process(target=worker_barrier_wait, args=(ch_ser, ))
            proc.start()
            proc_list.append(proc)

        for i in range(BARRIER_CHANNEL_CAPACITY):
            proc_list[i].join()
            self.assertEqual(proc_list[i].exitcode, 0, 'Non-zero exitcode from barrier proc')

    def test_blocking_recv(self):
        recvh = self.ch.recvh()
        with self.assertRaises(ChannelHandleNotOpenError):
            recvh.recv(timeout=1)

        recvh.open()

        with self.assertRaises(ChannelRecvTimeout) as ex:
            recvh.recv(timeout=1.5)

        with self.assertRaises(TimeoutError) as ex:
            recvh.recv(timeout=1)

        with self.assertRaises(ChannelEmpty) as ex:
            recvh.recv(timeout=0)

        sendh = self.ch.sendh()
        sendh.open()
        msg = Message.create_alloc(self.mpool, 32)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        sendh.send(msg, timeout=0)

        msg = recvh.recv()

        self.assertNotEqual(msg, None)

        msg.destroy()

        with self.assertRaises(ChannelEmpty) as ex:
            recvh.recv(timeout=0)

        sendh.close()
        recvh.close()

    def test_empty_buffer(self):
        recvh = self.ch.recvh()
        recvh.open()

        with self.assertRaises(ChannelEmpty) as ex:
            recvh.recv(blocking=False)

        recvh.close()

    def test_blocking_spin_wait_to_idle_wait(self):
        proc_list = []
        ser_ch = self.ch.serialize()
        ser_ch2 = self.ch2.serialize()
        ser_pool = self.mpool.serialize()

        for k in range(MAX_SPINNERS*2):
            proc = mp.Process(target=worker_send_recv, args=(k, ser_ch, ser_ch2, ser_pool))
            proc_list.append(proc)
            proc.start()

        recvh = self.ch.recvh(wait_mode=SPIN_WAIT)
        recvh.open()

        for _ in range(MAX_SPINNERS*2):
            msg = recvh.recv()
            msg.destroy()

        recvh.close()

        sendh = self.ch2.sendh(wait_mode=SPIN_WAIT)
        sendh.open()

        for _ in range(MAX_SPINNERS*2):
            msg = Message.create_alloc(self.mpool, 32)
            mview = msg.bytes_memview()
            mview[0:5] = b"Hello"
            sendh.send(msg, timeout=0)

        for proc in proc_list:
            proc.join()
            self.assertEqual(proc.exitcode, 0)

    def test_fill_buffer_exception(self):
        sendh = self.ch.sendh()
        sendh.open()
        msg = Message.create_alloc(self.mpool, 32)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        i = 0
        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLOUT), True)
        try:
            while True:
                sendh.send(msg, timeout=0)
                i += 1
        except ChannelFull:
            pass
        except:
            msg.destroy()
            sendh.close()
            self.fail("Expected ChannelFull exception")

        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLOUT), False)

        try:
            while True:
                sendh.send(msg, timeout=1.1)
                i += 1
        except ChannelSendTimeout:
            pass
        except:
            msg.destroy()
            sendh.close()
            self.fail("Expected ChannelSendTimeout exception")

        msg.destroy()
        sendh.close()

    def test_fill_then_empty_buffer(self):
        sendh = self.ch.sendh()
        sendh.open()
        msg = Message.create_alloc(self.mpool, 32)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        i = 0
        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLOUT), True)
        try:
            while True:
                sendh.send(msg, timeout=0)
                i += 1
        except ChannelFull:
            pass
        except:
            msg.destroy()
            sendh.close()
            self.fail("Expected ChannelFull exception")

        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLOUT), False)
        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLOUT|EventType.POLLIN), True)
        msg.destroy()
        sendh.close()
        recvh = self.ch.recvh()
        recvh.open()

        j = 0
        try:
            while True:
                msg = recvh.recv(blocking=False)
                mview = msg.bytes_memview()
                self.assertEqual(b"Hello", mview[0:5])
                msg.destroy()
                j += 1
        except ChannelEmpty:
            pass
        except:
            recvh.close()
            self.fail("Expected ChannelEmpty exception")

        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLIN), False)
        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLOUT), True)
        self.assertEqual(self.ch.poll(timeout=0, event_mask=EventType.POLLOUT|EventType.POLLIN), True)

        self.assertEqual(i, j)
        recvh.close()

    def test_create_attach_detach_destroy(self):
        # Create a channel, serialize it to a new process
        # Send a message from the new process
        # Receive the message from this process
        # Tests attach functionality from differing processes
        ch_ser = self.ch.serialize()
        pool_ser = self.mpool.serialize()

        proc = mp.Process(target=worker_attach_detach, args=(ch_ser, pool_ser))
        proc.start()
        proc.join()

        recvh = self.ch.recvh()
        recvh.open()
        msg = recvh.recv()
        mview = msg.bytes_memview()
        self.assertEqual(b"Hello", mview[0:5])
        msg.destroy()

    def test_create_attach_via_pickle_detach_destroy(self):
        # Create a channel, pickle it to a new process
        # Send a message from the new process
        # Receive the message from this process
        # Tests attach functionality from differing processes

        proc = mp.Process(target=worker_pickled_attach, args=(self.ch, self.mpool))
        proc.start()
        proc.join()

        recvh = self.ch.recvh()
        recvh.open()
        msg = recvh.recv()
        mview = msg.bytes_memview()
        self.assertEqual(b"Howdy", mview[:5])
        msg.destroy()

    def test_send_bytes(self):
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()
        x = bytearray(1024)
        x[0:5] = b"Hello"
        sendh.send_bytes(x, 1024)
        msg = recvh.recv()
        mview = msg.bytes_memview()
        self.assertEqual(mview[0:5], b"Hello")
        msg.destroy()

    def test_recv_bytes_timeout(self):
        recvh = self.ch.recvh()
        with self.assertRaises(ChannelHandleNotOpenError):
            recvh.recv_bytes(timeout=1)

        recvh.open()

        with self.assertRaises(ChannelRecvTimeout) as ex:
            recvh.recv_bytes(timeout=1.5)

        with self.assertRaises(TimeoutError) as ex:
            recvh.recv_bytes(timeout=1)

        with self.assertRaises(ChannelEmpty) as ex:
            recvh.recv_bytes(timeout=0)

        sendh = self.ch.sendh()
        sendh.open()
        my_bytes = b"Hello"
        sendh.send_bytes(my_bytes, timeout=0)

        msg = recvh.recv_bytes()

        self.assertEqual(msg, my_bytes)

        with self.assertRaises(ChannelEmpty) as ex:
            recvh.recv_bytes(timeout=0)

        sendh.close()
        recvh.close()

    def test_send_to_dest(self):
        # Assert messages can be delivered to specific payload objects
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()

        s_msg = Message.create_alloc(self.mpool, 512)
        r_msg = Message.create_alloc(self.mpool, 512)
        s_mview = s_msg.bytes_memview()
        s_mview[0:5] = b"Hello"
        sendh.send(s_msg, r_msg)
        s_msg.destroy()

        r_msg = recvh.recv()
        r_mview = r_msg.bytes_memview()

        self.assertEqual(b"Hello", r_mview[0:5])
        r_msg.destroy()

    def test_send_transfer(self):
        # Assert messages can be delivered by transferring ownership.
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()

        s_msg = Message.create_alloc(self.mpool, 512)
        s_mview = s_msg.bytes_memview()
        s_mview[0:5] = b"Hello"
        sendh.send(s_msg, ownership=sendh.transfer_ownership_on_send)

        r_msg = recvh.recv()
        r_mview = r_msg.bytes_memview()

        self.assertEqual(b"Hello", r_mview[0:5])
        r_msg.destroy()

    def test_send_copy(self):
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()

        s_msg = Message.create_alloc(self.mpool, 512)
        s_mview = s_msg.bytes_memview()
        s_mview[0:5] = b"Hello"
        sendh.send(s_msg, ownership=OwnershipOnSend.copy_on_send)
        s_msg.destroy()

        r_msg = recvh.recv()
        r_mview = r_msg.bytes_memview()

        self.assertEqual(b"Hello", r_mview[0:5])
        r_msg.destroy()

    def test_recv_to_dest(self):
        # Assert messages can be received into specific payload objects
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()

        s_msg = Message.create_alloc(self.mpool, 512)
        s_mview = s_msg.bytes_memview()
        s_mview[0:5] = b"Hello"
        sendh.send(s_msg)
        s_msg.destroy()

        r_msg = Message.create_alloc(self.mpool, 512)
        recvh.recv(r_msg)
        r_mview = r_msg.bytes_memview()

        self.assertEqual(b"Hello", r_mview[0:5])
        r_msg.destroy()

    def test_send_to_dest_and_recv_to_dest(self):
        # Make sure that if we receive a message to the same place we sent it
        #    that it works correctly
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()

        s_msg = Message.create_alloc(self.mpool, 512)
        r_msg = Message.create_alloc(self.mpool, 512)
        s_mview = s_msg.bytes_memview()
        s_mview[0:5] = b"Hello"
        sendh.send(s_msg, r_msg)
        s_msg.destroy()

        recvh.recv(r_msg)
        r_mview = r_msg.bytes_memview()

        self.assertEqual(b"Hello", r_mview[0:5])
        r_msg.destroy()

        sendh.close()
        recvh.close()

    def test_recv_bytes(self):
        # Assert that messages can be received as python byte objects
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()

        x = bytearray(1024)
        x[0:5] = b"Hello"
        sendh.send_bytes(x, 1024)
        recv_x = recvh.recv_bytes()
        self.assertEqual(recv_x[0:5], b"Hello")

        sendh.close()
        recvh.close()

    def test_send_bytes_sliced(self):
        # Assert that slices get sent with correct sizes
        sendh = self.ch.sendh()
        recvh = self.ch.recvh()
        sendh.open()
        recvh.open()

        x = bytearray(1024)
        x[0:11] = b"Hello World"
        sendh.send_bytes(x[0:5])
        recv_x = recvh.recv_bytes()
        self.assertEqual(recv_x[0:5], b"Hello")

        sendh.close()
        recvh.close()


class MessageTest(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        pool_name = "pydragon_channel_test"
        pool_size = 1073741824  # 1GB
        pool_uid = 1
        cls.mpool = MemoryPool(pool_size, pool_name, pool_uid)

    @classmethod
    def tearDownClass(cls):
        cls.mpool.destroy()

    def test_create_destroy(self):
        msg = Message.create_alloc(self.mpool, 512)
        msg.destroy()
        with self.assertRaises(ChannelError):
            msg.bytes_memview()

    def test_zero_size(self):
        with self.assertRaises(ChannelError):
            Message.create_alloc(self.mpool, 0)

    def test_negative_size(self):
        with self.assertRaises(OverflowError):
            Message.create_alloc(self.mpool, -1)

    def test_memview(self):
        msg = Message.create_alloc(self.mpool, 512)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        self.assertEqual(b"Hello", mview[0:5])
        msg.destroy()
        # Assert our memory is dead
        self.assertNotEqual(b"Hello", mview[0:5])


def worker_send_to_chset(ch_ser, pool_ser):
    # give some time to the head process to start
    # polling on the channel set

    try:
        ch = Channel.attach(ch_ser)
        sendh = ch.sendh()
        sendh.open()

        mpool = MemoryPool.attach(pool_ser)
        msg = Message.create_alloc(mpool, 32)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        sendh.send(msg)

        msg.destroy()
        sendh.close()

    except Exception as ex:
        print(repr(ex))


class ChannelSetTests(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        pool_name = f"pydragon_channelset_test_{os.getpid}"
        pool_size = 1073741824  # 1GB
        pool_uid = 1
        cls.mpool = MemoryPool(pool_size, pool_name, pool_uid)

    @classmethod
    def tearDownClass(cls):
        cls.mpool.destroy()

    def setUp(self):
        self.channel_list = []
        for i in range(3):
            self.channel_list.append(Channel(self.mpool, c_uid=i))
        self.ch_set = ChannelSet(self.mpool, self.channel_list)

    def tearDown(self):
        del self.ch_set
        for ch in self.channel_list:
            ch.destroy()

    def test_poll_timeout(self):
        timeout = 0.4
        start = time.monotonic()

        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=timeout)

        delta = time.monotonic() - start
        self.assertGreater(delta, timeout, 'delta should be bigger than timeout')

    def test_poll(self):
        ch_ser = self.channel_list[0].serialize()
        pool_ser = self.mpool.serialize()

        recvh = self.channel_list[0].recvh()
        recvh.open()

        # polling on an empty channelset
        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=0)

        # we are on DRAGON_SYNC mode
        # this process blocks on polling the channel set
        # while another process will send a message to a channel
        proc = mp.Process(target=worker_send_to_chset, args=(ch_ser, pool_ser))
        proc.start()
        # now, the first channel should have a message
        result, revent = self.ch_set.poll()
        self.assertEqual((result, revent), (self.channel_list[0], POLLIN))
        self.assertEqual(revent, POLLIN)

        # receive the message to empty the channel
        recvh.recv()

        # # Without a message it times out
        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=0.5)

        recvh.close()
        proc.join()

    def test_poll_more(self):
        # empty channelset for now
        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=0)

        pool_ser = self.mpool.serialize()

        for i in range(len(self.channel_list)):
            # send a message to channel i, poll, and
            # then receive
            recvh = self.channel_list[i].recvh()
            recvh.open()

            ch_ser = self.channel_list[i].serialize()

            proc = mp.Process(target=worker_send_to_chset, args=(ch_ser, pool_ser))
            proc.start()

            # there is a message in channel i
            result, revent = self.ch_set.poll()
            self.assertEqual((result, revent), (self.channel_list[i], POLLIN))

            # receive the message and poll again
            recvh.recv()
            # Without a message it times out
            with self.assertRaises(ChannelSetTimeout):
                self.ch_set.poll(timeout=0.5)

            recvh.close()
            proc.join()

        # verify we have an empty channel set
        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=0)

    def test_poll_with_dragon_no_sync(self):
        # Create a separate channel set with DRAGON_NO_SYNC
        # by initializing the set with capture_all_events=False
        channel_list = []
        for i in range(3):
            channel_list.append(Channel(self.mpool, c_uid=i+10))
        ch_set = ChannelSet(self.mpool, channel_list, capture_all_events=False)

        # poll an empty channel set
        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=0)

        # send a message to one of the channels
        sendh = channel_list[1].sendh()
        sendh.open()
        msg = Message.create_alloc(self.mpool, 32)
        mview = msg.bytes_memview()
        mview[0:5] = b"Hello"
        sendh.send(msg)

        # since the message was sent before polling
        # we should not be able to catch the message
        # and so the poll should timeout
        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=1)

        msg.destroy()
        sendh.close()
        for ch in channel_list:
            ch.destroy()

    def test_poll_when_set_is_full(self):
        # empty channelset for now
        with self.assertRaises(ChannelSetTimeout):
            self.ch_set.poll(timeout=0)

        pool_ser = self.mpool.serialize()

        proc_list = []
        # send a message to every channel
        for i in range(len(self.channel_list)):
            ch_ser = self.channel_list[i].serialize()
            proc = mp.Process(target=worker_send_to_chset, args=(ch_ser, pool_ser))
            proc.start()
            proc_list.append(proc)

        # Poll on a non-empty channel set.
        # Since the channel set is initialized with DRAGON_SYNC
        # we need to poll as many times as the sends, otherwise
        # the send operation will block
        for i in range(len(self.channel_list)):
            # this should return one of the channels
            # does not matter which and we cannot know which
            result, _ = self.ch_set.poll()
            self.assertNotEqual(result, False)

        for proc in proc_list:
            proc.join()


if __name__ == '__main__':
    mp.set_start_method('spawn', force=True)
    unittest.main(verbosity=2)

test/channels_subtests/test_basic_channels.py
DragonHPC/dragon
#!/usr/bin/env python3

import dragon.channels as dch
import dragon.managed_memory as dmm
import dragon.infrastructure.connection as dic

pool_name = "hyenas"
pool_size = 2 ** 30
pool_prealloc_blocks = None
pool_uid = 17
print('pool creating')
mpool = dmm.MemoryPool(pool_size, pool_name, pool_uid, pool_prealloc_blocks)
print('pool created')

channel_uid = 42
chan = dch.Channel(mpool, channel_uid)
print('channel created')

writer = dic.Connection(outbound_initializer==chan)
reader = dic.Connection(inbound_initializer=chan)

print('reading and writing connections made')

thing_to_send = 'hyena hyena hyena hyena'

writer.send(thing_to_send)
thing_got = reader.recv()

if thing_to_send == thing_got:
    print('hyenas received ok')
else:
    print('got {} and not {}'.format(thing_got, thing_to_send))

writer.close()

try:
    end_thing = reader.recv()
except EOFError as err:
    print('got the EOF')

reader.close()

chan.destroy()
mpool.destroy()

test/connection/basic_example.py
DragonHPC/dragon
#!/usr/bin/env python3

# This is a target program for a test of the Channels-based
# Connection object and argument delivery as needed
# for milestone 1.  It tests integration of the global
# services API, global services, and the shepherd in being
# able to use the API to ask GS for existing channels,
# communicate through them, use the Pipe constructor to
# make connection objects anonymously

# It starts as a managed process and looks at the delivered
# arguments for the name of input and output channels which the
# testbench should have already created.  It will query
# GS for those channels, make Connection objects out of them,
# then use the Pipe constructor anonymously to make two more
# Connection objects. It'll send the reading end of the new
# Connection through the original output Connection made
# from the pre-existing channel and expect to get another Connection
# through the original input Connection.  Then through these
# new ones it will send some data.

# This basically proves that we can get Channels by name, make Connections
# out of them, construct Connections anonymously, and serialize/deserialize
# them after construction and send data thru them.

# The test bench calling this script will start two mirror image copies.
# The script takes only one command line argument which is mainly used
# to differentiate the log files produced by either copy

import sys
import logging
import pickle

import dragon.channels as dch

import dragon.globalservices.api_setup as dapi
import dragon.globalservices.channel
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.parameters as dparm
import dragon.dlogging.util as dlog

if len(sys.argv) < 2:
    exit(1)

identity = sys.argv[1]

dlog.setup_logging(basename='catt-' + identity, level=logging.DEBUG)
log = logging.getLogger('')
log.info('hello from conn test worker {}'.format(identity))
log.info('p_uid is {}'.format(dparm.this_process.my_puid))

try:
    dapi.connect_to_infrastructure()
except Exception as err:
    log.info('failed to attach to infrastructure')
    log.exception(f'error is {err}')
    exit(1)

log.info('connected to infrastructure')

# now look at args to get inbound and outbound channel *names*
my_args = pickle.loads(dapi._ARG_PAYLOAD)

log.info(f'my args: {my_args}')

if len(my_args) != 2:
    log.error('unexpected number of args: {}'.format(len(my_args)))
    exit(1)

inbound_cname, outbound_cname = my_args

# ask for the descriptors for these args
try:
    inbound_cdesc = dragon.globalservices.channel.query(inbound_cname)
    outbound_cdesc = dragon.globalservices.channel.query(outbound_cname)
except Exception as e:
    log.exception("couldn't query existing channels")
    exit(1)

log.info('queried pre-existing channels')

try:
    inbound_ch = dch.Channel.attach(inbound_cdesc.sdesc)
    outbound_ch = dch.Channel.attach(outbound_cdesc.sdesc)
except Exception as e:
    log.exception("couldn't attach to existing channels")
    exit(1)

log.info('attached to pre-existing channels')

# should try this with a full duplex Connection object
orig_inbound_con = dconn.Connection(inbound_initializer=inbound_ch)
orig_outbound_con = dconn.Connection(outbound_initializer=outbound_ch)

first_msg = 'weasels'

try:
    orig_outbound_con.send(first_msg)
    partner = orig_inbound_con.recv()
except Exception as e:
    log.exception("couldn't send/recv initial msg", exc_info=True)
    exit(1)

if first_msg != partner:
    log.error('exchanged msgs do not match: {} vs {}'.format(first_msg, partner))
    exit(1)

log.info('we got our {}, now for a pipe and connection'.format(first_msg))

my_reader, my_writer = dconn.Pipe(duplex=False)

log.info('pipe called')

# should the process of sending my_reader put it in a
# state where it shouldn't be opened?
orig_outbound_con.send(my_reader)

log.info('reader obj sent')

others_reader = orig_inbound_con.recv()
log.info('reader obj received')

# need to do this, otherwise there will be
# a deadlock.
others_reader.open()
log.info('partner reader opened')

my_message = 'hyenas'

try:
    my_writer.send(my_message)
    log.info('my message is sent')
except Exception as e:
    log.exception("Couldn't send second message thru Connection")
    exit(1)

try:
    others_message = others_reader.recv()
    log.info('got the message from the other side')
except:
    log.exception("Couldn't get second message thru Connection")
    exit(1)

if my_message == others_message:
    log.info('messages are fine')
else:
    log.info('mismatch: {} vs {}'.format(my_message, others_message))
    exit(1)
# todo: some stuff here to do the cleanup explicitly

exit(0)

test/connection/conn_arg_test_target.py
DragonHPC/dragon
#!/usr/bin/env python3

"""Test script for basic Connection object"""
import os
import pickle
import random
import subprocess
import threading
import unittest
import multiprocessing

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.infrastructure.connection as dconn


class ConnectionTest(unittest.TestCase):
    def setUp(self):
        self.pool_name = 'conn_test_' + os.environ.get('USER', str(os.getuid()))
        self.pool_size = 2 ** 28
        self.pool_prealloc_blocks = None
        self.pool_uid = 17
        self.mpool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid, self.pool_prealloc_blocks)
        self.mpool_ser = self.mpool.serialize()

        self.first_channel_uid = 42
        self.first_chan = dch.Channel(self.mpool, self.first_channel_uid)
        self.first_chan_ser = self.first_chan.serialize()

        self.second_channel_uid = 43
        self.second_chan = dch.Channel(self.mpool, self.second_channel_uid)
        self.second_chan_ser = self.second_chan.serialize()

    def tearDown(self) -> None:
        self.second_chan.destroy()
        self.first_chan.destroy()
        self.mpool.destroy()

    def test_nothing(self):
        pass

    def pass_an_obj(self, obj):
        writer = dconn.Connection(outbound_initializer=self.second_chan)
        reader = dconn.Connection(inbound_initializer=self.second_chan)

        def send_obj():
            writer.send(obj)

        st = threading.Thread(target=send_obj)
        st.start()

        rec_obj = reader.recv()

        st.join()

        self.assertEqual(obj, rec_obj)

    def pass_some_bytes(self, some_bytes):
        writer = dconn.Connection(outbound_initializer=self.second_chan)
        reader = dconn.Connection(inbound_initializer=self.second_chan)

        def send_some_bytes():
            writer.send_bytes(some_bytes)

        st = threading.Thread(target=send_some_bytes)
        st.start()
        rec_bytes = reader.recv_bytes()
        st.join()
        self.assertEqual(some_bytes, rec_bytes)

    def test_bytes(self):
        self.pass_some_bytes(random.randbytes(17))

        for lg_size in range(25):
            self.pass_some_bytes(random.randbytes(2 ** lg_size))

    def test_objs(self):
        self.pass_an_obj('hyena')
        self.pass_an_obj({'hyena': 'hyena'})

        self.pass_an_obj(bytes(100000))

        self.pass_an_obj(bytes(200000))

        for lg_size in range(25):
            self.pass_an_obj(bytearray(2 ** lg_size))

        for lg_size in range(25):
            self.pass_an_obj(random.randbytes(2 ** lg_size))

    def test_recv_bytes_from_send(self):
        # required for multiprocessing.Connection unit tests...
        writer = dconn.Connection(outbound_initializer=self.second_chan)
        reader = dconn.Connection(inbound_initializer=self.second_chan)

        my_obj = {'hyenas': 'are awesome'}
        writer.send(my_obj)
        # would like this to raise, but....
        ser_obj = reader.recv_bytes()
        rec_obj = pickle.loads(ser_obj)
        self.assertEqual(my_obj, rec_obj)

    def test_recv_from_send_bytes(self):
        # again, weird usage but multiprocessing.Connection unit tests need
        # this functionality

        writer = dconn.Connection(outbound_initializer=self.second_chan)
        reader = dconn.Connection(inbound_initializer=self.second_chan)

        my_obj = {'hyenas': 'are awesome'}
        ser_obj = pickle.dumps(my_obj)
        writer.send_bytes(ser_obj)
        rec_obj = reader.recv()
        self.assertEqual(my_obj, rec_obj)

    def origin(self, write_options={}):
        some_bytes = random.randbytes(2 ** 24)

        try:
            writer = dch.Peer2PeerWritingChannelFile(self.first_chan, options=write_options)
            writer.open()
            writer.write_raw_header(len(some_bytes))
            writer.write(some_bytes)
        finally:
            writer.close()

        try:
            reader = dch.Peer2PeerReadingChannelFile(self.second_chan)
            reader.open()
            send_type, msg_len = reader.check_header()
            self.assertEqual(send_type, dch.ChannelAdapterMsgTypes.RAW_BYTES)
            self.assertEqual(msg_len, len(some_bytes))
            reader.advance_raw_header(msg_len)
            buf = bytearray(msg_len)
            reader.readinto(memoryview(buf))
            self.assertEqual(some_bytes, buf)
        finally:
            reader.close()

    def bounce(self, write_options={}):
        buf = bytearray(1)
        try:
            reader = dch.Peer2PeerReadingChannelFile(self.first_chan)
            reader.open()
            send_type, msg_len = reader.check_header()
            self.assertEqual(send_type, dch.ChannelAdapterMsgTypes.RAW_BYTES)
            reader.advance_raw_header(msg_len)
            buf = bytearray(msg_len)
            reader.readinto(memoryview(buf))
        finally:
            reader.close()

        try:
            writer = dch.Peer2PeerWritingChannelFile(self.second_chan, options=write_options)
            writer.open()
            writer.write_raw_header(len(buf))
            writer.write(buf)
        finally:
            writer.close()

    def test_channel_adapters_basic(self):

        bounce_th = threading.Thread(target=self.bounce)
        bounce_th.start()
        self.origin()
        bounce_th.join()

    def test_channel_adapters_tuned(self):

        options = { "write_options": {"small_blk_size": 2**8, "large_blk_size": 2**12, "huge_blk_size": 2**24,
                    "buffer_pool": self.mpool }}
        bounce_th = threading.Thread(target=self.bounce, args=(options))
        bounce_th.start()
        self.origin()
        bounce_th.join()

        broken_options = { "small_blk_size": 2**8, "large_blk_size": 2**24, "huge_blk_size": 2**12 }
        with self.assertRaises(ValueError):
            dch.Peer2PeerWritingChannelFile(self.first_chan, options=broken_options)

        empty_pool = dmm.MemoryPool.empty_pool()
        broken_options = { "buffer_pool": empty_pool }
        with self.assertRaises(dch.ChannelError):
            dch.Peer2PeerWritingChannelFile(self.first_chan, options=broken_options)

    def test_many2many_basic(self):
        some_bytes = random.randbytes(2 ** 24)

        try:
            writer = dch.Many2ManyWritingChannelFile(self.first_chan)
            writer.open()
            pickle.dump(some_bytes, file=writer, protocol=5)
        finally:
            writer.close()

        try:
            reader = dch.Many2ManyReadingChannelFile(self.first_chan)
            reader.open()
            buf = pickle.load(reader)
            self.assertEqual(some_bytes, buf)
        finally:
            reader.close()

    def _many2many_origin(self):
        some_bytes = random.randbytes(2 ** 24)

        try:
            writer = dch.Many2ManyWritingChannelFile(self.first_chan)
            writer.open()
            pickle.dump(some_bytes, file=writer, protocol=5)
        finally:
            writer.close()

        try:
            reader = dch.Many2ManyReadingChannelFile(self.second_chan)
            reader.open()
            buf = pickle.load(reader)
            self.assertEqual(some_bytes, buf)
        finally:
            reader.close()

    def _many2many_bounce(self):
        try:
            reader = dch.Many2ManyReadingChannelFile(self.first_chan)
            reader.open()
            buf = pickle.load(reader)
        finally:
            reader.close()

        try:
            writer = dch.Many2ManyWritingChannelFile(self.second_chan)
            writer.open()
            pickle.dump(buf, file=writer, protocol=5)
        finally:
            writer.close()

    def test_many2many_p2p(self):
        th = threading.Thread(target=self._many2many_bounce)
        th.start()
        msg = self._many2many_origin()
        th.join()


if __name__ == '__main__':
    multiprocessing.set_start_method('spawn')
    unittest.main()

test/connection/connection_obj_tests.py
DragonHPC/dragon
#!/usr/bin/env python3

import multiprocessing
import time

import dragon.channels as dch
import dragon.managed_memory as dmm
import dragon.infrastructure.connection as dic


def recv_n_objs(conn_or_init, num):
    if isinstance(conn_or_init, multiprocessing.connection.Connection):
        my_conn = conn_or_init
    else:
        input_chan = dch.Channel.attach(conn_or_init)
        my_conn = dic.Connection(inbound_initializer=input_chan)

    for _ in range(num):
        _ = my_conn.recv()


def send_n_obj(conn_or_init, num, obj_size):
    if isinstance(conn_or_init, multiprocessing.connection.Connection):
        my_conn = conn_or_init
    else:
        output_chan = dch.Channel.attach(conn_or_init)
        my_conn = dic.Connection(outbound_initializer=output_chan)

    start_obj = bytearray(obj_size)

    for _ in range(num):
        my_conn.send(start_obj)


def cheesy_measure_speed(num, obj_size, new_obj=True):
    if new_obj:
        mpool = dmm.MemoryPool(2 ** 30, 'hyenas', 17, None)
        chan = dch.Channel(mpool, 42)
        read_init = chan.serialize()
        write_init = read_init
    else:
        read_init, write_init = multiprocessing.Pipe(duplex=False)

    recv_proc = multiprocessing.Process(target=recv_n_objs, args=(read_init, num))
    send_proc = multiprocessing.Process(target=send_n_obj, args=(write_init, num, obj_size))

    recv_proc.start()

    st_time = time.time()
    send_proc.start()
    recv_proc.join()
    time_took = time.time() - st_time
    send_proc.join()

    if new_obj:
        chan.destroy()
        mpool.destroy()

    return time_took


if __name__ == "__main__":

    num_objs = 1000
    print('sending {} objects between 2 processes'.format(num_objs))

    for lg_size in range(10, 24):
        size = 2**lg_size
        existing = cheesy_measure_speed(num_objs, size, False)
        new = cheesy_measure_speed(num_objs, size)
        fmt = 'size 2**{}: old {:.3}\tnew {:.3}\tratio {:.3}'
        print(fmt.format(lg_size, existing, new, existing/new))

test/connection/simple_speed.py
DragonHPC/dragon
#!/usr/bin/env python3

'''
@MCB:
Intended to be used as a benchmark script to compare pipe multiprocessing vs dragon connections
'''

import multiprocessing as mp
import dragon.infrastructure.standalone_conn as dconn
import sys
import time
import csv

# This needs to be modified based on message size so we're not running a 15 minute test at 4MB with 8 duplex pipes
N_REPEATS = 1000
N_WARMUPS = int(N_REPEATS / 4)

# Divide ping-pong latency by 2 because MPI measurements are silly
# Adjust N_REPEAT const based on message size (huge for small, small for huge)
# Ping-pong latency will probably be wonky because of sleepy polling
# Mess with poll times and default block sizes in connection.py
# TODO: Look into fan & all to all latency tests

timings = []

def pingpong(writer, reader, ping_time, msg_size, write_first=True):
    x = bytearray(msg_size)

    # Warmup
    if write_first:
        for i in range(N_WARMUPS):
            writer.send(x)
            msg = reader.recv()
    else:
        for i in range(N_WARMUPS):
            msg = reader.recv()
            writer.send(x)

    # Measure
    start_time = time.monotonic_ns()
    if write_first:
        for i in range(N_REPEATS):
            writer.send(x)
            msg = reader.recv()
    else:
        for i in range(N_REPEATS):
            msg = reader.recv()
            writer.send(x)

    end_time = time.monotonic_ns()
    total = 1.0 * (end_time - start_time) * 1e-9 # Convert to seconds
    ping_time.value = total
    #writer.close()
    #reader.close()

# Unilateral write actor
def unilateral(actor, uni_time, msg_size, writer=True):
    x = bytearray(msg_size)

    # Warmup
    if writer:
        for i in range(N_WARMUPS):
            actor.send(x)
    else:
        for i in range(N_WARMUPS):
            msg = actor.recv()

    # Measure
    start_time = time.monotonic_ns()
    if writer:
        for i in range(N_REPEATS):
            actor.send(x)
    else:
        for i in range(N_REPEATS):
            msg = actor.recv()
    end_time = time.monotonic_ns()
    total = 1.0 * (end_time - start_time) * 1e-9 # Convert to seconds
    uni_time.value = total
    #actor.close() # this causes dragon to break but not mp.Pipe


def main(conn_class, num_pairs=1, msg_size=1024, is_pingpong=False):
    # Generate our read/write handles
    conns = []
    times = []
    procs = []
    for i in range(num_pairs):
        reader, writer = conn_class.Pipe(duplex=is_pingpong)
        # Store read/write pair as tuple
        conns.append((writer, reader))
        # Generate multiprocessing values to track times
        r_time = mp.Value("d", 0.0, lock=False)
        w_time = mp.Value("d", 0.0, lock=False)
        # Store as tuples similarly
        times.append((w_time, r_time))

    # Dumb way to do it but reiterate over pairs to match connections to processes
    for i in range(num_pairs):
        w_process = None
        r_process = None
        # Spawn our processes
        # TODO: This is a bit hard to read
        if is_pingpong:
            w_process = mp.Process(target=pingpong, args=(conns[i][0], conns[i+1][1], times[i][0], msg_size, True))
            r_process = mp.Process(target=pingpong, args=(conns[i+1][0], conns[i][1], times[i][1], 1, False))
        else:
            w_process = mp.Process(target=unilateral, args=(conns[i][0], times[i][0], msg_size, True))
            r_process = mp.Process(target=unilateral, args=(conns[i][1], times[i][1], 1, False))

        procs.append((w_process, r_process))

    for i in range(num_pairs):
        # Get processes
        w_process = procs[i][0]
        r_process = procs[i][1]
        # Start said processes
        w_process.start()
        r_process.start()

    total_0 = 0.0
    total_1 = 0.0

    for i in range(num_pairs):
        w_process = procs[i][0]
        r_process = procs[i][1]
        # Stop said processes
        w_process.join()
        r_process.join()
        # Routine finished. smiley_face.png
        total_0 += times[i][0].value
        total_1 += times[i][1].value

    # Average times
    total_0 = total_0 / num_pairs
    total_1 = total_1 / num_pairs

    timings.append((num_pairs, msg_size, total_0, total_1))
    # Spit them to stdout
    print("Size(b): {:8s} | {:0.5f}{:2s} | {:0.5f}".format(str(msg_size), total_0, " ", total_1))

if __name__ == "__main__":
    mp.set_start_method('spawn')
    conn_class = None
    is_pingpong = False
    if "-d" in sys.argv: # Check for dragon argument
        conn_class = dconn
        print("Using dragon connections")
    else: # Default to multiprocessing
        conn_class = mp
        print("Using multiprocessing connections")

    if "-pp" in sys.argv: # Check if we want to test pingpong
        is_pingpong = True

    # Chnage this to start at 8 bytes
    sizes = [2**x for x in range(6, 22)] # 8 bytes to 2MB
    #sizes = [2**x for x in range(10, 19)]
    #sizes = [1024, 2048, 4096]
    #sizes = [2**x for x in range(21, 26)]
    num_pairs = [1, 2, 4, 8]
    for p in num_pairs:
        print(f"Using {p} pairs")
        for sz in sizes:
            main(conn_class, p, sz, is_pingpong)

    with open("pipe.csv", 'w', newline='') as testfile:
        writer = csv.writer(testfile)
        for timing in timings:
            writer.writerow(timing)

test/connection/standalone_conn_bench.py
DragonHPC/dragon
#!/usr/bin/env python3

import multiprocessing as mp
import dragon.infrastructure.standalone_conn as dconn
import time
import argparse
import threading

# with the very naive approach to an all-all,
# using base multiprocessing Connection only works with very small message sizes
# without involving threading in the all-all, owing to short buffer depth in an
# OS-supplied socket.
# The dragon Connection version is much more forgiving and offers the option
# to get things flowing simply by making the underlying Channels deep enough
# and the underlying Pool big enough.
THREADED = True


def test_worker(links, result_link, msg_size, total_iterations):
    result_timings = []

    my_msg = bytearray(msg_size)

    def send_stuff():
        for link in links:
            link.send(my_msg)

    def recv_stuff():
        for link in links:
            msg = link.recv()
            assert len(msg) == msg_size

    # Each of these measures the time each worker
    # takes to send and recv a message of msg_size
    # through each link, which goes to each other worker.
    for _ in range(total_iterations):
        start = time.perf_counter()

        if THREADED:  # a slightly less dumb all-all
            send_th = threading.Thread(target=send_stuff)
            recv_th = threading.Thread(target=recv_stuff)

            send_th.start()
            recv_th.start()
            send_th.join()
            recv_th.join()
        else:  # the simplest all-all
            send_stuff()
            recv_stuff()

        result_timings.append(time.perf_counter() - start)

    result_link.send(result_timings)


def aa_test(nworkers, burn, iterations, msg_sz, my_pipe):
    total_iterations = burn + iterations

    # half duplex return links
    result_links = [my_pipe(duplex=False) for _ in range(nworkers)]

    all_links = {}

    # full duplex links; nworkers * (nworkers-1)/2 of them
    for i in range(nworkers):
        for j in range(i, nworkers):
            all_links[i, j] = my_pipe()

    all_workers = []
    for i in range(nworkers):
        # first end if other's index less than yours,
        # second end if other's index greater than yours
        links = []
        for k in range(nworkers):
            if k != i:
                links.append(all_links[min(i, k), max(i, k)][int(k < i)])

        proc = mp.Process(target=test_worker,
                          args=(links, result_links[i][1], msg_sz, total_iterations))

        all_workers.append(proc)

    for worker in all_workers:
        worker.start()

    avg_times = []
    total_times = []
    for i in range(nworkers):
        result = result_links[i][0].recv()[burn:]
        total_times.append(sum(result))
        avg_times.append(total_times[i] / iterations)

    avg_aa_completion_time = sum(avg_times) / nworkers  # seconds
    bandwidth = msg_sz * nworkers / avg_aa_completion_time  # bytes/second

    msg_size_in_k = msg_sz / (2 ** 10)
    bw_in_mb_sec = bandwidth / (2 ** 20)
    completion_time_in_ms = avg_aa_completion_time * 1000

    for worker in all_workers:
        worker.join()

    print(f'{nworkers=} {iterations=} {msg_size_in_k=} {completion_time_in_ms=:.3f} {bw_in_mb_sec=:.2f}')


if __name__ == "__main__":
    mp.set_start_method('spawn')

    parser = argparse.ArgumentParser(description='All-all channels/connection test')
    parser.add_argument('--num_workers', type=int, default=4,
                        help='number of workers to spawn at a time')

    parser.add_argument('--iterations', type=int, default=100,
                        help='number of iterations to do')

    parser.add_argument('--burn_iterations', type=int, default=5,
                        help='number of iterations to burn first time')

    parser.add_argument('--lg_message_size', type=int, default=4,
                        help='log base 2 of size of message to pass in')

    parser.add_argument('--dragon', action='store_true', help='run using dragon Connection')

    my_args = parser.parse_args()

    assert my_args.lg_message_size <= 25, 'limit: 32MB per worker to each other worker'

    if my_args.dragon:
        print('using dragon Connection')
        dconn.get_going(pool_size=2 ** 31)  # setup for standalone channels
        pipe_call = dconn.Pipe
    else:
        print('using multiprocessing Connection')
        pipe_call = mp.Pipe

    aa_test(nworkers=my_args.num_workers, burn=my_args.burn_iterations,
            iterations=my_args.iterations, msg_sz=2 ** my_args.lg_message_size,
            my_pipe=pipe_call)

test/connection/standalone_conn_bench_aa.py
DragonHPC/dragon
#!/usr/bin/env python3


import multiprocessing

import dragon.infrastructure.standalone_conn


def write_actor(writer):
    writer.send('hyenas')


def read_actor(reader):
    msg = reader.recv()
    if msg == 'hyenas':
        print('got hyenas')
    else:
        print('got {}, no hyenas'.format(msg))


def main():
    reader, writer = dragon.infrastructure.standalone_conn.Pipe(duplex=False)
    # reader, writer = multiprocessing.Pipe(duplex=False)
    write_proc = multiprocessing.Process(target=write_actor, args=(writer,))
    read_proc = multiprocessing.Process(target=read_actor, args=(reader,))

    write_proc.start()
    read_proc.start()

    write_proc.join()
    read_proc.join()


if __name__ == "__main__":
    multiprocessing.set_start_method('spawn')
    main()

test/connection/standalone_conn_example.py
DragonHPC/dragon
import atexit
import multiprocessing
import unittest

import dragon.infrastructure.standalone_conn as disc


def write_actor(writer):
    writer.send('hyenas')


SUCCESS = 0
FAILURE = 1


def read_actor(reader, expected):
    msg = reader.recv()

    if msg == expected:
        exit(SUCCESS)
    else:
        exit(FAILURE)


def echo_actor(handle):
    try:
        while msg := handle.recv():
            handle.send(msg)
    except EOFError:
        handle.close()
        exit(SUCCESS)
    except Exception as e:
        print(f'got some other exception: {e}')
        import traceback
        traceback.print_exc()

    exit(FAILURE)


def swap_actor(first, second):
    first_msg = first.recv()
    second_msg = second.recv()
    second.send(first_msg)
    first.send(second_msg)


class TestStandaloneConn(unittest.TestCase):

    def tearDown(self) -> None:
        # this runs what would normally happen at exit of the first process
        # to make a call to Pipe.
        disc.destroy_my_channels()
        disc.destroy_standalone()
        atexit.unregister(disc.destroy_my_channels)
        atexit.unregister(disc.destroy_standalone)
        disc._IS_INIT = False

    def test_single_oneway(self):
        reader, writer = disc.Pipe(duplex=False)
        send_msg = 'hyenas'
        read_proc = multiprocessing.Process(target=read_actor, args=(reader, send_msg))
        read_proc.start()
        writer.send(send_msg)
        read_proc.join()
        self.assertEqual(read_proc.exitcode, SUCCESS)

    def test_single_duplex(self):
        local, remote = disc.Pipe()

        echo_proc = multiprocessing.Process(target=echo_actor, args=(remote,))
        echo_proc.start()

        msg = 'hyenas'
        local.send(msg)
        recvmsg = local.recv()

        local.close()
        echo_proc.join()

        self.assertEqual(msg, recvmsg)
        self.assertEqual(echo_proc.exitcode, SUCCESS)

    def test_single_swap(self):
        firstlocal, firstremote = disc.Pipe()
        secondlocal, secondremote = disc.Pipe()
        first_msg = 'hyenas'
        second_msg = 'weasels'
        swap_proc = multiprocessing.Process(target=swap_actor, args=(firstremote, secondremote))
        swap_proc.start()
        firstlocal.send(first_msg)
        secondlocal.send(second_msg)
        frep = firstlocal.recv()
        srep = secondlocal.recv()
        swap_proc.join()
        self.assertEqual(frep, second_msg)
        self.assertEqual(srep, first_msg)


if __name__ == '__main__':
    multiprocessing.set_start_method('spawn')
    unittest.main()

test/connection/standalone_conn_test.py
DragonHPC/dragon
import dragon
import multiprocessing
import time
import sys

import dragon.infrastructure.parameters as dparm
import random


def fp(x):
    print(x)
    sys.stdout.flush()


def debug_it():
    delay = random.uniform(1, 10)
    fp(f'p_uid = {dparm.this_process.my_puid} will be delaying {delay} seconds')
    time.sleep(delay)
    breakpoint()  # people might want tools to conditionally enable breakpoints.
    fp(f'How fun it is here in {dparm.this_process.my_puid}')


class IHateFives(Exception):
    pass


def except_it():
    delay = random.uniform(1, 10)
    fp(f'p_uid = {dparm.this_process.my_puid} will be delaying {delay} seconds')
    time.sleep(delay)

    if dparm.this_process.my_puid == 5:
        raise IHateFives

    fp(f'Leaving {dparm.this_process.my_puid}')


def main():
    multiprocessing.set_start_method('dragon')
    # multiprocessing.set_start_method('spawn')

    nproc = 3

    procs = [multiprocessing.Process(target=debug_it) for _ in range(nproc)]

    # procs = [multiprocessing.Process(target=except_it) for _ in range(nproc)]

    for p in procs:
        p.start()

    # will need to join and release all of these in the debugger
    for p in procs:
        p.join()

    return 0


if __name__ == '__main__':
    exit(main())

test/debug/example.py
DragonHPC/dragon
#!/usr/bin/env python3

import pdb
import sys

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.infrastructure.parameters as dp
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.debug_support as dds
import dragon.utils as du


def make_fake_inf_pool():
    pool = dmm.MemoryPool.create(2 ** 30, 'mtp', dfacts.infrastructure_pool_muid_from_index(0))
    dp.this_process.inf_pd = du.B64.bytes_to_str(pool.serialize())
    be_ch = dch.Channel(mem_pool=pool, c_uid=dfacts.launcher_cuid_from_index(0))
    dp.this_process.local_be_cd = du.B64.bytes_to_str(be_ch.serialize())


def main():
    print('start')
    x = 17
    sys.breakpointhook = dds.dragon_debug_hook
    dds._TESTING_DEBUG_HOOK = True
    make_fake_inf_pool()

    print('breakpoint the first')
    breakpoint()

    y = 27

    print('breakpoint the second')
    breakpoint()

    print('done')
    return 0


if __name__ == '__main__':
    exit(main())

test/debug/tester.py
DragonHPC/dragon
#!/usr/bin/env python3
"""Test script for the GS group API in single node mode"""

import copy
import inspect
import logging
import multiprocessing
import os
import threading
import unittest

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.globalservices.api_setup as dapi
import dragon.globalservices.server as dserver
import dragon.infrastructure.facts as dfacts
import dragon.dlogging.util as dlog
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.parameters as dparm
from dragon.infrastructure.node_desc import NodeDescriptor

import support.util as tsu

from dragon.globalservices.process import get_create_message, multi_join
from dragon.globalservices.process import create as process_create
from dragon.globalservices.group import create as group_create
from dragon.globalservices.group import GroupError, add_to, create_add_to, remove_from, destroy_remove_from, kill, destroy
from dragon.infrastructure.group_desc import GroupDescriptor
from dragon.infrastructure.process_desc import ProcessDescriptor
from dragon.infrastructure.policy import Policy

from dragon.utils import B64


def bringup_channels(gs_stdout, env_updates, channel_overrides, logname=""):
    dlog.setup_logging(basename="gs_" + logname, level=logging.DEBUG, force=True)
    log = logging.getLogger("gs entry")
    log.info("starting")

    dparm.this_process = dparm.LaunchParameters.from_env(env_updates)

    # reconstitute overrides here.
    test_connections = {}
    for uid, stuff in channel_overrides.items():
        _, ser_chan, reading = stuff

        chan = dch.Channel.attach(ser_chan)

        if reading:
            test_connections[uid] = dconn.Connection(inbound_initializer=chan)
        else:
            test_connections[uid] = dconn.Connection(outbound_initializer=chan)

    the_ctx = dserver.GlobalContext()

    try:
        the_ctx.run_startup(mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE, test_gs_stdout=gs_stdout)

        the_ctx.run_global_server(
            mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE,
            test_gs_stdout=gs_stdout,
            test_conns=test_connections,
        )

        the_ctx.run_teardown(test_gs_stdout=gs_stdout)
        log.info("normal exit")
    except:
        log.exception("fatal")
        gs_stdout.send(dmsg.AbnormalTermination(tag=0).serialize())


class GSGroupAPI(unittest.TestCase):
    def setUp(self) -> None:
        self.gs_stdout_rh, self.gs_stdout_wh = multiprocessing.Pipe(duplex=False)

        self.node_sdesc = NodeDescriptor.make_for_current_node(is_primary=True).sdesc

        self.some_parms = copy.copy(dparm.this_process)

        username = os.environ.get("USER", str(os.getuid()))
        self.pool_name = "spac_" + username
        self.pool_size = 2**30
        self.pool_prealloc_blocks = None
        self.pool_uid = 17

        self.mpool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid, self.pool_prealloc_blocks)
        self.some_parms.inf_pd = B64.bytes_to_str(self.mpool.serialize())

        def mk_handles(cuid):
            chan = dch.Channel(self.mpool, cuid)
            rh = dconn.Connection(inbound_initializer=chan)
            wh = dconn.Connection(outbound_initializer=chan)
            envser = B64.bytes_to_str(chan.serialize())
            return chan, rh, wh, envser

        handles = mk_handles(dfacts.GS_INPUT_CUID)
        self.gs_input_chan, self.gs_input_rh, self.gs_input_wh, self.some_parms.gs_cd = handles

        handles = mk_handles(dfacts.BASE_BE_CUID)
        self.bela_input_chan, self.bela_input_rh, self.bela_input_wh, self.some_parms.local_be_cd = handles

        handles = mk_handles(dfacts.BASE_SHEP_CUID)
        self.shep_input_chan, self.shep_input_rh, self.shep_input_wh, self.some_parms.local_shep_cd = handles

        self.gs_return_cuid = 17
        handles = mk_handles(self.gs_return_cuid)
        (
            self.proc_gs_return_chan,
            self.proc_gs_return_rh,
            self.proc_gs_return_wh,
            self.some_parms.gs_ret_cd,
        ) = handles

        dapi.test_connection_override(
            test_gs_input=self.gs_input_wh,
            test_gs_return=self.proc_gs_return_rh,
            test_gs_return_cuid=self.gs_return_cuid,
            test_shep_input=self.shep_input_wh,
        )

        self.tag = 0
        self.dut = None
        self._start_dut()
        self._start_head()

    def tearDown(self) -> None:
        self._kill_head()
        self._teardown_dut()
        self.gs_stdout_rh.close()
        self.gs_stdout_wh.close()

        # connections
        self.gs_input_rh.close()
        self.gs_input_wh.close()
        self.bela_input_rh.close()
        self.bela_input_wh.close()
        self.shep_input_rh.close()
        self.shep_input_wh.close()

        self.proc_gs_return_rh.close()
        self.proc_gs_return_wh.close()

        # underlying channels
        self.gs_input_chan.destroy()
        self.bela_input_chan.destroy()
        self.shep_input_chan.destroy()
        self.proc_gs_return_chan.destroy()
        self.mpool.destroy()

        if self.dut is not None:
            self.dut.join(0)

    def next_tag(self):
        tmp = self.tag
        self.tag += 1
        return tmp

    def _start_dut(self):
        test_name = self.__class__.__name__ + "_" + inspect.stack()[1][0].f_code.co_name

        reconst_gs_return_wh = (self.mpool.serialize(), self.proc_gs_return_chan.serialize(), False)

        test_overrides = {self.gs_return_cuid: reconst_gs_return_wh}

        the_env_wanted = self.some_parms.env()

        self.dut = multiprocessing.Process(
            target=bringup_channels,
            args=(self.gs_stdout_wh, the_env_wanted, test_overrides),
            kwargs={"logname": test_name},
            name="globalservices",
        )

        self.dut.start()

        tsu.get_and_check_type(self.shep_input_rh, dmsg.GSPingSH)

        self.gs_input_wh.send(dmsg.SHPingGS(tag=0, node_sdesc=self.node_sdesc).serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSIsUp)

    def _teardown_dut(self):
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self.next_tag()).serialize())

        tsu.get_and_check_type(self.gs_stdout_rh, dmsg.GSHalted)

    def _start_head(self):
        create_msg = dmsg.GSProcessCreate(
            tag=self.next_tag(),
            p_uid=dfacts.LAUNCHER_PUID,
            r_c_uid=dfacts.BASE_BE_CUID,
            exe="dummy",
            args=[],
            env={},
            user_name="dummy",
        )

        self.gs_input_wh.send(create_msg.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
        self.assertEqual(shep_msg.exe, create_msg.exe)
        self.assertEqual(shep_msg.args, create_msg.args)

        shep_reply_msg = dmsg.SHProcessCreateResponse(
            tag=0, ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_reply_msg = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessCreateResponse)
        self.assertEqual(create_reply_msg.ref, create_msg.tag, "tag-ref mismatch")
        self.assertEqual(create_reply_msg.err, dmsg.GSProcessCreateResponse.Errors.SUCCESS)
        self.assertEqual(shep_msg.t_p_uid, create_reply_msg.desc.p_uid)

        self.head_puid = create_reply_msg.desc.p_uid

        # pretend the test driver is the head process.
        dparm.this_process.my_puid = self.head_puid

    def _kill_head(self):
        death_msg = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=self.head_puid)

        self.gs_input_wh.send(death_msg.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)

    def _create_proc(self, proc_name):
        def create_wrap(the_exe, the_run_dir, the_args, the_env, the_name, result_list):
            res = process_create(
                exe=the_exe, run_dir=the_run_dir, args=the_args, env=the_env, user_name=the_name
            )
            result_list.append(res)

        create_result = []
        create_thread = threading.Thread(
            target=create_wrap, args=("test", "/tmp", ["foo", "bar"], {}, proc_name, create_result)
        )
        create_thread.start()

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
        shep_reply_msg = dmsg.SHProcessCreateResponse(
            tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_thread.join()
        desc = create_result[0]
        return desc

    def _send_get_responses(self, nitems, result):
        if result == 'fail':
            for i in range(nitems):
                shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
                shep_reply_msg = dmsg.SHProcessCreateResponse(
                    tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.FAIL,
                    err_info='simulated failure'
                )
                self.gs_input_wh.send(shep_reply_msg.serialize())
        elif result == 'success':
            for i in range(nitems):
                shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
                shep_reply_msg = dmsg.SHProcessCreateResponse(
                    tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
                )
                self.gs_input_wh.send(shep_reply_msg.serialize())
        else:
            # we create half failed processes and half successful ones
            for i in range(0, nitems//2):
                shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
                shep_reply_msg = dmsg.SHProcessCreateResponse(
                    tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
                )
                self.gs_input_wh.send(shep_reply_msg.serialize())
            for i in range(nitems//2, nitems):
                shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
                shep_reply_msg = dmsg.SHProcessCreateResponse(
                    tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.FAIL,
                    err_info='simulated failure'
                )
                self.gs_input_wh.send(shep_reply_msg.serialize())

    def _create_group(self, group_items, group_policy, group_name, existing=False):
        def create_wrap(items, policy, the_name, result_list):
            res = group_create(items, policy, user_name=the_name)
            result_list.append(res)

        create_result = []
        create_thread = threading.Thread(
            target=create_wrap, args=(group_items, group_policy, group_name, create_result)
        )
        create_thread.start()

        # get the number of members inside this group in total
        nitems = 0
        for lst in group_items:
            nitems += lst[0]

        # in the case that we're trying to create a group that already exists,
        # we don't need the following code
        if not existing: # not already existing group
            self._send_get_responses(nitems, 'success')

        create_thread.join()

        desc = create_result[0]
        return desc

    def _create_group_with_failed_processes(self, group_items, group_policy, group_name, fail):
        def create_wrap(items, policy, the_name, result_list):
            res = group_create(items, policy, user_name=the_name)
            result_list.append(res)

        create_result = []
        create_thread = threading.Thread(
            target=create_wrap, args=(group_items, group_policy, group_name, create_result)
        )
        create_thread.start()

        # get the number of members inside this group in total
        nitems = 0
        for lst in group_items:
            nitems += lst[0]

        self._send_get_responses(nitems, fail)

        create_thread.join()

        desc = create_result[0]
        return desc

    def test_create(self):
        n = 10
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")

        self.assertEqual(descr.name, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(type(descr.sets[0]), list)
        self.assertEqual(len(descr.sets), 1) # 1 list in the sets
        self.assertEqual(len(descr.sets[0]), n) # n members inside the list
        self.assertEqual(all(isinstance(item, GroupDescriptor.GroupMember) for item in descr.sets[0]), True)
        self.assertEqual(all(isinstance(item.desc, ProcessDescriptor) for item in descr.sets[0]), True)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in descr.sets[0]), True)
        self.assertEqual(all(item.desc.p_uid == item.uid for item in descr.sets[0]), True)
        self.assertEqual(all(item.error_code == dmsg.GSProcessCreateResponse.Errors.SUCCESS.value for item in descr.sets[0]), True)

        self.assertEqual(descr.policy.placement, policy.Placement.DEFAULT)
        self.assertEqual(descr.policy.host_name, '')
        self.assertEqual(descr.policy.host_id, -1)
        self.assertEqual(descr.policy.distribution, Policy.Distribution.DEFAULT)

        # test each member's name when a user_name is provided by the user
        for lst_idx, lst in enumerate(descr.sets):
            for item_idx, item in enumerate(lst):
                self.assertEqual(item.desc.name, f'solver.{descr.g_uid}.{lst_idx}.{item_idx}')

    def test_create_multiple_tuples(self):
        n1 = 10
        n2 = 20
        process_msg1 = get_create_message(exe="test", run_dir="/tmp", args=["foo1", "bar1"], env={})
        process_msg2 = get_create_message(exe="test", run_dir="/tmp", args=["foo2", "bar2"], env={})
        items = [(n1, process_msg1.serialize()), (n2, process_msg2.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, None)

        self.assertEqual(descr.name, dfacts.DEFAULT_GROUP_NAME_BASE + str(dfacts.FIRST_GUID)) # check that the default group name was given
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(all(isinstance(lst, list) for lst in descr.sets), True)
        self.assertEqual(len(descr.sets), 2) # 2 lists in the sets
        self.assertEqual(len(descr.sets[0]), n1) # n1 members inside the first list
        self.assertEqual(len(descr.sets[1]), n2) # n2 members inside the second list
        self.assertEqual(len(descr.sets[0])+len(descr.sets[1]), n1+n2) # n1+n2 members in total

        group_puids = []
        for lst in descr.sets:
            self.assertEqual(all(isinstance(item, GroupDescriptor.GroupMember) for item in lst), True)
            self.assertEqual(all(isinstance(item.desc, ProcessDescriptor) for item in lst), True)
            self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in lst), True)
            self.assertEqual(all(item.error_code == dmsg.GSProcessCreateResponse.Errors.SUCCESS.value for item in lst), True)
            self.assertEqual(all(item.desc.p_uid == item.uid for item in lst), True)

            group_puids.extend([item.uid for item in lst if item.desc.state == ProcessDescriptor.State.ACTIVE])

            # test each member's name when a user_name is NOT provided by the user and the default is used
            for item in lst:
                self.assertEqual(item.desc.name, f'dragon_process_{item.desc.p_uid}')

                # now, kill the process
                death_msg = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=item.uid)
                self.gs_input_wh.send(death_msg.serialize())

        multi_join(group_puids, join_all=True)

        self.assertEqual(descr.policy.placement, policy.Placement.DEFAULT)
        self.assertEqual(descr.policy.host_name, '')
        self.assertEqual(descr.policy.host_id, -1)
        self.assertEqual(descr.policy.distribution, Policy.Distribution.DEFAULT)

    def test_create_when_group_already_exists(self):
        n = 5
        descriptors = []
        # attempt to create two identical groups
        # the first group will be created successfully
        # the second will just return the group descriptor immediately
        existing = False
        for i in range(2):
            if i == 1:
                existing = True
            process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
            items = [(n, process_msg.serialize())]
            policy = Policy()
            descr = self._create_group(items, policy, "bob", existing)
            descriptors.append(descr)
        self.assertEqual(descriptors[0].name, descriptors[1].name)
        self.assertEqual(descriptors[0].g_uid, descriptors[1].g_uid)

    def test_create_with_failure(self):
        """All the processes/members of the group fail to create"""

        n = 5
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group_with_failed_processes(items, policy, "bob", "fail")

        self.assertEqual(descr.name, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(type(descr.sets[0]), list)
        self.assertEqual(len(descr.sets), 1) # 1 list in the sets
        self.assertEqual(len(descr.sets[0]), n) # n members inside the list
        self.assertEqual(all(isinstance(item, GroupDescriptor.GroupMember) for item in descr.sets[0]), True)
        self.assertEqual(all(isinstance(item.desc, ProcessDescriptor) for item in descr.sets[0]), True)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.DEAD for item in descr.sets[0]), True)
        self.assertEqual(all(item.error_code == dmsg.GSProcessCreateResponse.Errors.FAIL.value for item in descr.sets[0]), True)
        self.assertEqual(all(item.desc.p_uid == item.uid for item in descr.sets[0]), True)

        self.assertEqual(descr.policy.placement, policy.Placement.DEFAULT)
        self.assertEqual(descr.policy.host_name, '')
        self.assertEqual(descr.policy.host_id, -1)
        self.assertEqual(descr.policy.distribution, Policy.Distribution.DEFAULT)

        # test each member's name when a user_name is provided by the user
        for lst_idx, lst in enumerate(descr.sets):
            for item_idx, item in enumerate(lst):
                self.assertEqual(item.desc.name, f'solver.{descr.g_uid}.{lst_idx}.{item_idx}')

    def test_create_with_half_failures(self):
        n = 10
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group_with_failed_processes(items, policy, "bob", "mix")

        self.assertEqual(descr.name, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(type(descr.sets[0]), list)
        self.assertEqual(len(descr.sets), 1) # 1 list in the sets
        self.assertEqual(len(descr.sets[0]), n) # n members inside the list
        self.assertEqual(all(isinstance(item, GroupDescriptor.GroupMember) for item in descr.sets[0]), True)
        self.assertEqual(all(isinstance(item.desc, ProcessDescriptor) for item in descr.sets[0]), True)
        self.assertEqual(all(item.desc.p_uid == item.uid for item in descr.sets[0]), True)

        self.assertEqual(all(descr.sets[0][i].desc.state == ProcessDescriptor.State.ACTIVE for i in range(0, n//2)), True)
        self.assertEqual(all(descr.sets[0][i].desc.state == ProcessDescriptor.State.DEAD for i in range(n//2, n)), True)
        self.assertEqual(all(descr.sets[0][i].error_code == dmsg.GSProcessCreateResponse.Errors.SUCCESS.value for i in range(0, n//2)), True)
        self.assertEqual(all(descr.sets[0][i].error_code == dmsg.GSProcessCreateResponse.Errors.FAIL.value for i in range(n//2, n)), True)

        self.assertEqual(descr.policy.placement, policy.Placement.DEFAULT)
        self.assertEqual(descr.policy.host_name, '')
        self.assertEqual(descr.policy.host_id, -1)
        self.assertEqual(descr.policy.distribution, Policy.Distribution.DEFAULT)

        # test each member's name when a user_name is provided by the user
        for lst_idx, lst in enumerate(descr.sets):
            for item_idx, item in enumerate(lst):
                self.assertEqual(item.desc.name, f'solver.{descr.g_uid}.{lst_idx}.{item_idx}')

    def test_create_invalid_members(self):
        n = 10
        # send inappropriate message type
        process_msg = dmsg.SHProcessCreate(tag=dapi.next_tag(),
                                           p_uid=dfacts.GS_PUID,
                                           r_c_uid=dfacts.GS_INPUT_CUID,
                                           t_p_uid=dfacts.FIRST_PUID+1000,
                                           exe="test",
                                           args=["foo", "bar"])
        items = [(n, process_msg.serialize())]
        policy = Policy()
        with self.assertRaises(GroupError):
            descr = group_create(items, policy, "bob")

    def test_create_multiple_groups(self):
        n = 10
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        descr1 = self._create_group(items, policy, "alice")

        self.assertEqual(descr.name, "bob")
        self.assertEqual(descr1.name, "alice")

    def _kill_a_thing(self, identifier, n):
        def kill_wrap(ident, result_list):
            # send a kill request for the group
            res = kill(ident)
            result_list.append(res)

        kill_result = []
        kill_thread = threading.Thread(target=kill_wrap, args=(identifier, kill_result))
        kill_thread.start()

        for _ in range(n):
            sh_kill_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessKill)
            sh_kill_reply = dmsg.SHProcessKillResponse(
                tag=self.next_tag(), ref=sh_kill_msg.tag, err=dmsg.SHProcessKillResponse.Errors.SUCCESS
            )
            self.gs_input_wh.send(sh_kill_reply.serialize())

        kill_thread.join()
        return kill_result[0]

    def test_kill(self):
        n = 10
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in descr.sets[0]), True)
        self.assertEqual(descr.name, "bob")

        # kill an existing group of resources
        descr = self._kill_a_thing(descr.g_uid, n)
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in descr.sets[0]), True)

    def test_kill_non_existing_group(self):
        # try to destroy a non-existing group
        with self.assertRaises(GroupError):
            kill("doe")

    def test_kill_a_pending_group(self):
        # try to destroy a group that is still pending
        n = 20
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()

        def create_wrap(items, policy, the_name, result_list):
            res = group_create(items, policy, user_name=the_name)
            result_list.append(res)

        create_result = []
        create_thread = threading.Thread(
            target=create_wrap, args=(items, policy, "bob", create_result)
        )
        create_thread.start()

        # immediately try to destroy the group, that is still pending
        # because the create is not complete yet; this should raise
        with self.assertRaises(GroupError):
            kill("bob")

        # complete the construction of the group
        self._send_get_responses(n, 'success')
        create_thread.join()

        descr = create_result[0]
        self.assertEqual(descr.name, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)

    def _destroy_a_thing(self, identifier, n, sets):
        def destroy_wrap(ident, result_list):
            res = destroy(ident)
            result_list.append(res)

        destroy_result = []
        destroy_thread = threading.Thread(target=destroy_wrap, args=(identifier, destroy_result))
        destroy_thread.start()

        for _ in range(n):
            sh_kill_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessKill)
            sh_kill_reply = dmsg.SHProcessKillResponse(
                tag=self.next_tag(), ref=sh_kill_msg.tag, err=dmsg.SHProcessKillResponse.Errors.SUCCESS
            )
            self.gs_input_wh.send(sh_kill_reply.serialize())

        for item in sets:
            death_msg = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=item.uid)
            self.gs_input_wh.send(death_msg.serialize())

        destroy_thread.join()
        return destroy_result[0]

    def test_destroy(self):
        n = 10
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in descr.sets[0]), True)
        self.assertEqual(descr.name, "bob")

        # destroy an existing group of resources
        descr = self._destroy_a_thing(descr.g_uid, n, descr.sets[0])
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.DEAD for item in descr.sets[0]), True)
        self.assertEqual(int(descr.state), GroupDescriptor.State.DEAD)

    def test_destroy_alt(self):
        # first call kill on the group and then destroy
        n = 10
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in descr.sets[0]), True)
        self.assertEqual(descr.name, "bob")

        # kill an existing group of resources
        descr = self._kill_a_thing(descr.g_uid, n)
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in descr.sets[0]), True)

        # now, actually kill the processes
        for item in descr.sets[0]:
            death_msg = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=item.uid)
            self.gs_input_wh.send(death_msg.serialize())

        # destroy an existing group of resources
        descr = destroy("bob")
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.DEAD for item in descr.sets[0]), True)
        self.assertEqual(int(descr.state), GroupDescriptor.State.DEAD)

    def test_destroy_non_existing_group(self):
        # try to destroy a non-existing group
        with self.assertRaises(GroupError):
            destroy("doe")

    def test_destroy_a_pending_group(self):
        # try to destroy a group that is still pending
        n = 20
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()

        def create_wrap(items, policy, the_name, result_list):
            res = group_create(items, policy, user_name=the_name)
            result_list.append(res)

        create_result = []
        create_thread = threading.Thread(
            target=create_wrap, args=(items, policy, "bob", create_result)
        )
        create_thread.start()

        # immediately try to destroy the group, that is still pending
        # because the create is not complete yet; this should raise
        with self.assertRaises(GroupError):
            destroy("bob")

        # complete the construction of the group
        self._send_get_responses(n, 'success')
        create_thread.join()

        descr = create_result[0]
        self.assertEqual(descr.name, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)

    def _send_responses(self, nitems):
            for i in range(nitems):
                shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
                shep_reply_msg = dmsg.SHProcessCreateResponse(
                    tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
                )
                self.gs_input_wh.send(shep_reply_msg.serialize())

    def test_create_add_to(self):
        # first create a group
        n = 5
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")

        # now, let's add processes to the group
        n1 = 3
        add_thread = threading.Thread(target=self._send_responses, args=(n1,))
        add_thread.start()

        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='addition')
        items = [(n1, process_msg.serialize())]
        descr = create_add_to(descr.g_uid, items, policy)
        add_thread.join()

        self.assertEqual(descr.name, "bob")
        self.assertEqual(int(descr.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(type(descr.sets[0]), list)
        self.assertEqual(len(descr.sets), 2) # 2 lists in the sets
        self.assertEqual(len(descr.sets[0]), n) # n members inside the 1st list
        self.assertEqual(len(descr.sets[1]), n1) # n1 members inside the 2nd list
        self.assertEqual(len(descr.sets[0])+len(descr.sets[1]), n1+n) # n1+n members in total

        for lst in descr.sets:
            self.assertEqual(all(isinstance(item, GroupDescriptor.GroupMember) for item in lst), True)
            self.assertEqual(all(isinstance(item.desc, ProcessDescriptor) for item in lst), True)
            self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in lst), True)
            self.assertEqual(all(item.error_code == dmsg.GSProcessCreateResponse.Errors.SUCCESS.value for item in lst), True)
            self.assertEqual(all(item.desc.p_uid == item.uid for item in lst), True)

        # test each member's name when a user_name is provided by the user
        for lst_idx, lst in enumerate(descr.sets):
            for item_idx, item in enumerate(lst):
                if lst_idx == 0:
                    self.assertEqual(item.desc.name, f'solver.{descr.g_uid}.{lst_idx}.{item_idx}')
                elif lst_idx == 1:
                    self.assertEqual(item.desc.name, f'addition.{descr.g_uid}.{lst_idx}.{item_idx}')

    def test_add_to(self):
        # first create a group
        n = 5
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")

        # create some processes
        items = []
        n1 = 3
        for i in range(n1):
            pdescr = self._create_proc(f"alice{i}")
            self.assertEqual(pdescr.name, f"alice{i}")
            items.append(pdescr.p_uid)

        # now, let's add these processes to the group
        descr1 = add_to(descr.g_uid, items)

        self.assertEqual(descr1.name, "bob")
        self.assertEqual(int(descr1.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(len(descr1.sets), 4) # 4 lists in the sets: 1+n1
        self.assertEqual(len(descr1.sets[0]), n) # n members inside the 1st list

        for lst in descr1.sets:
            self.assertEqual(all(isinstance(item, GroupDescriptor.GroupMember) for item in lst), True)
            self.assertEqual(all(isinstance(item.desc, ProcessDescriptor) for item in lst), True)
            self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in lst), True)
            self.assertEqual(all(item.error_code == dmsg.GSProcessCreateResponse.Errors.SUCCESS.value for item in lst), True)
            self.assertEqual(all(item.desc.p_uid == item.uid for item in lst), True)

        # test each member's name when a user_name is provided by the user
        for lst_idx, lst in enumerate(descr1.sets):
            for item_idx, item in enumerate(lst):
                if lst_idx == 0:
                    self.assertEqual(item.desc.name, f'solver.{descr.g_uid}.{lst_idx}.{item_idx}')
                else:
                    self.assertEqual(item.desc.name, f"alice{lst_idx-1}")

    def test_remove_from(self):
        # first create a group
        n = 8
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n) # descr should have 8 items in descr.sets

        group_puids = []
        for lst in descr.sets:
            group_puids.extend([item.uid for item in lst if item.desc.state == ProcessDescriptor.State.ACTIVE])

        # now, let's remove a few processes
        # remove specific processes
        name = f'solver.{descr.g_uid}.0.0' # this should be the name of the first created process
        descr1 = remove_from(descr.g_uid, [name, group_puids[6]])
        self.assertEqual(len(descr1.sets[0]), 6) # descr1 should have 6 items in descr1.sets

        # update group_puids
        del group_puids[6]
        del group_puids[0]

        # remove the first n1 processes
        n1 = 3
        descr2 = remove_from(descr1.g_uid, group_puids[0:n1])
        self.assertEqual(len(descr2.sets[0]), 3) # descr2 should have 3 items in descr2.sets

        # update group_puids
        del group_puids[0:n1]

        # compare group_puids with the final/remaining members of the group
        for i in range(len(descr2.sets[0])):
            assert group_puids[i] == descr2.sets[0][i].uid

        self.assertEqual(len(descr2.sets[0]), len(group_puids))

    def test_remove_when_nonexist(self):
        # try to remove processes that are not members of the group
        # first create a group
        n = 5
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        # try to remove non-existing processes
        name = f'solver1.{descr.g_uid}.0.0'
        with self.assertRaises(GroupError):
            remove_from(descr.g_uid, [name])

    def test_remove_non_existing_procs(self):
        # first create a group
        n = 3
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        # try to remove one valid and two non-valid procs
        with self.assertRaises(GroupError):
            remove_from(descr.g_uid, [descr.sets[0][0].uid, "rand1", "rand2"])

    def _send_kill_messages(self, n):
        for _ in range(n):
            sh_kill_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessKill)
            sh_kill_reply = dmsg.SHProcessKillResponse(
                tag=self.next_tag(), ref=sh_kill_msg.tag, err=dmsg.SHProcessKillResponse.Errors.SUCCESS
            )
            self.gs_input_wh.send(sh_kill_reply.serialize())

    def test_destroy_remove_from(self):
        # first create a group
        n = 8
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n) # descr should have 8 items in descr.sets

        group_puids = []
        for lst in descr.sets:
            group_puids.extend([item.uid for item in lst if item.desc.state == ProcessDescriptor.State.ACTIVE])

        # now, let's remove a few processes
        n1 = 2
        kill_thread = threading.Thread(target=self._send_kill_messages, args=(n1,))
        kill_thread.start()
        # remove specific processes
        name = f'solver.{descr.g_uid}.0.0' # this should be the name of the first created process
        descr1 = destroy_remove_from(descr.g_uid, [name, group_puids[6]])
        self.assertEqual(len(descr1.sets[0]), 6) # descr1 should have 6 items in descr1.sets
        kill_thread.join()

        # update group_puids
        del group_puids[6]
        del group_puids[0]

        # remove the first n1 processes
        n1 = 3
        kill_thread = threading.Thread(target=self._send_kill_messages, args=(n1,))
        kill_thread.start()
        descr2 = destroy_remove_from(descr1.g_uid, group_puids[0:n1])
        self.assertEqual(len(descr2.sets[0]), 3) # descr2 should have 3 items in descr2.sets
        kill_thread.join()

        # update group_puids
        del group_puids[0:n1]

        # compare group_puids with the final/remaining members of the group
        for i in range(len(descr2.sets[0])):
            assert group_puids[i] == descr2.sets[0][i].uid

        self.assertEqual(len(descr2.sets[0]), len(group_puids))
        self.assertEqual(int(descr2.state), GroupDescriptor.State.ACTIVE)
        for lst in descr2.sets:
            group_puids.extend([item.uid for item in lst if item.desc.state == ProcessDescriptor.State.ACTIVE])

    def test_destroy_remove_from_non_existing_group(self):
        # try to destroy_remove_from a non-existing group
        with self.assertRaises(GroupError):
            destroy_remove_from("doe", ["rand"])

    def test_destroy_remove_from_already_dead_procs(self):
        # first create a group
        n = 5
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        puids_list = []
        for item in descr.sets[0]:
            puids_list.append(item.uid)

        # kill all the procs that belong to the group without letting the group know
        for item in puids_list:
            death_msg = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=item)
            self.gs_input_wh.send(death_msg.serialize())

        # and now try to remove them from the group
        descr1 = destroy_remove_from(descr.g_uid, puids_list)
        # we end up with an active but empty group
        self.assertEqual(int(descr1.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(len(descr1.sets), 0)

    def test_destroy_remove_from_comb_already_dead_procs(self):
        # we'll ask to remove procs from which some are already dead
        # and some are normal

        # first create a group
        n = 5
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        puids_list = []
        for item in descr.sets[0]:
            puids_list.append(item.uid)

        # kill some of the procs that belong to the group without letting the group know
        for item in puids_list[:n-2]:
            death_msg = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=item)
            self.gs_input_wh.send(death_msg.serialize())

        kill_thread = threading.Thread(target=self._send_kill_messages, args=(1,))
        kill_thread.start()

        # remove all the members except for the last one
        descr1 = destroy_remove_from(descr.g_uid, puids_list[:n-1])
        kill_thread.join()
        self.assertEqual(int(descr1.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(len(descr1.sets[0]), 1)

        # compare with the initial group.sets
        for i, item in enumerate(descr.sets[0]):
            if i < n-1:
                assert item not in descr1.sets[0]
            else:
                self.assertEqual(item.uid, descr1.sets[0][0].uid)

    def test_destroy_remove_from_non_existing_comb_procs(self):
        # first create a group
        n = 3
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        # try to remove one valid and two non-valid procs
        with self.assertRaises(GroupError):
            destroy_remove_from(descr.g_uid, [descr.sets[0][0].uid, "rand1", "rand2"])

    def test_destroy_remove_from_non_existing_procs(self):
        # first create a group
        n = 3
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        # try to remove one valid and two non-valid procs
        with self.assertRaises(GroupError):
            destroy_remove_from(descr.g_uid, ["rand1", "rand2"])

    def test_remove_from_and_then_destroy_group(self):
        # first create a group
        n = 8
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        group_puids = []
        for lst in descr.sets:
            group_puids.extend([item.uid for item in lst if item.desc.state == ProcessDescriptor.State.ACTIVE])

        # now, let's remove a few processes
        # remove specific processes
        name = f'solver.{descr.g_uid}.0.0' # this should be the name of the first created process
        descr1 = remove_from(descr.g_uid, [name, group_puids[6]])
        self.assertEqual(len(descr1.sets[0]), 6) # descr1 should have 6 items in descr1.sets

        # update group_puids
        del group_puids[6]
        del group_puids[0]

        # now, let's destroy the group
        descr2 = self._destroy_a_thing(descr1.g_uid, n-2, descr1.sets[0])
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.DEAD for item in descr2.sets[0]), True)
        self.assertEqual(int(descr2.state), GroupDescriptor.State.DEAD)

    def test_kill_then_add_to(self):
        # create, kill, createAdd

        # first create a group
        n = 8
        process_msg = get_create_message(exe="test", run_dir="/tmp", args=["foo", "bar"], env={}, user_name='solver')
        items = [(n, process_msg.serialize())]
        policy = Policy()
        descr = self._create_group(items, policy, "bob")
        self.assertEqual(len(descr.sets[0]), n)

        # next call kill on the group's processes
        descr1 = self._kill_a_thing(descr.g_uid, n)
        self.assertEqual(int(descr1.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in descr1.sets[0]), True)

        # last, try to add a process to the group
        # create some processes
        items = []
        n1 = 3
        for i in range(n1):
            pdescr = self._create_proc(f"alice{i}")
            self.assertEqual(pdescr.name, f"alice{i}")
            items.append(pdescr.p_uid)

        # now, let's add these processes to the group
        descr2 = add_to(descr.g_uid, items)

        self.assertEqual(descr2.name, "bob")
        self.assertEqual(int(descr2.state), GroupDescriptor.State.ACTIVE)
        self.assertEqual(len(descr2.sets), 4) # 4 lists in the sets: 1+n1
        self.assertEqual(len(descr2.sets[0]), n) # n members inside the 1st list

        for lst in descr2.sets:
            self.assertEqual(all(isinstance(item, GroupDescriptor.GroupMember) for item in lst), True)
            self.assertEqual(all(isinstance(item.desc, ProcessDescriptor) for item in lst), True)
            self.assertEqual(all(item.desc.state == ProcessDescriptor.State.ACTIVE for item in lst), True)
            self.assertEqual(all(item.error_code == dmsg.GSProcessCreateResponse.Errors.SUCCESS.value for item in lst), True)
            self.assertEqual(all(item.desc.p_uid == item.uid for item in lst), True)

        # test each member's name when a user_name is provided by the user
        for lst_idx, lst in enumerate(descr2.sets):
            for item_idx, item in enumerate(lst):
                if lst_idx == 0:
                    self.assertEqual(item.desc.name, f'solver.{descr.g_uid}.{lst_idx}.{item_idx}')
                else:
                    self.assertEqual(item.desc.name, f"alice{lst_idx-1}")

test/globalservices/group_api.py
DragonHPC/dragon
#!/usr/bin/env python3

# This is a target program for a managed process test
# of large argument delivery.
# it returns with 0 if everything is good
# and with 1 if not, and creates its own logfile too.

import sys

import dragon.globalservices.api_setup as dapi
import dragon.infrastructure.parameters as dparms
import dragon.dlogging.util as dlog
import logging


def show_desc(desc):
    return f"len {len(desc)}: {desc[:4]} ... {desc[-4:]}"


dlog.setup_logging(basename=f"jatt_id{dparms.this_process.my_puid}".format(), level=logging.DEBUG)
log = logging.getLogger("")
log.info("hello from jumbo arg test")

expected_size = int(sys.argv[1])
log.info(f"expecting {expected_size} bytes")

dapi.connect_to_infrastructure()

if len(dapi._ARG_PAYLOAD) == expected_size:
    log.info("payload arrived ok")
else:
    log.error(f"payload was {len(dapi._ARG_PAYLOAD)} bytes")
    exit(1)

log.info("test finished successfully")
exit(0)

test/globalservices/jumbo_arg_test_target.py
DragonHPC/dragon
#!/usr/bin/env python3

# This is a target program for a managed process test.
# It attempts to verify that certain channels are
# constructed and that their descriptors
# appear in the environment in the expected way.

# it returns with 0 if everything is good
# and with 1 if not, and creates its own logfile too.

# by default this checks for the GS channel descriptor
# and whether one can attach to it
# options. pick at most one
# --gs-ret  : check for the gs return channel too
# --args : connect to infrastructure and receive ping message thru api
# If --args is chosen, can add 3 more integer arguments to check
# against the bytes like objects in which immediate args are
# deposited.

import binascii
import logging
import sys

import dragon.channels as dch
import dragon.managed_memory as dmm
import dragon.infrastructure.parameters as dparms
import dragon.dlogging.util as dlog
import dragon.utils as du

import pickle


def show_desc(desc):
    return "len {}: {} ... {}".format(len(desc), desc[:4], desc[-4:])


dlog.setup_logging(basename="mptt_id{}".format(dparms.this_process.my_puid), level=logging.DEBUG)
log = logging.getLogger("")
log.info("hello from test util")

do_args = len(sys.argv) > 1 and sys.argv[1] == "--args"
check_gs_ret = do_args or (len(sys.argv) > 1 and sys.argv[1] == "--gs-ret")

if check_gs_ret:
    log.info("checking gs return channel")

if do_args:
    log.info("checking full infrastructure conn and arg delivery")

gs_desc = None
gs_return_desc = None

log.info(f"assigned uid is {dparms.this_process.my_puid}")

if 0 == len(dparms.this_process.inf_pd):
    log.info("inf pd absent")
    exit(1)

try:
    infpool_desc = du.B64.str_to_bytes(dparms.this_process.inf_pd)
except binascii.Error as err:
    log.info("failed to decode inf pool desc")
    log.exception(f"error is {err}", exc_info=True)
    exit(1)

try:
    infpool = dmm.MemoryPool.attach(infpool_desc)
except Exception as err:
    log.info("Failed to attach to inf pool")
    log.exception(f"error is {err}", exc_info=True)
    exit(1)

if 0 == len(dparms.this_process.default_pd):
    log.info("default pd absent")
    exit(1)

try:
    defp_desc = du.B64.str_to_bytes(dparms.this_process.default_pd)
except binascii.Error as err:
    log.info("failed to decode default pool desc")
    log.exception(f"error is {err}", exc_info=True)
    exit(1)

try:
    defpool = dmm.MemoryPool.attach(defp_desc)
except Exception as err:
    log.info("Failed to attach to default pool")
    log.exception(f"error is {err}", exc_info=True)
    exit(1)


if 0 == len(dparms.this_process.gs_cd):
    log.info("GS cd absent")
    exit(1)

log.info("gs desc is: {}".format(show_desc(dparms.this_process.gs_cd)))

try:
    gs_desc = du.B64.str_to_bytes(dparms.this_process.gs_cd)
except binascii.Error as err:
    log.info("Failed to decode GS cd")
    log.exception(f"error is {err}", exc_info=True)
    exit(1)

try:
    gs_chan = dch.Channel.attach(gs_desc)
    log.info("attached to gs channel")
except Exception as err:
    log.info("Failed to attach to gs channel")
    log.exception(f"error is {err}", exc_info=True)
    exit(1)

if check_gs_ret:
    log.info("checking gs ret channel")

    if 0 == len(dparms.this_process.gs_ret_cd):
        log.info("GS ret cd absent")
        exit(1)

    log.info("gs return desc is: {}".format(show_desc(dparms.this_process.gs_ret_cd)))

    try:
        gs_return_desc = du.B64.str_to_bytes(dparms.this_process.gs_ret_cd)
    except binascii.Error as err:
        log.info("Failed to decode GS return cd")
        log.exception(f"error is {err}", exc_info=True)
        exit(1)

    try:
        gs_ret_chan = dch.Channel.attach(gs_return_desc)
        log.info("attached to gs return channel")
    except Exception as err:
        log.info("Failed to attach to gs return channel")
        log.exception(f"error is {err}", exc_info=True)
        exit(1)

if do_args:
    log.info("trying to attach to infrastructure and args")
    import dragon.globalservices.api_setup as dapi

    try:
        dapi.connect_to_infrastructure()
    except Exception as err:
        log.info("failed to attach to infrastructure")
        log.exception(f"error is {err}", exc_info=True)
        exit(1)

    if 5 == len(sys.argv):
        log.info("checking argument delivery")
        ex_item1 = int(sys.argv[2])
        ex_item2 = int(sys.argv[3])
        ex_item3 = int(sys.argv[4])

        if dapi._ARG_PAYLOAD is None:
            log.info("arg payload not updated")

        item1, item2, item3 = pickle.loads(dapi._ARG_PAYLOAD)

        if item1 != ex_item1:
            log.info("item1 expected: {} got: {}".format(ex_item1, item1))
            exit(1)

        if item2 != ex_item2:
            log.info("args expected: {} got: {}".format(ex_item2, item2))
            exit(1)

        if item3 != ex_item3:
            log.info("kwargs expected: {} got: {}".format(ex_item3, item3))
            exit(1)

log.info("test finished successfully")
exit(0)

test/globalservices/managed_proc_test_target.py
DragonHPC/dragon
#!/usr/bin/env python3
"""Test script for the GS process API in single node mode"""

import time

import copy
import inspect
import logging
import multiprocessing
import os
import threading
import unittest

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.globalservices.api_setup as dapi
import dragon.globalservices.channel as dchannel
import dragon.globalservices.pool as dpool
import dragon.globalservices.process as dproc
import dragon.globalservices.server as dserver
import dragon.infrastructure.facts as dfacts
import dragon.dlogging.util as dlog
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.parameters as dparm
from dragon.infrastructure.node_desc import NodeDescriptor
import support.util as tsu
from dragon.utils import B64


def bringup_channels(gs_stdout, env_updates, channel_overrides, logname=""):
    dlog.setup_logging(basename="gs_" + logname, level=logging.DEBUG, force=True)
    log = logging.getLogger("gs entry")
    log.info("starting")

    dparm.this_process = dparm.LaunchParameters.from_env(env_updates)

    # reconstitute overrides here.
    test_connections = {}
    for uid, stuff in channel_overrides.items():
        _, ser_chan, reading = stuff

        chan = dch.Channel.attach(ser_chan)

        if reading:
            test_connections[uid] = dconn.Connection(inbound_initializer=chan)
        else:
            test_connections[uid] = dconn.Connection(outbound_initializer=chan)

    the_ctx = dserver.GlobalContext()

    try:
        the_ctx.run_startup(mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE, test_gs_stdout=gs_stdout)

        the_ctx.run_global_server(
            mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE,
            test_gs_stdout=gs_stdout,
            test_conns=test_connections,
        )

        the_ctx.run_teardown(test_gs_stdout=gs_stdout)
        log.info("normal exit")
    except:
        log.exception("fatal")
        gs_stdout.send(dmsg.AbnormalTermination(tag=0).serialize())


class SingleProcAPIChannels(unittest.TestCase):
    def setUp(self) -> None:
        self.gs_stdout_rh, self.gs_stdout_wh = multiprocessing.Pipe(duplex=False)

        self.node_sdesc = NodeDescriptor.make_for_current_node(is_primary=True).sdesc

        self.some_parms = copy.copy(dparm.this_process)

        username = os.environ.get("USER", str(os.getuid()))
        self.pool_name = "spac_" + username
        self.pool_size = 2**30
        self.pool_prealloc_blocks = None
        self.pool_uid = 17

        self.mpool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid, self.pool_prealloc_blocks)
        self.some_parms.inf_pd = B64.bytes_to_str(self.mpool.serialize())

        def mk_handles(cuid):
            chan = dch.Channel(self.mpool, cuid)
            rh = dconn.Connection(inbound_initializer=chan)
            wh = dconn.Connection(outbound_initializer=chan)
            envser = B64.bytes_to_str(chan.serialize())
            return chan, rh, wh, envser

        handles = mk_handles(dfacts.GS_INPUT_CUID)
        self.gs_input_chan, self.gs_input_rh, self.gs_input_wh, self.some_parms.gs_cd = handles

        handles = mk_handles(dfacts.BASE_BE_CUID)
        self.bela_input_chan, self.bela_input_rh, self.bela_input_wh, self.some_parms.local_be_cd = handles

        handles = mk_handles(dfacts.BASE_SHEP_CUID)
        self.shep_input_chan, self.shep_input_rh, self.shep_input_wh, self.some_parms.local_shep_cd = handles

        self.gs_return_cuid = 17
        handles = mk_handles(self.gs_return_cuid)
        (
            self.proc_gs_return_chan,
            self.proc_gs_return_rh,
            self.proc_gs_return_wh,
            self.some_parms.gs_ret_cd,
        ) = handles

        dapi.test_connection_override(
            test_gs_input=self.gs_input_wh,
            test_gs_return=self.proc_gs_return_rh,
            test_gs_return_cuid=self.gs_return_cuid,
            test_shep_input=self.shep_input_wh,
        )

        self.tag = 0
        self.dut = None
        self._start_dut()
        self._start_head()

    def tearDown(self) -> None:
        self._kill_head()
        self._teardown_dut()
        self.gs_stdout_rh.close()
        self.gs_stdout_wh.close()

        # connections
        self.gs_input_rh.close()
        self.gs_input_wh.close()
        self.bela_input_rh.close()
        self.bela_input_wh.close()
        self.shep_input_rh.close()
        self.shep_input_wh.close()

        self.proc_gs_return_rh.close()
        self.proc_gs_return_wh.close()

        # underlying channels
        self.gs_input_chan.destroy()
        self.bela_input_chan.destroy()
        self.shep_input_chan.destroy()
        self.proc_gs_return_chan.destroy()
        self.mpool.destroy()

        if self.dut is not None:
            self.dut.join(0)

    def next_tag(self):
        tmp = self.tag
        self.tag += 1
        return tmp

    def _start_dut(self):
        test_name = self.__class__.__name__ + "_" + inspect.stack()[1][0].f_code.co_name

        reconst_gs_return_wh = (self.mpool.serialize(), self.proc_gs_return_chan.serialize(), False)

        test_overrides = {self.gs_return_cuid: reconst_gs_return_wh}

        the_env_wanted = self.some_parms.env()

        self.dut = multiprocessing.Process(
            target=bringup_channels,
            args=(self.gs_stdout_wh, the_env_wanted, test_overrides),
            kwargs={"logname": test_name},
            name="globalservices",
        )

        self.dut.start()

        tsu.get_and_check_type(self.shep_input_rh, dmsg.GSPingSH)

        self.gs_input_wh.send(dmsg.SHPingGS(tag=0, node_sdesc=self.node_sdesc).serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSIsUp)

    def _teardown_dut(self):
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self.next_tag()).serialize())

        tsu.get_and_check_type(self.gs_stdout_rh, dmsg.GSHalted)

    def _start_head(self):
        create_msg = dmsg.GSProcessCreate(
            tag=self.next_tag(),
            p_uid=dfacts.LAUNCHER_PUID,
            r_c_uid=dfacts.BASE_BE_CUID,
            exe="dummy",
            args=[],
            env={},
            user_name="dummy",
        )

        self.gs_input_wh.send(create_msg.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
        self.assertEqual(shep_msg.exe, create_msg.exe)
        self.assertEqual(shep_msg.args, create_msg.args)

        shep_reply_msg = dmsg.SHProcessCreateResponse(
            tag=0, ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_reply_msg = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessCreateResponse)
        self.assertEqual(create_reply_msg.ref, create_msg.tag, "tag-ref mismatch")
        self.assertEqual(create_reply_msg.err, dmsg.GSProcessCreateResponse.Errors.SUCCESS)
        self.assertEqual(shep_msg.t_p_uid, create_reply_msg.desc.p_uid)

        self.head_puid = create_reply_msg.desc.p_uid

        # pretend the test driver is the head process.
        dparm.this_process.my_puid = self.head_puid

    def _kill_head(self):
        death_msg = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=self.head_puid)

        self.gs_input_wh.send(death_msg.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)

    def _kill_a_thing(self, identifier):
        def kill_wrap(ident, result_list):
            res = dproc.kill(ident)
            result_list.append(res)

        kill_result = []
        kill_thread = threading.Thread(target=kill_wrap, args=(identifier, kill_result))
        kill_thread.start()

        sh_kill_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessKill)
        sh_kill_reply = dmsg.SHProcessKillResponse(
            tag=self.next_tag(), ref=sh_kill_msg.tag, err=dmsg.SHProcessKillResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(sh_kill_reply.serialize())

        kill_thread.join()
        return kill_result[0]

    def _create_proc(self, proc_name):
        def create_wrap(the_exe, the_run_dir, the_args, the_env, the_name, result_list):
            res = dproc.create(
                exe=the_exe, run_dir=the_run_dir, args=the_args, env=the_env, user_name=the_name
            )
            result_list.append(res)

        create_result = []
        create_thread = threading.Thread(
            target=create_wrap, args=("test", "/tmp", ["foo", "bar"], {}, proc_name, create_result)
        )
        create_thread.start()

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)

        shep_reply_msg = dmsg.SHProcessCreateResponse(
            tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_thread.join()

        desc = create_result[0]

        return desc

    def _make_a_pool(self, fake_pool_name=""):
        def create_wrap(the_size, the_name, result_list):
            res = dpool.create(size=the_size, user_name=the_name)
            result_list.append(res)

        fake_pool_size = 2**30
        create_result = []
        create_thread = threading.Thread(
            target=create_wrap, args=(fake_pool_size, fake_pool_name, create_result)
        )
        create_thread.start()
        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHPoolCreate)

        dummy_sdesc = B64.bytes_to_str("xxxTESTDUMMYxxx".encode())
        shep_reply_msg = dmsg.SHPoolCreateResponse(
            tag=self.next_tag(),
            ref=shep_msg.tag,
            err=dmsg.SHPoolCreateResponse.Errors.SUCCESS,
            desc=dummy_sdesc,
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_thread.join()

        desc = create_result[0]

        self.pool_m_uid = desc.m_uid

        return desc

    def _make_a_channel(self, m_uid, fake_channel_name=""):
        def create_wrap(the_m_uid, the_name, result_list):
            res = dchannel.create(m_uid=the_m_uid, user_name=the_name)
            result_list.append(res)

        create_result = []
        create_thread = threading.Thread(target=create_wrap, args=(m_uid, fake_channel_name, create_result))
        create_thread.start()
        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHChannelCreate)
        dummy_sdesc = B64.bytes_to_str("xxxTESTDUMMYxxx".encode())
        shep_reply_msg = dmsg.SHChannelCreateResponse(
            tag=self.next_tag(),
            ref=shep_msg.tag,
            err=dmsg.SHChannelCreateResponse.Errors.SUCCESS,
            desc=dummy_sdesc,
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_thread.join()

        desc = create_result[0]

        self.channel_c_uid = desc.m_uid

        return desc

    def test_list(self):
        plist = dproc.get_list()
        self.assertEqual(plist, [self.head_puid])

    def test_query(self):
        plist = dproc.get_list()
        descr = dproc.query(plist[0])
        self.assertEqual(descr.p_uid, self.head_puid)

    def test_create(self):
        descr = self._create_proc("bob")

        self.assertEqual(descr.name, "bob")

    def test_dump_smoke(self):
        descr = self._create_proc("bob")

        self.assertEqual(descr.name, "bob")

        dump_msg = dmsg.GSDump(tag=self.next_tag(), filename="dump_file")
        self.gs_input_wh.send(dump_msg.serialize())

    def test_query_by_name(self):
        proc_name = "bob"
        desc = self._create_proc(proc_name)

        self.assertEqual(desc.name, proc_name)

        desc2 = dproc.query(proc_name)
        self.assertEqual(desc.p_uid, desc2.p_uid)

    def test_heritance(self):
        bobdesc = self._create_proc("bob")
        freddesc = self._create_proc("fred")

        the_children = {bobdesc.p_uid, freddesc.p_uid}

        mydesc = dproc.query(self.head_puid)

        first_live_children = mydesc.live_children

        bob_dies = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=bobdesc.p_uid)
        self.gs_input_wh.send(bob_dies.serialize())

        mydesc_later = dproc.query(self.head_puid)

        self.assertEqual(the_children, first_live_children)
        self.assertEqual(mydesc_later.live_children, {freddesc.p_uid})
        self.assertEqual(bobdesc.p_p_uid, self.head_puid)

    def test_kill(self):
        proc_name = "bob"
        desc = self._create_proc(proc_name)

        self.assertEqual(desc.name, proc_name)

        result = self._kill_a_thing(proc_name)
        self.assertEqual(result, None)

    def test_join(self):
        proc_name = "bob"
        desc = self._create_proc(proc_name)
        self.assertEqual(desc.name, proc_name)
        test_exit_code = 17

        def join_wrap(identifier, result_list):
            res = dproc.join(identifier)
            result_list.append(res)

        join_result = []
        join_thread = threading.Thread(target=join_wrap, args=(proc_name, join_result))
        join_thread.start()

        sh_kill_reply = dmsg.SHProcessExit(tag=self.next_tag(), p_uid=desc.p_uid, exit_code=test_exit_code)

        self.gs_input_wh.send(sh_kill_reply.serialize())
        join_thread.join()
        self.assertEqual(join_result[0], test_exit_code)

    def test_join_timeout(self):
        proc_name = "bob"
        desc = self._create_proc(proc_name)

        self.assertEqual(desc.name, proc_name)

        start = time.monotonic()
        result = dproc.join(proc_name, timeout=1.0)
        self.assertEqual(result, None)
        duration = time.monotonic() - start
        # should time out in about a second
        self.assertGreater(duration, 1.0)

    def test_join_timeout_mult_req(self):
        proc_name = "bob"
        desc = self._create_proc(proc_name)

        self.assertEqual(desc.name, proc_name)

        # join on the same puid multiple times
        for _ in range(5):
            start = time.monotonic()
            res = dproc.join(proc_name, timeout=0.5)
            duration = time.monotonic() - start
            self.assertEqual(res, None)
            self.assertGreater(duration, 0.5)

    def test_pool_create(self):
        fake_pool_name = "swimming"
        desc = self._make_a_pool(fake_pool_name)
        self.assertEqual(desc.name, fake_pool_name)

    def test_pool_destroy(self):
        self._make_a_pool()

        def destroy_wrap(identifier, result_list):
            try:
                dpool.destroy(identifier)
            except dpool.PoolError as pe:
                result_list.append(pe)

        destroy_result = []
        destroy_thread = threading.Thread(target=destroy_wrap, args=(self.pool_m_uid, destroy_result))
        destroy_thread.start()
        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHPoolDestroy)
        shep_reply_msg = dmsg.SHPoolDestroyResponse(
            tag=self.next_tag(), ref=shep_msg.tag, err=dmsg.SHPoolDestroyResponse.Errors.SUCCESS
        )
        self.gs_input_wh.send(shep_reply_msg.serialize())
        destroy_thread.join()

        self.assertEqual(destroy_result, [])

    def test_pool_query(self):
        fake_pool_name = "swimming"
        desc = self._make_a_pool(fake_pool_name)

        desc2 = dpool.query(fake_pool_name)
        self.assertEqual(desc.m_uid, desc2.m_uid)

    def test_pool_list(self):
        fake_pool_name = "swimming"
        desc = self._make_a_pool(fake_pool_name)

        plist = dpool.get_list()

        self.assertEqual([desc.m_uid], plist)

    def test_channel_create(self):
        pool_desc = self._make_a_pool()

        the_channel_name = "my_channel"
        channel_desc = self._make_a_channel(pool_desc.m_uid, the_channel_name)
        self.assertEqual(channel_desc.name, the_channel_name)

    def test_channel_destroy(self):
        pool_desc = self._make_a_pool()

        the_channel_name = "my_channel"
        channel_desc = self._make_a_channel(pool_desc.m_uid, the_channel_name)
        self.assertEqual(channel_desc.name, the_channel_name)

        def destroy_wrap(identifier, result_list):
            try:
                dchannel.destroy(identifier)
            except dpool.PoolError as pe:
                result_list.append(pe)

        destroy_result = []
        destroy_thread = threading.Thread(target=destroy_wrap, args=(the_channel_name, destroy_result))
        destroy_thread.start()
        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHChannelDestroy)
        shep_reply_msg = dmsg.SHChannelDestroyResponse(
            tag=178343, ref=shep_msg.tag, err=dmsg.SHChannelDestroyResponse.Errors.SUCCESS
        )
        self.gs_input_wh.send(shep_reply_msg.serialize())
        destroy_thread.join()

        self.assertEqual(destroy_result, [])

    def test_channel_query(self):
        pool_desc = self._make_a_pool()

        the_channel_name = "my_channel"
        channel_desc = self._make_a_channel(pool_desc.m_uid, the_channel_name)
        self.assertEqual(channel_desc.name, the_channel_name)

        channel_desc2 = dchannel.query(the_channel_name)
        self.assertEqual(channel_desc.c_uid, channel_desc2.c_uid)

    def test_channel_list(self):
        pool_desc = self._make_a_pool()

        the_channel_name = "my_channel"
        channel_desc = self._make_a_channel(pool_desc.m_uid, the_channel_name)
        self.assertEqual(channel_desc.name, the_channel_name)

        clist = dchannel.get_list()
        self.assertEqual([channel_desc.c_uid], clist)

    def test_channel_join(self):
        pool_desc = self._make_a_pool()

        the_channel_name = "my_channel"

        def join_wrap(identifier, result_list):
            res = dchannel.join(identifier)
            result_list.append(res)

        join_result = []
        join_thread = threading.Thread(target=join_wrap, args=(the_channel_name, join_result))
        join_thread.start()

        channel_desc = self._make_a_channel(pool_desc.m_uid, the_channel_name)
        join_thread.join()
        self.assertEqual(join_result[0].name, channel_desc.name)

    def test_channel_join_timeout(self):
        the_channel_name = "intentionally_futile"

        def join_wrap(identifier, result_list):
            start = time.monotonic()
            try:
                _ = dchannel.join(identifier, 1.0)
            except TimeoutError:
                pass

            result_list.append(time.monotonic() - start)

        join_result = []
        join_thread = threading.Thread(target=join_wrap, args=(the_channel_name, join_result))
        join_thread.start()
        join_thread.join()
        self.assertGreater(join_result[0], 1.0)

    def test_multi_join(self):
        """Simple multi join test with infinite timeout.

        Join two processes, from which only one exits. The join returns and the
        result includes only the process exited.
        """

        proc_name = "first"
        desc_first = self._create_proc(proc_name)
        self.assertEqual(desc_first.name, proc_name)
        test_exit_code = 17
        proc_name = "second"
        desc_sec = self._create_proc(proc_name)
        self.assertEqual(desc_sec.name, proc_name)

        def join_wrap(identifiers, result_list):
            res = dproc.multi_join(identifiers)
            result_list += res[0]

        the_one_exiting = desc_first.p_uid
        join_result = []
        join_thread = threading.Thread(
            target=join_wrap, args=([desc_first.p_uid, desc_sec.p_uid], join_result)
        )
        join_thread.start()

        sh_kill_reply = dmsg.SHProcessExit(
            tag=self.next_tag(), p_uid=the_one_exiting, exit_code=test_exit_code
        )
        self.gs_input_wh.send(sh_kill_reply.serialize())

        join_thread.join()

        # the join result should only include the process that exited
        self.assertEqual(join_result[0][1], test_exit_code)
        self.assertEqual(join_result[0][0], the_one_exiting)
        self.assertEqual(len(join_result), 1)

    def test_multi_join_alt(self):
        """Multi join with infinite timeout.

        Join in multiple processes. Each process in the list waits for some time and
        then exits. We join on the same list of processes repeatedly and the join result
        includes only those processes exited. The join returns as soon as there is at least
        one process in the list that has exited.
        """

        identifiers = []
        nprocs = 5
        test_exit_code = 17

        def join_wrap(identifiers, result_list):
            while True:
                res = dproc.multi_join(identifiers)
                if res[0] is not None:
                    if len(res[0]) == nprocs:
                        result_list += res[0]
                        break
                # Wait before joining again
                time.sleep(0.1)

        def exit_wrap(tag, p_uid, exit_code, handler, sleep_time):
            time.sleep(sleep_time)
            sh_kill_reply = dmsg.SHProcessExit(tag=tag, p_uid=p_uid, exit_code=exit_code)
            try:
                handler.send(sh_kill_reply.serialize())
            except ConnectionError as ex:
                # If it had this error, it is because things are shutting down. Ignore it.
                pass

        for i in range(nprocs):
            proc_name = "proc" + str(i)
            desc = self._create_proc(proc_name)
            self.assertEqual(desc.name, proc_name)
            identifiers.append(desc.p_uid)

            exit_thread = threading.Thread(
                target=exit_wrap, args=(self.next_tag(), desc.p_uid, test_exit_code, self.gs_input_wh, i + 1)
            )
            exit_thread.start()

        join_result = []
        join_thread = threading.Thread(target=join_wrap, args=(identifiers, join_result))
        join_thread.start()

        join_thread.join()
        self.assertEqual(len(join_result), nprocs)

    def test_multi_join_timeout(self):
        """Multi join with timeout.

        Join when using user_name and p_uid for the target processes.
        """

        proc_name = "first"
        desc_first = self._create_proc(proc_name)
        self.assertEqual(desc_first.name, proc_name)
        proc_name = "second"
        desc_sec = self._create_proc(proc_name)
        self.assertEqual(desc_sec.name, proc_name)

        start = time.monotonic()
        # join using both process name and p_uid
        result = dproc.multi_join(["first", desc_sec.p_uid], timeout=1.0)

        duration = time.monotonic() - start
        # should time out in about one second
        self.assertGreater(duration, 1.0)

        for entry in result[1].values():
            self.assertEqual(entry, [2, None])
        self.assertEqual(result[0], None)

    def test_multi_join_with_names_timeout(self):
        """Multi join with timeout.

        Join when using user_names for the target processes.
        """

        proc_name = "first"
        desc_first = self._create_proc(proc_name)
        self.assertEqual(desc_first.name, proc_name)
        proc_name = "second"
        desc_sec = self._create_proc(proc_name)
        self.assertEqual(desc_sec.name, proc_name)

        start = time.monotonic()
        # join using process names instead of p_uid
        result = dproc.multi_join(["first", "second"], timeout=1.0)

        duration = time.monotonic() - start
        # should time out in about one second
        self.assertGreater(duration, 1.0)

        for entry in result[1].values():
            self.assertEqual(entry, [2, None])
        self.assertEqual(result[0], None)

    def test_multi_join_alt_all(self):
        """Multi join with option 'all' active and infinite timeout.

        Join in multiple processes. Each process in the list waits for some time and
        then exits. We join on this list of processes. The join call returns as soon
        as every process in the list has exited.
        """

        identifiers = []
        nprocs = 5
        test_exit_code = 17

        def join_wrap(identifiers, result_list):
            res = dproc.multi_join(identifiers, join_all=True)
            if res[0] is not None:
                result_list += res[0]

        def exit_wrap(tag, p_uid, exit_code, handler, sleep_time):
            time.sleep(sleep_time)
            sh_kill_reply = dmsg.SHProcessExit(tag=tag, p_uid=p_uid, exit_code=exit_code)
            try:
                handler.send(sh_kill_reply.serialize())
            except ConnectionError as ex:
                # If it had this error, it is because things are shutting down. Ignore it.
                pass

        for i in range(nprocs):
            proc_name = "proc" + str(i)
            desc = self._create_proc(proc_name)
            self.assertEqual(desc.name, proc_name)
            identifiers.append(desc.p_uid)
            exit_thread = threading.Thread(
                target=exit_wrap, args=(self.next_tag(), desc.p_uid, test_exit_code, self.gs_input_wh, i + 1)
            )
            if i == 0:
                start = time.monotonic()
            exit_thread.start()

        join_result = []
        join_thread = threading.Thread(target=join_wrap, args=(identifiers, join_result))
        join_thread.start()
        join_thread.join()
        end = time.monotonic()

        self.assertEqual(len(join_result), nprocs)
        self.assertGreater(end - start, 5.0)

    def test_multi_join_timeout_all(self):
        """Multi join with 'all' and timeout.

        Join when using user_name and p_uid for the target processes.
        """

        proc_name = "first"
        desc_first = self._create_proc(proc_name)
        self.assertEqual(desc_first.name, proc_name)
        proc_name = "second"
        desc_sec = self._create_proc(proc_name)
        self.assertEqual(desc_sec.name, proc_name)

        start = time.monotonic()
        # join using both process name and p_uid
        result = dproc.multi_join(["first", desc_sec.p_uid], timeout=1.0, join_all=True)

        duration = time.monotonic() - start
        # should time out in about one second
        self.assertGreater(duration, 1.0)

        for entry in result[1].values():
            self.assertEqual(entry, [2, None])
        self.assertEqual(result[0], None)

    def test_multi_join_any_all(self):
        """Multi join without timeout.

        Create two different wait requests on the same puid set.
        One comes with 'any' option and the other with 'all' option.
        The multi_join with 'any' option returns when one process from
        the set exits. The multi_join with 'all' option returns when all
        the processes from the set exit.
        """

        identifiers = []
        nprocs = 5
        test_exit_code = 17

        def join_wrap(identifiers, result_list, join_all):
            res = dproc.multi_join(identifiers, join_all=join_all)
            if res[0] is not None:
                result_list += res[0]

        def exit_wrap(tag, p_uid, exit_code, handler, sleep_time):
            time.sleep(sleep_time)
            sh_kill_reply = dmsg.SHProcessExit(tag=tag, p_uid=p_uid, exit_code=exit_code)
            try:
                handler.send(sh_kill_reply.serialize())
            except ConnectionError as ex:
                # If it had this error, it is because things are shutting down. Ignore it.
                pass

        for i in range(nprocs):
            proc_name = "proc" + str(i)
            desc = self._create_proc(proc_name)
            self.assertEqual(desc.name, proc_name)
            identifiers.append(desc.p_uid)
            exit_thread = threading.Thread(
                target=exit_wrap, args=(self.next_tag(), desc.p_uid, test_exit_code, self.gs_input_wh, i + 1)
            )
            if i == 0:
                start_any = time.monotonic()  # this is for the 'any' case
            if i == 4:
                start_all = time.monotonic()  # this is for timing the 'all' case
            exit_thread.start()

        join_result = []
        join_result_all = []
        join_thread = threading.Thread(target=join_wrap, args=(identifiers, join_result, False))
        join_thread.start()

        join_thread_all = threading.Thread(target=join_wrap, args=(identifiers, join_result_all, True))
        join_thread_all.start()

        join_thread.join()
        end_any = time.monotonic()
        join_thread_all.join()
        end_all = time.monotonic()

        self.assertGreater(end_all - start_all, 5.0)
        self.assertGreater(end_any - start_any, 1.0)
        self.assertEqual(len(join_result), 1)
        self.assertEqual(len(join_result_all), nprocs)

    def test_multi_join_timeout_all_with_error(self):
        """Multi join with 'all' and without timeout.

        Ask to join on a set of puids and one of them does not exist.
        The rest exit normally. The result from multi_join should be None.
        """

        identifiers = []
        nprocs = 5
        test_exit_code = 17

        def join_wrap(identifiers, result_list, join_all):
            res = dproc.multi_join(identifiers, join_all=join_all)
            if res[0] is not None:
                result_list += res[0]
            else:
                result_list.append(res[0])  # for the errored case

        def exit_wrap(tag, p_uid, exit_code, handler, sleep_time):
            time.sleep(sleep_time)
            sh_kill_reply = dmsg.SHProcessExit(tag=tag, p_uid=p_uid, exit_code=exit_code)
            try:
                handler.send(sh_kill_reply.serialize())
            except ConnectionError as ex:
                # If it had this error, it is because things are shutting down. Ignore it.
                pass

        for i in range(nprocs):
            proc_name = "proc" + str(i)
            desc = self._create_proc(proc_name)
            self.assertEqual(desc.name, proc_name)
            identifiers.append(desc.p_uid)
            exit_thread = threading.Thread(
                target=exit_wrap, args=(self.next_tag(), desc.p_uid, test_exit_code, self.gs_input_wh, i + 1)
            )
            if i == 0:
                start = time.monotonic()
            exit_thread.start()

        join_result = []
        identifiers.append("ghost")
        join_thread = threading.Thread(target=join_wrap, args=(identifiers, join_result, True))
        join_thread.start()
        join_thread.join()
        end = time.monotonic()

        self.assertEqual(join_result, [None])
        self.assertGreater(end - start, 5.0)

    def test_multi_join_timeout_all_diff_sets(self):
        """Multi join with 'all' and without timeout.

        Create two different requests with 'all' option that wait
        on two sets. The one set is two items less than the
        other set.
        """

        identifiers1 = []
        identifiers2 = []
        nprocs = 6
        test_exit_code = 17

        def join_wrap(identifiers, result_list, join_all):
            res = dproc.multi_join(identifiers, join_all=join_all)
            if res[0] is not None:
                result_list += res[0]

        def exit_wrap(tag, p_uid, exit_code, handler, sleep_time):
            time.sleep(sleep_time)
            sh_kill_reply = dmsg.SHProcessExit(tag=tag, p_uid=p_uid, exit_code=exit_code)
            try:
                handler.send(sh_kill_reply.serialize())
            except ConnectionError as ex:
                # If it had this error, it is because things are shutting down. Ignore it.
                pass

        for i in range(nprocs):
            proc_name = "proc" + str(i)
            desc = self._create_proc(proc_name)
            self.assertEqual(desc.name, proc_name)
            identifiers1.append(desc.p_uid)
            identifiers2.append(desc.p_uid)
            exit_thread = threading.Thread(
                target=exit_wrap, args=(self.next_tag(), desc.p_uid, test_exit_code, self.gs_input_wh, i + 1)
            )
            if i == 0:
                start1 = time.monotonic()  # this is for the full set
            if i == 1:
                start2 = time.monotonic()  # this is for timing the set including fewer items
            exit_thread.start()

        join_result1 = []
        join_result2 = []

        join_thread1 = threading.Thread(target=join_wrap, args=(identifiers1, join_result1, True))
        join_thread1.start()

        # remove two items from the list of identifiers for the 2nd request to wait on
        identifiers2.pop(0)
        identifiers2.pop(-1)
        join_thread2 = threading.Thread(target=join_wrap, args=(identifiers2, join_result2, True))
        join_thread2.start()

        join_thread1.join()
        end1 = time.monotonic()
        join_thread2.join()
        end2 = time.monotonic()

        self.assertEqual(len(join_result1), nprocs)
        self.assertEqual(len(join_result2), nprocs - 2)
        self.assertGreater(end1 - start1, 6.0)
        self.assertGreater(end2 - start2, 5.0)

    def test_multi_join_all_test_new_timeout(self):
        """Multi join with option 'all' option and timeout.

        Join in multiple processes. Each process in the list waits for some time and
        then exits. We join on this list of processes. The join call returns as soon
        as every process in the list has exited, since the timeout used is larger
        than the sleep time of the processes. In this test we exercise the new entries
        to pending_join_list DS with the new timeout (keep shrinking the puid_set that
        we wait on, as a result to a process exiting).
        """

        identifiers = []
        nprocs = 3
        test_exit_code = 17

        def join_wrap(identifiers, result_list, timeout):
            res = dproc.multi_join(identifiers, timeout=timeout, join_all=True)
            if res[0] is not None:
                result_list += res[0]

        def exit_wrap(tag, p_uid, exit_code, handler, sleep_time):
            time.sleep(sleep_time)
            sh_kill_reply = dmsg.SHProcessExit(tag=tag, p_uid=p_uid, exit_code=exit_code)
            try:
                handler.send(sh_kill_reply.serialize())
            except ConnectionError as ex:
                # If it had this error, it is because things are shutting down. Ignore it.
                pass

        for i in range(nprocs):
            proc_name = "proc" + str(i)
            desc = self._create_proc(proc_name)
            self.assertEqual(desc.name, proc_name)
            identifiers.append(desc.p_uid)
            exit_thread = threading.Thread(
                target=exit_wrap, args=(self.next_tag(), desc.p_uid, test_exit_code, self.gs_input_wh, i + 1)
            )
            if i == 0:
                start = time.monotonic()
            exit_thread.start()

        join_result = []
        join_thread = threading.Thread(target=join_wrap, args=(identifiers, join_result, 4.0))
        join_thread.start()
        join_thread.join()
        end = time.monotonic()

        self.assertEqual(len(join_result), nprocs)
        self.assertGreater(end - start, 3.0)

    def test_multi_join_mult_req_on_same_key(self):
        """Multi join on the same puid set with timeout.

        Create two different wait requests on the same puid set with
        timeout. The first request times out and then we make the
        second one.
        """

        identifiers = []
        nprocs = 3

        for i in range(nprocs):
            proc_name = "proc" + str(i)
            desc = self._create_proc(proc_name)
            self.assertEqual(desc.name, proc_name)
            identifiers.append(desc.p_uid)

        # it doesn't really matter if it is 'any' or 'all' requests or any combination of those
        res1 = dproc.multi_join(identifiers, timeout=0.1, join_all=False)
        # want the first request to return/time out by the time we create the next request
        res2 = dproc.multi_join(identifiers, timeout=0.1, join_all=True)

        self.assertEqual(res1[0], None)
        self.assertEqual(res2[0], None)


if __name__ == "__main__":
    unittest.main()

test/globalservices/process_api.py
DragonHPC/dragon
#!/usr/bin/env python3
"""Single node gs startup/teardown smoke test"""

import logging
import multiprocessing
import threading
import time
import unittest

from dragon.infrastructure.node_desc import NodeDescriptor

import dragon.globalservices.server as dserver
import dragon.dlogging.util as dlog
import dragon.infrastructure.messages as dmsg
import support.util as tsu


def startup(gs_input, shep_input, bela_input, gs_stdout, logname=""):
    dlog.setup_logging(basename="gs_" + logname, level=logging.DEBUG, force=True)
    log = logging.getLogger("single startup entry")
    log.info("starting")
    the_ctx = dserver.GlobalContext()

    try:
        the_ctx.run_startup(
            mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE,
            test_gs_input=gs_input,
            test_shep_inputs=[shep_input],
            test_bela_input=bela_input,
            test_gs_stdout=gs_stdout
        )

        the_ctx.run_global_server(
            mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE,
            test_gs_input=gs_input,
            test_shep_inputs=[shep_input],
            test_bela_input=bela_input,
            test_gs_stdout=gs_stdout
        )

        log.info("normal exit")
    except:
        log.exception("fatal startup error")
        gs_stdout.send(dmsg.AbnormalTermination(tag=0).serialize())

    logging.shutdown()


def main_loop_exit(gs_input, shep_input, bela_input, gs_stdout, logname=""):
    dlog.setup_logging(basename="gs_" + logname, level=logging.DEBUG, force=True)
    log = logging.getLogger("main loop exit test")
    log.info("starting")

    the_ctx = dserver.GlobalContext()
    the_ctx._launch_mode = the_ctx.LaunchModes.TEST_STANDALONE_SINGLE

    try:
        the_ctx.run_global_server(
            mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE,
            test_gs_input=gs_input,
            test_shep_inputs=[shep_input],
            test_bela_input=bela_input,
            test_gs_stdout=gs_stdout,
        )
        log.info("normal exit")
    except:
        log.exception("fatal main loop error")
        gs_stdout.send(dmsg.AbnormalTermination(tag=0).serialize())

    logging.shutdown()


def updown(gs_input, shep_input, bela_input, gs_stdout, logname=""):
    dlog.setup_logging(basename="gs_" + logname, level=logging.DEBUG, force=True)
    log = logging.getLogger("up down test")
    log.info("starting")

    the_ctx = dserver.GlobalContext()

    try:
        the_ctx.run_startup(
            mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE,
            test_gs_input=gs_input,
            test_shep_inputs=[shep_input],
            test_bela_input=bela_input,
            test_gs_stdout=gs_stdout,
        )

        the_ctx.run_global_server(
            mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE,
            test_gs_input=gs_input,
            test_shep_inputs=[shep_input],
            test_bela_input=bela_input,
            test_gs_stdout=gs_stdout,
        )

        the_ctx.run_teardown(
            test_gs_input=gs_input,
            test_shep_inputs=[shep_input],
            test_bela_input=bela_input,
            test_gs_stdout=gs_stdout,
        )
    except:
        log.exception("fatal error")
        gs_stdout.send(dmsg.AbnormalTermination(tag=0).serialize())

    log.info("normal exit")
    logging.shutdown()


class SingleInternal(unittest.TestCase):
    def setUp(self) -> None:
        self.gs_input_rh, self.gs_input_wh = multiprocessing.Pipe(duplex=False)
        self.gs_stdout_rh, self.gs_stdout_wh = multiprocessing.Pipe(duplex=False)
        self.shep_input_rh, self.shep_input_wh = multiprocessing.Pipe(duplex=False)
        self.bela_input_rh, self.bela_input_wh = multiprocessing.Pipe(duplex=False)
        self.node_sdesc = NodeDescriptor.make_for_current_node(is_primary=True).sdesc
        self.dut = None
        self.tag = 0

    def tearDown(self) -> None:
        self.gs_input_rh.close()
        self.gs_input_wh.close()
        self.gs_stdout_rh.close()
        self.gs_stdout_wh.close()
        self.bela_input_rh.close()
        self.bela_input_wh.close()
        self.shep_input_rh.close()
        self.shep_input_wh.close()
        self.dut.join(1)

    def next_tag(self):
        tmp = self.tag
        self.tag += 1
        return tmp


    def test_startup_lifecycle_err(self):
        self.dut = threading.Thread(
            target=startup,
            args=(self.gs_input_rh, self.shep_input_wh, self.bela_input_wh, self.gs_stdout_wh),
            kwargs={"logname": self.id()},
            daemon=True,
            name="globalservices",
        )

        self.dut.start()
        tsu.get_and_check_type(self.shep_input_rh, dmsg.GSPingSH)
        self.gs_input_wh.send("crap")  # triggers error
        tsu.get_and_check_type(self.gs_stdout_rh, dmsg.AbnormalTermination)


    def test_main_loop_exit(self):
        self.dut = threading.Thread(
            target=main_loop_exit,
            args=(self.gs_input_rh, self.shep_input_wh, self.bela_input_wh, self.gs_stdout_wh),
            kwargs={"logname": self.id()},
            daemon=True,
            name="globalservices",
        )
        self.dut.start()
        self.gs_input_wh.send(dmsg.GSTeardown(tag=0).serialize())
        time.sleep(1)
        self.assertFalse(self.dut.is_alive())

    def test_main_loop_bad_exit(self):
        self.dut = threading.Thread(
            target=main_loop_exit,
            args=(self.gs_input_rh, self.shep_input_wh, self.bela_input_wh, self.gs_stdout_wh),
            kwargs={"logname": self.id()},
            daemon=True,
            name="globalservices",
        )
        self.dut.start()
        self.gs_input_wh.send("crap")
        tsu.get_and_check_type(self.gs_stdout_rh, dmsg.AbnormalTermination)

    def test_updown(self):
        self.dut = threading.Thread(
            target=updown,
            args=(self.gs_input_rh, self.shep_input_wh, self.bela_input_wh, self.gs_stdout_wh),
            kwargs={"logname": self.id()},
            daemon=True,
            name="globalservices",
        )
        self.dut.start()
        tsu.get_and_check_type(self.shep_input_rh, dmsg.GSPingSH)
        self.gs_input_wh.send(dmsg.SHPingGS(tag=self.next_tag(), node_sdesc=self.node_sdesc).serialize())
        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSIsUp)
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self.next_tag()).serialize())
        tsu.get_and_check_type(self.gs_stdout_rh, dmsg.GSHalted)


if __name__ == "__main__":
    unittest.main()

test/globalservices/single_internal.py
DragonHPC/dragon
#!/usr/bin/env python3

"""Test script for single node gs process messages"""
import signal

import copy
import inspect
import logging
import multiprocessing
import os
import unittest

from dragon import channels as dch
from dragon import managed_memory as dmm

from dragon.globalservices import server as dserver

from dragon.infrastructure import process_desc
from dragon.infrastructure import facts as dfacts
from dragon.dlogging import util as dlog
from dragon.infrastructure import messages as dmsg
from dragon.infrastructure import connection as dconn
from dragon.infrastructure import parameters as dparm
import support.util as tsu
from dragon.utils import B64

from dragon.infrastructure.node_desc import NodeDescriptor


def bringup_channels(gs_stdout, env_updates, logname=""):
    dlog.setup_logging(basename="gs_" + logname, level=logging.DEBUG, force=True)
    log = logging.getLogger("channels bringup test")
    log.info("start")

    # patch up the environment
    dparm.this_process = dparm.LaunchParameters.from_env(env_updates)

    the_ctx = dserver.GlobalContext()

    try:
        the_ctx.run_startup(mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE, test_gs_stdout=gs_stdout)

        the_ctx.run_global_server(mode=the_ctx.LaunchModes.TEST_STANDALONE_SINGLE, test_gs_stdout=gs_stdout)

        the_ctx.run_teardown(test_gs_stdout=gs_stdout)

    except:
        log.exception("fatal")
        gs_stdout.send(dmsg.AbnormalTermination(tag=0).serialize())

    log.info("normal exit")
    logging.shutdown()


class SingleProcMsgChannels(unittest.TestCase):
    def setUp(self) -> None:
        self.gs_stdout_rh, self.gs_stdout_wh = multiprocessing.Pipe(duplex=False)
        self.some_parms = copy.copy(dparm.this_process)

        username = os.environ.get("USER", str(os.getuid()))
        self.pool_name = "mctest_" + username
        self.pool_size = 2**30
        self.pool_prealloc_blocks = None
        self.pool_uid = 17
        self.mpool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid, self.pool_prealloc_blocks)
        self.mpool_ser = self.mpool.serialize()
        self.some_parms.inf_pd = B64.bytes_to_str(self.mpool_ser)

        self.gs_input_chan = dch.Channel(self.mpool, dfacts.GS_INPUT_CUID)
        self.gs_input_rh = dconn.Connection(inbound_initializer=self.gs_input_chan)
        self.gs_input_wh = dconn.Connection(outbound_initializer=self.gs_input_chan)
        self.some_parms.gs_cd = B64.bytes_to_str(self.gs_input_chan.serialize())

        self.bela_input_chan = dch.Channel(self.mpool, dfacts.BASE_BE_CUID)
        self.bela_input_rh = dconn.Connection(inbound_initializer=self.bela_input_chan)
        self.bela_input_wh = dconn.Connection(outbound_initializer=self.bela_input_chan)
        self.some_parms.local_be_cd = B64.bytes_to_str(self.bela_input_chan.serialize())

        self.shep_input_chan = dch.Channel(self.mpool, dfacts.BASE_SHEP_CUID)
        self.shep_input_rh = dconn.Connection(inbound_initializer=self.shep_input_chan)
        self.shep_input_wh = dconn.Connection(outbound_initializer=self.shep_input_chan)
        self.some_parms.local_shep_cd = B64.bytes_to_str(self.shep_input_chan.serialize())

        self.node_sdesc = NodeDescriptor.make_for_current_node(is_primary=True).sdesc

        self.dut = None
        self.head_uid = None
        self.tag_cnt = 4

    def tearDown(self) -> None:
        # wrapped queue ends
        self.gs_stdout_rh.close()
        self.gs_stdout_wh.close()

        # connections
        self.gs_input_rh.close()
        self.gs_input_wh.close()

        self.bela_input_rh.close()
        self.bela_input_wh.close()

        self.shep_input_rh.close()
        self.shep_input_wh.close()

        # underlying channels
        self.gs_input_chan.destroy()
        self.bela_input_chan.destroy()
        self.shep_input_chan.destroy()
        self.mpool.destroy()

        if self.dut is not None:
            self.dut.join(timeout=1)

    def _start_dut(self, test_name=None):
        if test_name is None:
            test_name = self.__class__.__name__ + "_" + inspect.stack()[1][0].f_code.co_name

        the_env_wanted = self.some_parms.env()

        dut = multiprocessing.Process(
            target=bringup_channels,
            args=(self.gs_stdout_wh, the_env_wanted),
            kwargs={"logname": test_name},
            name="globalservices",
        )
        dut.start()

        tsu.get_and_check_type(self.shep_input_rh, dmsg.GSPingSH)

        self.gs_input_wh.send(dmsg.SHPingGS(tag=self._tag_inc(), node_sdesc=self.node_sdesc).serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSIsUp)

        return dut

    def _tag_inc(self):
        self.tag_cnt += 1
        return self.tag_cnt

    def _start_a_process(self, exe_name, the_args, the_name, the_tag, the_puid, the_rcuid, pmi_required=False, pmi_info=None):
        if pmi_required:
            self.assertIsNotNone(pmi_info)

        create_msg = dmsg.GSProcessCreate(
            tag=the_tag,
            p_uid=the_puid,
            r_c_uid=the_rcuid,
            exe=exe_name,
            args=the_args,
            env={},
            user_name=the_name,
            pmi_required=pmi_required,
            _pmi_info=pmi_info
        )

        self.gs_input_wh.send(create_msg.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
        self.assertEqual(shep_msg.exe, create_msg.exe)
        self.assertEqual(shep_msg.args, create_msg.args)

        if pmi_required:
            self.assertIsNotNone(shep_msg.pmi_info)
            self.assertEqual(shep_msg.pmi_info.job_id, pmi_info.job_id)
            self.assertEqual(shep_msg.pmi_info.lrank, pmi_info.lrank)
            self.assertEqual(shep_msg.pmi_info.ppn, pmi_info.ppn)
            self.assertEqual(shep_msg.pmi_info.nid, pmi_info.nid)
            self.assertEqual(shep_msg.pmi_info.nnodes, pmi_info.nnodes)
            self.assertEqual(shep_msg.pmi_info.nranks, pmi_info.nranks)
            self.assertEqual(shep_msg.pmi_info.nidlist, pmi_info.nidlist)
            self.assertEqual(shep_msg.pmi_info.hostlist, pmi_info.hostlist)
            self.assertEqual(shep_msg.pmi_info.control_port, pmi_info.control_port)

        shep_reply_msg = dmsg.SHProcessCreateResponse(
            tag=self._tag_inc(), ref=shep_msg.tag, err=dmsg.SHProcessCreateResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_reply_msg = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessCreateResponse)
        self.assertEqual(create_reply_msg.ref, create_msg.tag, "tag-ref mismatch")
        self.assertEqual(create_reply_msg.err, dmsg.GSProcessCreateResponse.Errors.SUCCESS)
        self.assertEqual(shep_msg.t_p_uid, create_reply_msg.desc.p_uid)

        return create_msg, create_reply_msg

    def _teardown_dut(self):
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self._tag_inc()).serialize())

        tsu.get_and_check_type(self.gs_stdout_rh, dmsg.GSHalted)

    def _bringup_head_and_dut(self):
        test_name = self.__class__.__name__ + "_" + inspect.stack()[1][0].f_code.co_name

        self._start_dut(test_name)
        test_exe = "head"
        test_args = ["foo"]
        test_name = "first one"
        create_msg, create_reply_msg = self._start_a_process(
            exe_name=test_exe,
            the_args=test_args,
            the_name=test_name,
            the_tag=self._tag_inc(),
            the_puid=dfacts.LAUNCHER_PUID,
            the_rcuid=dfacts.BASE_BE_CUID,
        )

        self.assertEqual(
            create_reply_msg.desc.state, process_desc.ProcessDescriptor.State.ACTIVE
        )
        self.assertEqual(create_msg.user_name, create_reply_msg.desc.name)
        self.head_uid = create_reply_msg.desc.p_uid
        self.head_name = test_name

    def _teardown_head_and_dut(self):
        death_msg = dmsg.SHProcessExit(tag=self._tag_inc(), p_uid=self.head_uid)

        self.gs_input_wh.send(death_msg.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)
        self._teardown_dut()

    def _make_head_fail(self):
        test_exe = "dummy"
        test_args = ["foo", "bar"]
        test_name = "first one"

        create_msg = dmsg.GSProcessCreate(
            tag=self._tag_inc(),
            p_uid=dfacts.LAUNCHER_PUID,
            r_c_uid=dfacts.BASE_BE_CUID,
            exe=test_exe,
            args=test_args,
            env={},
            user_name=test_name,
        )

        self.gs_input_wh.send(create_msg.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessCreate)
        self.assertEqual(shep_msg.exe, create_msg.exe)
        self.assertEqual(shep_msg.args, create_msg.args)

        shep_reply_msg = dmsg.SHProcessCreateResponse(
            tag=self._tag_inc(),
            ref=shep_msg.tag,
            err=dmsg.SHProcessCreateResponse.Errors.FAIL,
            err_info="whoa",
        )

        self.gs_input_wh.send(shep_reply_msg.serialize())

        create_reply_msg = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessCreateResponse)
        self.assertEqual(create_reply_msg.ref, create_msg.tag, "tag-ref mismatch")
        self.assertEqual(create_reply_msg.err, dmsg.GSProcessCreateResponse.Errors.FAIL)

    def _make_a_pool(self):
        the_tag = self._tag_inc()

        pcm = dmsg.GSPoolCreate(
            tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID, size=2**30
        )
        self.gs_input_wh.send(pcm.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHPoolCreate)

        dummy_sdesc = B64.bytes_to_str("xxxTESTDUMMYxxx".encode())
        shep_response = dmsg.SHPoolCreateResponse(
            tag=self._tag_inc(),
            ref=shep_msg.tag,
            err=dmsg.SHPoolCreateResponse.Errors.SUCCESS,
            desc=dummy_sdesc,
        )
        self.gs_input_wh.send(shep_response.serialize())

        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSPoolCreateResponse)

        self.assertEqual(gs_response.err, dmsg.GSPoolCreateResponse.Errors.SUCCESS)
        self.assertEqual(gs_response.ref, the_tag)
        return gs_response.desc.m_uid

    def _destroy_a_pool(self, the_muid):
        the_tag = self._tag_inc()
        pdm = dmsg.GSPoolDestroy(
            tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID, m_uid=the_muid
        )
        self.gs_input_wh.send(pdm.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHPoolDestroy)

        shep_response = dmsg.SHPoolDestroyResponse(
            tag=self._tag_inc(), ref=shep_msg.tag, err=dmsg.SHPoolDestroyResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(shep_response.serialize())

        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSPoolDestroyResponse)
        self.assertEqual(gs_response.err, dmsg.GSPoolDestroyResponse.Errors.SUCCESS)
        self.assertEqual(gs_response.ref, the_tag)

    def _make_a_channel(self, the_muid):
        the_tag = self._tag_inc()

        pcm = dmsg.GSChannelCreate(
            tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, m_uid=the_muid, r_c_uid=dfacts.BASE_BE_CUID
        )

        self.gs_input_wh.send(pcm.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHChannelCreate)

        dummy_sdesc = B64.bytes_to_str("xxxTESTDUMMYxxx".encode())
        shep_response = dmsg.SHChannelCreateResponse(
            tag=self._tag_inc(),
            ref=shep_msg.tag,
            err=dmsg.SHChannelCreateResponse.Errors.SUCCESS,
            desc=dummy_sdesc,
        )
        self.gs_input_wh.send(shep_response.serialize())

        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSChannelCreateResponse)

        self.assertEqual(gs_response.err, dmsg.GSChannelCreateResponse.Errors.SUCCESS)
        self.assertEqual(gs_response.ref, the_tag)

        return gs_response.desc.c_uid

    def _destroy_a_channel(self, the_cuid):
        the_tag = self._tag_inc()

        cdm = dmsg.GSChannelDestroy(
            tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID, c_uid=the_cuid
        )
        self.gs_input_wh.send(cdm.serialize())

        shep_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHChannelDestroy)

        shep_response = dmsg.SHChannelDestroyResponse(
            tag=self._tag_inc(), ref=shep_msg.tag, err=dmsg.SHChannelDestroyResponse.Errors.SUCCESS
        )
        self.gs_input_wh.send(shep_response.serialize())

        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSChannelDestroyResponse)
        self.assertEqual(gs_response.err, dmsg.GSChannelDestroyResponse.Errors.SUCCESS)
        self.assertEqual(gs_response.ref, the_tag)

    def test_proc_updown(self):
        self._bringup_head_and_dut()
        self._teardown_head_and_dut()

    def test_null_bringup(self):
        self._start_dut()
        self._teardown_dut()

    # regression for PE-33554
    def test_make_head_proc_fail_to_teardown(self):
        self._start_dut()
        self._make_head_fail()
        self._teardown_dut()

    def test_make_head_proc_list(self):
        self._bringup_head_and_dut()
        list_msg = dmsg.GSProcessList(
            tag=self._tag_inc(), p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID
        )
        self.gs_input_wh.send(list_msg.serialize())

        list_reply_msg = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessListResponse)
        self.assertEqual(self.head_uid, list_reply_msg.plist[0])

        self._teardown_head_and_dut()

    def test_make_head_proc_query(self):
        self._bringup_head_and_dut()

        query_msg = dmsg.GSProcessQuery(
            tag=self._tag_inc(),
            p_uid=dfacts.LAUNCHER_PUID,
            t_p_uid=self.head_uid,
            r_c_uid=dfacts.BASE_BE_CUID,
        )
        query_msg_byname = dmsg.GSProcessQuery(
            tag=self._tag_inc(),
            user_name=self.head_name,
            p_uid=dfacts.LAUNCHER_PUID,
            r_c_uid=dfacts.BASE_BE_CUID,
        )
        query_msg_unknown = dmsg.GSProcessQuery(
            tag=self._tag_inc(), user_name="absent", p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID
        )

        self.gs_input_wh.send(query_msg.serialize())
        query_reply_msg = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessQueryResponse)
        self.assertEqual(query_reply_msg.ref, query_msg.tag)
        self.assertEqual(query_reply_msg.err, dmsg.GSProcessQueryResponse.Errors.SUCCESS)
        self.assertEqual(query_reply_msg.desc.name, self.head_name)

        self.gs_input_wh.send(query_msg_byname.serialize())
        query_byname_reply_message = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessQueryResponse)
        self.assertEqual(query_byname_reply_message.ref, query_msg_byname.tag)
        self.assertEqual(query_byname_reply_message.err, dmsg.GSProcessQueryResponse.Errors.SUCCESS)
        self.assertEqual(query_byname_reply_message.desc.p_uid, self.head_uid)

        self.gs_input_wh.send(query_msg_unknown.serialize())
        query_unknown_reply = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessQueryResponse)
        self.assertEqual(query_unknown_reply.ref, query_msg_unknown.tag)
        self.assertEqual(query_unknown_reply.err, dmsg.GSProcessQueryResponse.Errors.UNKNOWN)

        self._teardown_head_and_dut()

    def test_up_down_with_proc(self):
        self._start_dut()

        test_exe = "dummy"
        test_args = ["foo", "bar"]
        test_name = "first one"

        create_msg, create_reply_msg = self._start_a_process(
            exe_name=test_exe,
            the_args=test_args,
            the_name=test_name,
            the_tag=self._tag_inc(),
            the_puid=dfacts.LAUNCHER_PUID,
            the_rcuid=dfacts.BASE_BE_CUID,
        )

        death_msg = dmsg.SHProcessExit(tag=self._tag_inc(), p_uid=create_reply_msg.desc.p_uid)

        self.gs_input_wh.send(death_msg.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)

        self._teardown_dut()

    def test_up_down_with_pmi_proc(self):
        self._start_dut()

        test_exe = "dummy"
        test_args = ["foo", "bar"]
        test_name = "first one"

        create_msg, create_reply_msg = self._start_a_process(
            exe_name=test_exe,
            the_args=test_args,
            the_name=test_name,
            the_tag=self._tag_inc(),
            the_puid=dfacts.LAUNCHER_PUID,
            the_rcuid=dfacts.BASE_BE_CUID,
            pmi_required=True,
            pmi_info=dmsg.PMIInfo(
                job_id=12323,
                lrank=0,
                ppn=1,
                nid=1,
                nnodes=2,
                nranks=2,
                nidlist=[0,1],
                hostlist=['head', 'worker'],
                control_port=1023,
                pid_base=1
            )
        )

        death_msg = dmsg.SHProcessExit(tag=self._tag_inc(), p_uid=create_reply_msg.desc.p_uid)

        self.gs_input_wh.send(death_msg.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)

        self._teardown_dut()

    def test_controlled_kill(self):
        self._start_dut()

        test_exe = "dummy"
        test_args = ["foo", "bar"]
        test_name = "first one"

        create_msg, create_reply_msg = self._start_a_process(
            exe_name=test_exe,
            the_args=test_args,
            the_name=test_name,
            the_tag=self._tag_inc(),
            the_puid=dfacts.LAUNCHER_PUID,
            the_rcuid=dfacts.BASE_BE_CUID,
        )

        kill_msg = dmsg.GSProcessKill(
            tag=self._tag_inc(),
            p_uid=dfacts.LAUNCHER_PUID,
            r_c_uid=dfacts.BASE_BE_CUID,
            t_p_uid=create_reply_msg.desc.p_uid,
            sig=signal.SIGKILL,
        )

        self.gs_input_wh.send(kill_msg.serialize())

        sh_kill_msg = tsu.get_and_check_type(self.shep_input_rh, dmsg.SHProcessKill)
        self.assertEqual(sh_kill_msg.t_p_uid, kill_msg.t_p_uid)

        sh_kill_reply = dmsg.SHProcessKillResponse(
            tag=17, ref=sh_kill_msg.tag, err=dmsg.SHProcessKillResponse.Errors.SUCCESS
        )

        self.gs_input_wh.send(sh_kill_reply.serialize())

        kill_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessKillResponse)
        self.assertEqual(kill_response.err, dmsg.GSProcessKillResponse.Errors.SUCCESS)

        exit_msg = dmsg.SHProcessExit(tag=self._tag_inc(), p_uid=sh_kill_msg.t_p_uid)

        self.gs_input_wh.send(exit_msg.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)

        self._teardown_dut()

    def test_up_down_twice_with_proc(self):
        self._start_dut()

        test_exe = "dummy"
        test_args = ["foo", "bar"]
        test_name = "first one"

        create_msg, create_reply_msg = self._start_a_process(
            exe_name=test_exe,
            the_args=test_args,
            the_name=test_name,
            the_tag=self._tag_inc(),
            the_puid=dfacts.LAUNCHER_PUID,
            the_rcuid=dfacts.BASE_BE_CUID,
        )

        death_msg = dmsg.SHProcessExit(tag=self._tag_inc(), p_uid=create_reply_msg.desc.p_uid)

        self.gs_input_wh.send(death_msg.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)

        test_args = ["foo2", "bar2"]
        test_name = "second one"

        create_msg2, create_reply_msg2 = self._start_a_process(
            exe_name=test_exe,
            the_args=test_args,
            the_name=test_name,
            the_tag=self._tag_inc(),
            the_puid=dfacts.LAUNCHER_PUID,
            the_rcuid=dfacts.BASE_BE_CUID,
        )

        death_msg2 = dmsg.SHProcessExit(tag=self._tag_inc(), p_uid=create_reply_msg2.desc.p_uid)

        self.gs_input_wh.send(death_msg2.serialize())

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)

        self._teardown_dut()

    def test_pool_create(self):
        self._bringup_head_and_dut()
        self._make_a_pool()
        self._teardown_head_and_dut()

    def test_pool_create_destroy(self):
        self._bringup_head_and_dut()

        the_muid = self._make_a_pool()
        self._destroy_a_pool(the_muid)

        self._teardown_head_and_dut()

    def test_channel_create(self):
        self._bringup_head_and_dut()

        the_muid = self._make_a_pool()
        self._make_a_channel(the_muid)
        self._teardown_head_and_dut()

    def test_channel_destroy(self):
        self._bringup_head_and_dut()

        the_muid = self._make_a_pool()
        the_cuid = self._make_a_channel(the_muid)
        self._destroy_a_channel(the_cuid)
        self._teardown_head_and_dut()

    def test_channel_pool_list(self):
        self._bringup_head_and_dut()

        the_muid = self._make_a_pool()
        the_cuid = self._make_a_channel(the_muid)

        the_tag = self._tag_inc()
        pool_list_msg = dmsg.GSPoolList(tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID)
        self.gs_input_wh.send(pool_list_msg.serialize())
        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSPoolListResponse)
        self.assertEqual(gs_response.err, dmsg.GSPoolListResponse.Errors.SUCCESS)
        self.assertEqual(gs_response.mlist, [the_muid])

        the_tag = self._tag_inc()
        channel_list_msg = dmsg.GSChannelList(
            tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID
        )
        self.gs_input_wh.send(channel_list_msg.serialize())
        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSChannelListResponse)
        self.assertEqual(gs_response.err, dmsg.GSChannelListResponse.Errors.SUCCESS)
        self.assertEqual(gs_response.clist, [the_cuid])

        self._teardown_head_and_dut()

    def test_pool_query(self):
        self._bringup_head_and_dut()

        the_muid = self._make_a_pool()

        the_tag = self._tag_inc()
        pool_query_msg = dmsg.GSPoolQuery(
            m_uid=the_muid, tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID
        )
        self.gs_input_wh.send(pool_query_msg.serialize())
        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSPoolQueryResponse)

        self.assertEqual(the_muid, gs_response.desc.m_uid)

        self._teardown_head_and_dut()

    def test_channel_query(self):
        self._bringup_head_and_dut()

        the_muid = self._make_a_pool()
        the_cuid = self._make_a_channel(the_muid)

        the_tag = self._tag_inc()
        channel_query_msg = dmsg.GSChannelQuery(
            c_uid=the_cuid, tag=the_tag, p_uid=dfacts.LAUNCHER_PUID, r_c_uid=dfacts.BASE_BE_CUID
        )
        self.gs_input_wh.send(channel_query_msg.serialize())
        gs_response = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSChannelQueryResponse)

        self.assertEqual(the_cuid, gs_response.desc.c_uid)
        self._teardown_head_and_dut()


if __name__ == "__main__":
    unittest.main()

test/globalservices/single_process_msg.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import random
import string
import json

from dragon.globalservices.policy_eval import Policy, PolicyEvaluator, ResourceLayout
from dragon.infrastructure.node_desc import NodeDescriptor

class TestPolicy(unittest.TestCase):

    def test_sdict(self):
        policy = Policy()
        sdict = policy.sdesc
        policy2 = Policy.from_sdict(sdict)
        sdict2 = policy2.sdesc
        for k in sdict:
            self.assertEqual(sdict[k], sdict2[k])

    def test_json_sdict(self):
        # Policy can be json-ified, then un-json'd, and values match
        policy = Policy()
        orig_sdict = policy.sdesc
        ser_sdict = json.dumps(Policy.get_sdict(policy))
        unser_sdict = json.loads(ser_sdict)
        for k in orig_sdict:
            self.assertEqual(orig_sdict[k], unser_sdict[k])


    def test_default(self):
        policy = Policy()
        self.assertEqual(policy.placement, Policy.Placement.DEFAULT)
        self.assertEqual(policy.host_name, '')
        self.assertEqual(policy.host_id, -1)
        self.assertEqual(policy.distribution, Policy.Distribution.DEFAULT)

    def test_host_name(self):
        # Test a random hostname and make sure the policy holds it
        name = ''.join(random.choices(string.ascii_letters + string.digits,
                                      k=random.randint(1, 236)))

        policy = Policy(host_name=name)
        self.assertEqual(policy.placement, Policy.Placement.DEFAULT)
        self.assertEqual(policy.host_name, name)
        self.assertEqual(policy.host_id, -1)
        self.assertEqual(policy.distribution, Policy.Distribution.DEFAULT)

    def test_host_id(self):
        policy = Policy(host_id=100)
        self.assertEqual(policy.placement, Policy.Placement.DEFAULT)
        self.assertEqual(policy.host_name, '')
        self.assertEqual(policy.host_id, 100)
        self.assertEqual(policy.distribution, Policy.Distribution.DEFAULT)

    def test_merge(self):
        high = Policy(placement=Policy.Placement.HOST_NAME, host_name='high_policy')
        low = Policy(affinity=Policy.Affinity.SPECIFIC, specific_affinity=[1,3,5])
        out = PolicyEvaluator.merge(high, low)
        self.assertEqual(high.placement, out.placement)
        self.assertEqual(high.host_name, out.host_name)
        self.assertEqual(low.affinity, out.affinity)
        self.assertEqual(low.specific_affinity, out.specific_affinity)


class TestEvaluator(unittest.TestCase):

    def setUp(self):
        self.node_cpus = 4
        self.nodes = []
        self.num_nodes = 4
        for i in range(1, self.num_nodes+1):
            node = NodeDescriptor()
            node.host_name = str(i)
            node.host_id = i
            node.h_uid = i
            node.num_cpus = self.node_cpus
            self.nodes.append(node)

        self.RR = Policy.Distribution.ROUNDROBIN
        self.BLOCK = Policy.Distribution.BLOCK
        self.policy = Policy(
            placement=Policy.Placement.ANYWHERE,
            host_name="",
            host_id=-1,
            distribution=Policy.Distribution.ROUNDROBIN,
            device=Policy.Device.CPU,
            affinity=Policy.Affinity.ANY,
            specific_affinity=[],
            wait_mode=Policy.WaitMode.IDLE,
            refcounted=True
        )

    def test_roundrobin(self):
        policies = []
        for i in range(len(self.nodes)*2):
            p = Policy(distribution=self.RR)
            policies.append(p)

        eval = PolicyEvaluator(self.nodes, self.policy)
        layouts = eval.evaluate(policies=policies)
        self.assertEqual(len(layouts), len(policies))
        for node in self.nodes:
            self.assertEqual(node.num_policies, 2)

    def test_block(self):
        policies = []
        for i in range(self.node_cpus * 2):
            p = Policy(distribution=self.BLOCK)
            policies.append(p)

        eval = PolicyEvaluator(self.nodes, self.policy)
        layouts = eval.evaluate(policies=policies)
        self.assertEqual(len(layouts), len(policies))

        # First half of nodes should be full
        for i in range(self.num_nodes // 2):
            self.assertEqual(self.nodes[i].num_policies, 4)

        # Last half of nodes should be empty
        for i in range(self.num_nodes // 2, self.num_nodes):
            self.assertEqual(self.nodes[i].num_policies, 0)

    def test_mix(self):
        policies = []
        policies.append(Policy(distribution=self.RR)) # Node 0, 1 policy, move to next node
        policies.append(Policy(distribution=self.BLOCK)) # Node 1, 1 policy, stay on node
        policies.append(Policy(distribution=self.BLOCK)) # Node 1, 2 policies, stay on node
        policies.append(Policy(distribution=self.BLOCK)) # Node 1, 3 policies, stay on node
        policies.append(Policy(distribution=self.RR)) # Node 1, 4 policies, move to next node
        policies.append(Policy(distribution=self.BLOCK)) # Node 2, 1 policy, stay on node
        policies.append(Policy(distribution=self.RR)) # Node 2, 2 policies, move to next node
        policies.append(Policy(distribution=self.BLOCK)) # Node 3, 1 policy, stay on node

        eval = PolicyEvaluator(self.nodes, self.policy)
        layouts = eval.evaluate(policies=policies)
        self.assertEqual(self.nodes[0].num_policies, 1)
        self.assertEqual(self.nodes[1].num_policies, 4)
        self.assertEqual(self.nodes[2].num_policies, 2)
        self.assertEqual(self.nodes[3].num_policies, 1)

    def test_host_id(self):
        p = Policy(placement=Policy.Placement.HOST_ID, host_id=2)
        eval = PolicyEvaluator(self.nodes, self.policy)
        layouts = eval.evaluate(policies=[p])
        self.assertEqual(layouts[0].h_uid, 2)

    def test_host_name(self):
        p = Policy(placement=Policy.Placement.HOST_NAME, host_name='2')
        eval = PolicyEvaluator(self.nodes, self.policy)
        layouts = eval.evaluate(policies=[p])
        self.assertEqual(layouts[0].host_name, '2')

    def test_overprovision(self):
        # Generate N policies for N total CPUs across nodes
        # Apply N policies to N/2 nodes
        policies =[]
        for i in range(self.num_nodes * self.node_cpus):
            p = Policy(distribution=self.BLOCK)
            policies.append(p)

        eval = PolicyEvaluator(self.nodes[:len(self.nodes)//2], self.policy)
        layouts = eval.evaluate(policies=policies)

        self.assertEqual(self.nodes[0].num_policies, self.num_nodes * 2)
        self.assertEqual(self.nodes[1].num_policies, self.num_nodes * 2)

    def test_affinity_any(self):
        node = NodeDescriptor()
        p = Policy()
        eval = PolicyEvaluator([node], self.policy)
        layout = eval.evaluate([p])
        self.assertEqual(layout[0].core, node.cpu_devices)

    def test_affinity_specific(self):
        node = NodeDescriptor()
        expected_cores = [1, 3, 5]
        p = Policy(affinity=Policy.Affinity.SPECIFIC, specific_affinity=expected_cores)
        eval = PolicyEvaluator([node], self.policy)
        layout = eval.evaluate([p])
        self.assertEqual(layout[0].core, expected_cores)

if __name__ == '__main__':
    unittest.main()

test/globalservices/test_policy_eval.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import sys
import os
import time

import dragon
import dragon.channels as dch
import dragon.globalservices.channel as dgchan
import dragon.globalservices.process as dgprocess

from dragon.infrastructure.channel_desc import ChannelOptions
from dragon.infrastructure.parameters import this_process
from dragon.infrastructure.facts import default_pool_muid_from_index
from dragon.infrastructure.process_desc import ProcessOptions
import dragon.utils as du


def attach_proc(ser_desc_str):
    sdesc = du.B64.str_to_bytes(ser_desc_str)
    channel = dch.Channel.attach(sdesc)
    dgchan.get_refcnt(channel.cuid)
    dgchan.release_refcnt(channel.cuid)
    channel.detach()

@unittest.skip(f"PE-42136 : Fails if run with other GS unit tests.")
class TestGSRefcounting(unittest.TestCase):

    def test_refcnt(self):

        m_uid = default_pool_muid_from_index(this_process.index)

        the_options = ChannelOptions(ref_count=True)
        the_options.local_opts.capacity = 1

        descr = dgchan.create(m_uid=m_uid, user_name=f"test_refcnt_{os.getpid()}", options=the_options)

        # refcount == 1

        cname = descr.name
        cuid = descr.c_uid

        channel = dch.Channel.attach(descr.sdesc)

        channel.detach()

        descr = dgchan.query(cuid)  # channel should be alive
        self.assertTrue(descr.state == descr.State.ACTIVE)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        puids = []

        for _ in range(0, 4):
            p = dgprocess.create(
                cmd,
                wdir,
                [os.path.basename(__file__), "attach_proc", du.B64.bytes_to_str(descr.sdesc)],
                None,
                options=options,
            )
            puids.append(p.p_uid)

        for puid in puids:
            dgprocess.join(puid, timeout=None)
            descr = dgprocess.query(puids[0])
            self.assertTrue(descr.state == descr.State.DEAD)

        # refcount == 1

        descr = dgchan.query(cuid)  # channel is alive
        self.assertTrue(descr.state == descr.State.ACTIVE)

        dgchan.release_refcnt(cuid)  # channel needs to die now

        for i in range(0, 10):
            descr = dgchan.query(cuid)  # channel is not alive
            if descr.state == descr.State.DEAD:
                break
            time.sleep(0.1)

        self.assertTrue(descr.state == descr.State.DEAD)


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "attach_proc":
        attach_proc(sys.argv[2])
    else:
        unittest.main()

test/globalservices/test_refcounting.py
DragonHPC/dragon
import unittest
import random
from parameterized import parameterized
from dragon.malloc.heapmanager import BitSet


# @MCB TODO: Need more tests and data sets

class BitSetTest(unittest.TestCase):

    @classmethod
    def setUpClass(self) -> None:
        self.max_bits = 1024
        mem_size = BitSet.size(self.max_bits)
        self.memobj = bytearray(mem_size)

    @parameterized.expand([
        ["512", 512, [0, 350, 36, 55]],
        ["1024", 1024, [0, 400, 916, 42]]
    ])
    def test_bitset_handle(self, name: str, num_bits: int, bits_to_set: [int]):
        self.assertTrue(num_bits <= self.max_bits, "Trying to test too many bits")

        bset1 = BitSet.init(num_bits, self.memobj)
        bset2 = BitSet.attach(self.memobj)

        self.assertEqual(bset1.dump_to_str("", ""), bset2.dump_to_str("", ""))

        for bit in bits_to_set:
            bset1.set_bit(bit)
            self.assertEqual(1, bset1.get_bit(bit), "Bit expected to be one, was not")
            self.assertEqual(1, bset2.get_bit(bit), "Bit expected to be one, was not")

        self.assertEqual(bset1.dump_to_str("", ""), bset2.dump_to_str("", ""))

        bset1.reset_bit(bits_to_set[-1])
        self.assertEqual(bset1.right_zeroes(0), bset2.right_zeroes(0))

        for bit in bits_to_set:
            bset2.reset_bit(bit)
            self.assertEqual(0, bset1.get_bit(bit), "Bit expected to be zero, was not")
            self.assertEqual(0, bset2.get_bit(bit), "Bit expected to be zero, was not")

        self.assertEqual(bset1.right_zeroes(0), bset2.right_zeroes(0))

    @parameterized.expand([
        ["32", 32],
        ["64", 64],
        ["128", 128],
        ["256", 256],
        ["512", 512],
        ["1024", 1024]
    ])
    def test_random_set_bit(self, name: str, num_bits: int):
        self.assertTrue(num_bits <= self.max_bits, "Trying to test too many bits")

        # Get a random shuffle of bits to set to 1
        num_ones = random.randint(1, num_bits)
        bits = [0] * (num_bits - num_ones) + [1] * num_ones
        random.shuffle(bits)

        bset = BitSet.init(num_bits, self.memobj)

        for i in range(num_bits):
            if bits[i] == 1:
                bset.set_bit(i)

        for i in range(num_bits):
            self.assertEqual(bits[i], bset.get_bit(i),
                             "Bit expected to be set to {}, is actually set to {}".format(bits[i], bset.get_bit(i)))

        for i in range(num_bits):
            bset.reset_bit(i)

        for i in range(num_bits):
            self.assertEqual(0, bset.get_bit(i), "Bit expected to be zero, was not")


if __name__ == '__main__':
    unittest.main()

test/heapmanager/test_bitset.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import mmap
from multiprocessing import shared_memory as shm
from parameterized import parameterized

from dragon.heapmanager import Heap
from dragon.locks import Type


# @MCB: Handy for debugging failing tests since we might run into segfaults and that will crash
#       the whole unittest runner.  Uncomment the line in __main__ to use.
def trace(frame, event, arg):
    print("%s, %s:%d" % (event, frame.f_code.co_filename, frame.f_lineno))
    return trace

# @MCB: Intended util function for comapring stringified heap dumps against known-good dump files.
#       Currently does nothing but spam stdout.
def check_heap_dump(expected, actual):
    expected_lines = expected.split('\n')
    actual_lines = actual.split('\n')
    for idx in range(len(expected_lines)):
        print(expected_lines[idx])
        print(actual_lines[idx])


class MallocTest():

    # Test basic init and verify stat values for a few different heap sizes
    @parameterized.expand([
        ["128mb", 27, 5, 2048],
        ["1gb", 30, 12, 4096],
        ["4gb", 32, 12, 4096]
    ])
    def test_init_and_stats(self, name: str, max_pwr: int, min_pwr: int, alignment: int):
        # Destroy the handle we got from setUp so we can re-init to test different sizes
        self.heap_hdl.destroy()
        self.heap_hdl = Heap.init(
            max_pwr, min_pwr, alignment, Type.FIFO, self.memobj)

        expct_num_freelists = max_pwr - min_pwr + 1
        self.assertEqual(1 << min_pwr, self.heap_hdl.segment_size,
                         "Handle: Incorrect segment size")
        self.assertEqual(1 << (expct_num_freelists - 1), self.heap_hdl.num_segments,
                         "Handle: Incorrect segment count")
        self.assertEqual(expct_num_freelists, self.heap_hdl.num_freelists,
                         "Handle: Incorrect number of free lists")
        self.assertEqual(0, self.heap_hdl.recovery_needed,
                         "Handle: Unexpected needed recovery")
        self.assertEqual(1, self.heap_hdl.init_handle,
                         "Handle: Expected init handle to be set")

        for i in range(self.heap_hdl.num_freelists):
            self.assertNotEqual(None, self.heap_hdl.free_lists(i),
                                "Handle: Free list {} shouldn't be empty but is".format(i))

        stats = self.heap_hdl.get_stats()
        self.assertEqual(1 << max_pwr, stats.total_size,
                         "Stats: Incorrect heap size")
        self.assertEqual(1 << min_pwr, stats.segment_size,
                         "Stats: Incorrect segment size")
        self.assertEqual(0, stats.utilization_pct,
                         "Stats: Incorrect expected utilization")

    # Verify we can stick more than one handle onto an initialized heap
    # @MCB TODO: Parameterize this to attach/detach random numbers of extra handles
    def test_attach_detach(self):
        heap_hdl2 = Heap.attach(self.memobj)

        self.assertEqual(self.heap_hdl.num_segments, heap_hdl2.num_segments,
                         "Number of segments don't match")
        self.assertEqual(self.heap_hdl.num_freelists, heap_hdl2.num_freelists,
                         "Number of free lists don't match")
        self.assertEqual(self.heap_hdl.exclusive_access, heap_hdl2.exclusive_access,
                         "Exclusive access pointers should be the same")
        for i in range(self.heap_hdl.num_freelists):
            self.assertEqual(self.heap_hdl.free_lists(i), heap_hdl2.free_lists(i),
                             "Free lists being empty doesn't match up")

        heap_hdl2.detach()

    # Do a basic allocation on the heap, verify we can acutally write into it, validate stats are correct
    def test_dragon_malloc(self):
        tmp_obj = self.heap_hdl.malloc(16)

        # Assert that we can actually write into our malloc'd space
        tmp_obj[0:16] = b'abcdefghijklmnop'
        self.assertEqual(b'abcdefghijklmnop', tmp_obj.tobytes())

        # Assert stats are correct based on a 4gb heap
        stats = self.heap_hdl.get_stats()

        self.assertTrue(stats.utilization_pct < 0.000096,
                        "Utilization percent greater than 0.000096, expected about 0.000095")
        self.assertTrue(stats.utilization_pct > 0.000094,
                        "Utilization percent lower than 0.000094, expected about 0.000095")
        self.assertEqual(21, stats.num_block_sizes,
                         "Incorrect number of block sizes after malloc")

        for i in range(stats.num_block_sizes - 1):
            self.assertEqual(1, stats.free_block_count(i),
                             "Expected number of free blocks to be 1, but got {}".format(stats.free_block_count(i)))

        self.assertEqual(0, stats.free_block_count(stats.num_block_sizes - 1),
                         "Expected final free block size to have no free blocks, but has {}".format(stats.free_block_count(stats.num_block_sizes - 1)))


    # Allocate, verify we can write, verify we can free, verify stats are correct
    def test_dragon_free(self):
        tmp_obj = self.heap_hdl.malloc(16)

        # Assert that we can actually write into our malloc'd space
        tmp_obj[0:16] = b'abcdefghijklmnop'
        self.assertEqual(b'abcdefghijklmnop', tmp_obj.tobytes())

        # Free it
        self.heap_hdl.free(tmp_obj)

        stats = self.heap_hdl.get_stats()

        self.assertEqual(0.0, stats.utilization_pct,
                         "Utilization percent not 0")

        for i in range(stats.num_block_sizes - 1):
            self.assertEqual(0, stats.free_block_count(i),
                             "Expected number of free blocks to be 0, but got {}".format(stats.free_block_count(i)))

        self.assertEqual(1, stats.free_block_count(stats.num_block_sizes - 1),
                         "Expected final free block size to have 1 free block, but got {}".format(stats.free_block_count(stats.num_block_sizes - 1)))

    # Intent of this test is to overwrite and corrupt allocations.  Because python wraps memoryview pointers into actual objects,
    #    we can't actually test this, since the returned memoryview objects know their own bounds.  Prior testing with raw pointers
    #    verified this worked.  Implemented an internal cython method to continue testing that corruption process?
    @unittest.skip("Skipping corrupt for now due to python bounds checking")
    def test_dragon_corrupt(self):
        mem1 = self.heap_hdl.malloc(32)
        mem2 = self.heap_hdl.malloc(5000)

        for i in range(5100):
            mem1[i] = b'e'

        # Don't catch exception here, should free correctly
        self.heap_hdl.free(mem1)

        # Should catch a recovery exception since we corrupted mem2
        with self.assertRaises(RuntimeError):
            self.heap_hdl.free(mem2)

        self.heap_hdl.recover()
        self.heap_hdl.free(mem2)

        stats = self.heap_hdl.get_stats()
        self.assertEqual(0.0, stats.utilization_pct,
                         "Utilization percent after recovery and free should be 0")

        for i in range(stats.num_block_sizes - 1):
            self.assertEqual(0, stats.free_block_count(i),
                             "Expected number of free blocks to be 0, but got {}".format(stats.free_block_count(i)))

        self.assertEqual(1, stats.free_block_count(stats.num_block_sizes - 1),
                         "Expected final free block size to have 1 free block, but got {}".format(stats.free_block_count(stats.num_block_sizes - 1)))


    # Make a bunch of garbage allocations, assert that all sizes that should fit into a 4GB heap are valid and
    #      all that aren't are invalid (throw exceptions)
    def test_random_malloc(self):
        good_values = [1382, 250915223, 1916, 2121290834, 449, 177372358, 2115, 360690951, 3524, 286665256,
                       191, 1241134, 3778, 29086950, 1927, 36346056, 703, 9029225, 3507, 84239849, 3085,
                       4864365, 81, 2227456, 2455, 2991, 40179276, 3798, 35426764, 832, 41696064,
                       3579, 3584099, 1329, 2473419, 2783, 314163, 1926, 207293, 1339, 20486641, 452,
                       710, 25659, 3314, 1828, 812, 957, 3262, 807,
                       3862, 4096]

        bad_values = [810373453, 1546034581, 79469219, 65647013,
                      2911824890, 3223138597, 3072378300, 31477418]

        valid_allocs = 0
        num_tests = len(good_values)
        allocs = [None] * num_tests

        # Perform mallocs with good values
        for i in range(num_tests):
            mem = self.heap_hdl.malloc(good_values[i])
            allocs[valid_allocs] = mem
            valid_allocs += 1

        # Perform mallocs with bad values, expect and catch raised exceptions
        for i in range(len(bad_values)):
            with self.assertRaises(RuntimeError):
                self.heap_hdl.malloc(bad_values[i])

        self.assertEqual(52, valid_allocs, "Expected 52 valid allocations")

        # Free valid allocations
        for i in range(valid_allocs):
            self.heap_hdl.free(allocs[i])


class MMAPTest(MallocTest, unittest.TestCase):

    def setUp(self) -> None:
        self.max_pwr = 32
        self.min_pwr = 12
        self.alignment = 4096
        self.memobj = mmap.mmap(-1, Heap.size(self.max_pwr, self.min_pwr, self.alignment, Type.FIFO))
        self.heap_hdl = Heap.init(
            self.max_pwr, self.min_pwr, self.alignment, Type.FIFO, self.memobj)

    def tearDown(self) -> None:
        if self.heap_hdl:
            self.heap_hdl.destroy()

        self.memobj.close()
        del self.memobj


class SharedMemTest(MallocTest, unittest.TestCase):

    def setUp(self) -> None:
        self.max_pwr = 32
        self.min_pwr = 12
        self.alignment = 4096
        self.shmemobj = shm.SharedMemory(create=True, size=Heap.size(self.max_pwr, self.min_pwr, self.alignment, Type.FIFO))
        self.memobj = self.shmemobj.buf
        self.heap_hdl = Heap.init(
            self.max_pwr, self.min_pwr, self.alignment, Type.FIFO, self.memobj)

    def tearDown(self) -> None:
        if self.heap_hdl:
            self.heap_hdl.destroy()

        self.shmemobj.close()
        self.shmemobj.unlink()
        del self.shmemobj


if __name__ == '__main__':
    unittest.main(verbosity=2)
    # sys.settrace(trace)

test/heapmanager/test_heapmanager.py
DragonHPC/dragon
import copy
import os
import unittest
import dragon.infrastructure.parameters as dip


class LaunchParameterTest(unittest.TestCase):
    def setUp(self) -> None:
        self.all_names = dip.LaunchParameters.all_parm_names()

        # need this to keep this test case from poisoning
        # other tests that might be run together with other ones
        # in the same process
        self.old_env = dict(os.environ)
        self.old_this_proc = copy.copy(dip.this_process)

    def tearDown(self) -> None:
        for k, v in self.old_env.items():
            os.environ[k] = v

        for k, v in self.old_this_proc.env().items():
            os.environ[k] = v

        dip.this_process = self.old_this_proc

    def test_smoke(self):
        mlp = dip.LaunchParameters.from_env()
        mlp2 = dip.LaunchParameters.from_env()

        for name in self.all_names:
            self.assertEqual(getattr(mlp, name), getattr(mlp2, name), f'looking for {name}')

    def test_from_env(self):
        key_to_test = 'INF_SEG_SZ'
        old_val = None
        ev_name = dip.LaunchParameters._ev_name_from_pt(key_to_test)

        if ev_name in os.environ:
            old_val = os.environ[ev_name]

        os.environ[ev_name] = str(17)

        mlp = dip.LaunchParameters.from_env()

        self.assertEqual(mlp.inf_seg_sz, 17)

        if old_val is not None:
            os.environ[ev_name] = old_val

    def test_to_env(self):
        key_to_test = 'INF_SEG_SZ'
        ev_name = dip.LaunchParameters._ev_name_from_pt(key_to_test)
        mlp = dip.LaunchParameters.from_env()
        mlp.inf_seg_sz = 1001
        the_env = mlp.env()

        self.assertEqual(mlp.inf_seg_sz, int(the_env[ev_name]))


if __name__ == "__main__":
    unittest.main()

test/infrastructure/env_parameter_tests.py
DragonHPC/dragon
from contextlib import closing
import io
import os
import sys
import unittest

from dragon.infrastructure.util import NewlineStreamWrapper, _get_access_modes


class NewlineStreamWrapperTest(unittest.TestCase):

    def _make_pipe(self):
        r, w = os.pipe()
        return os.fdopen(r), os.fdopen(w, 'w')

    def test_auto_intents(self):
        reader, writer = map(NewlineStreamWrapper, self._make_pipe())
        with closing(reader):
            self.assertTrue(reader.read_intent)
            self.assertFalse(reader.write_intent)
        with closing(writer):
            self.assertFalse(writer.read_intent)
            self.assertTrue(writer.write_intent)

    def test_inverse_intents(self):
        r, w = self._make_pipe()
        with self.assertWarns(UserWarning):
            reader = NewlineStreamWrapper(r, read_intent=False, write_intent=True)
        with closing(reader):
            self.assertFalse(reader.read_intent)
            self.assertTrue(reader.write_intent)
        with self.assertWarns(UserWarning):
            writer = NewlineStreamWrapper(w, read_intent=True, write_intent=False)
        with closing(writer):
            self.assertTrue(writer.read_intent)
            self.assertFalse(writer.write_intent)

    def test_stdin(self):
        readable, writeable = _get_access_modes(sys.stdin)
        self.assertTrue(readable)
        # XXX Unless redirected from a file, stdin actually has access
        # XXX mode O_RDWR even though sys.stdin.mode is 'r'.
        if sys.stdin.isatty():
            self.assertTrue(writeable)
        else:
            self.assertFalse(writeable)

    def test_stdin_wrapper(self):
        stream = NewlineStreamWrapper(sys.stdin)
        # Do not call stream.close() otherwise sys.stdin will be closed
        self.assertTrue(stream.read_intent)
        # XXX Unless redirected from a file, stdin actually has access
        # XXX mode O_RDWR even though sys.stdin.mode is 'r'.
        if sys.stdin.isatty():
            self.assertTrue(stream.write_intent)
        else:
            self.assertFalse(stream.write_intent)
        
    def test_stdout(self):
        readable, writeable = _get_access_modes(sys.stdout)
        # XXX Unless redirected from a file, stdout actually has access
        # XXX mode O_RDWR even though sys.stdout.mode is 'w'.
        if sys.stdout.isatty():
            self.assertTrue(readable)
        else:
            self.assertFalse(readable)
        self.assertTrue(writeable)

    def test_stdout_wrapper(self):
        stream = NewlineStreamWrapper(sys.stdout)
        # Do not call stream.close() otherwise sys.stdout will be closed
        # XXX Unless redirected from a file, stdout actually has access
        # XXX mode O_RDWR even though sys.stdout.mode is 'w'.
        if sys.stdout.isatty():
            self.assertTrue(stream.read_intent)
        else:
            self.assertFalse(stream.read_intent)
        self.assertTrue(stream.write_intent)

    def test_poll(self):
        r, w = self._make_pipe()
        reader = NewlineStreamWrapper(r)
        with closing(w):
            w.write("line 1\nline 2\n")
        with closing(reader):
            self.assertTrue(reader.poll())
        with self.assertRaises(ValueError):
            reader.poll()

    def test_poll_stringio(self):
        # Although one might think io.StringIO() is a valid file-like object
        # that can be polled, their fileno() method raises
        # io.UnsupportedOperation, e.g.:
        #
        #   >>> import io
        #   >>> f = io.StringIO()
        #   >>> hasattr(f, 'fileno')
        #   True
        #   >>> f.fileno
        #   <built-in method fileno of _io.StringIO object at 0x7f399d5363a0>
        #   >>> f.fileno()
        #   Traceback (most recent call last):
        #     File "<stdin>", line 1, in <module>
        #   io.UnsupportedOperation: fileno
        #
        stream = NewlineStreamWrapper(io.StringIO())
        with closing(stream):
            with self.assertRaises(ValueError):
                stream.poll()

    def test_poll_bytesio(self):
        # Although one might think io.BytesIO() is a valid file-like object
        # that can be polled, their fileno() method raises
        # io.UnsupportedOperation, e.g.:
        #
        #   >>> import io
        #   >>> f = io.BytesIO()
        #   >>> hasattr(f, 'fileno')
        #   True
        #   >>> f.fileno
        #   <built-in method fileno of _io.BytesIO object at 0x7efef9feac70>
        #   >>> f.fileno()
        #   Traceback (most recent call last):
        #     File "<stdin>", line 1, in <module>
        #   io.UnsupportedOperation: fileno
        #
        stream = NewlineStreamWrapper(io.BytesIO())
        with closing(stream):
            with self.assertRaises(ValueError):
                stream.poll()

    def test_recv(self):
        r, w = self._make_pipe()
        reader = NewlineStreamWrapper(r)
        with closing(w):
            w.write("line 1\nline 2\n")
        with closing(reader):
            self.assertTrue(reader.poll())
            self.assertEqual(reader.recv(), "line 1\n")
            self.assertTrue(reader.poll())
            self.assertEqual(reader.recv(), "line 2\n")
            # Once all data has been read, poll() will return True and recv()
            # will return an empty string.
            self.assertTrue(reader.poll())
            self.assertEqual(reader.recv(), '')

    def test_send(self):
        r, w = self._make_pipe()
        writer = NewlineStreamWrapper(w)
        with closing(writer):
            writer.send("line 1")
            writer.send("line 2")
        with closing(r):
            self.assertEqual(r.read(), "line 1\nline 2\n")


if __name__ == "__main__":
    unittest.main()

test/infrastructure/newline_stream_wrapper_test.py
DragonHPC/dragon
import dragon

import argparse
import multiprocessing
import time
import threading


def get_args():
    parser = argparse.ArgumentParser(description='All-all connection/queue test')

    parser.add_argument('--num_workers', type=int, default=4,
                        help='number of workers in all-all')

    parser.add_argument('--iterations', type=int, default=100,
                        help='number of iterations to do')

    parser.add_argument('--burn_iterations', type=int, default=5,
                        help='number of iterations to burn first time')

    parser.add_argument('--lg_message_size', type=int, default=4,
                        help='log base 2 of msg size between each pair')

    parser.add_argument('--dragon', action='store_true', help='run with dragon objs')
    parser.add_argument('--queues', action='store_true', help='measure aa with per-worker queues')
    my_args = parser.parse_args()

    assert my_args.lg_message_size <= 24, 'arbitrary limit: 16MB message size'

    return my_args


def test_worker_conn(links, result_link, msg_size, total_iterations):
    result_timings = []

    my_msg = bytearray(msg_size)

    def send_stuff():
        for link in links:
            link.send(my_msg)

    def recv_stuff():
        for link in links:
            _ = link.recv()

    for _ in range(total_iterations):
        start = time.perf_counter()

        send_th = threading.Thread(target=send_stuff)
        recv_th = threading.Thread(target=recv_stuff)

        send_th.start()
        recv_th.start()
        send_th.join()
        recv_th.join()

        result_timings.append(time.perf_counter() - start)

    result_link.send(result_timings)


# messages can mix between workers between iterations - there really
# ought to be a barrier - but the books balance at the end and this
# is meant to be more of a stress/comparison test anyway than
# a serious all-all.
def test_worker_queue(queues, my_idx, result_link, msg_size, total_iterations):
    result_timings = []
    my_msg = bytearray(msg_size)
    others = set(range(len(queues)))
    others.discard(my_idx)

    def send_stuff():
        for idx, q in enumerate(queues):
            if idx != my_idx:
                q.put(my_msg)

    def recv_stuff():
        for _ in range(len(others)):
            msg = queues[my_idx].get()
            assert msg_size == len(msg)

    for it_cnt in range(total_iterations):
        start = time.perf_counter()

        send_th = threading.Thread(target=send_stuff)
        recv_th = threading.Thread(target=recv_stuff)

        send_th.start()
        recv_th.start()
        send_th.join()
        recv_th.join()

        result_timings.append(time.perf_counter() - start)

    result_link.send(result_timings)


def run_test(with_queues, nworkers, its, burns, msg_sz):
    total_iterations = burns + its

    # half duplex return links
    result_links = [multiprocessing.Pipe(duplex=False) for _ in range(nworkers)]

    if with_queues:
        print(f'using {nworkers} Queues')
        queues = [multiprocessing.Queue(maxsize=nworkers) for _ in range(nworkers)]

        all_workers = [multiprocessing.Process(target=test_worker_queue,
                                               args=(queues, i,
                                                     result_links[i][1], msg_sz,
                                                     total_iterations))
                       for i in range(nworkers)]
    else:
        print(f'using {nworkers * (nworkers - 1) // 2} Pipes')
        all_links = {}

        # full duplex links; nworkers * (nworkers-1)/2 of them
        for i in range(nworkers):
            for j in range(i, nworkers):
                all_links[i, j] = multiprocessing.Pipe()

        all_workers = []
        for i in range(nworkers):
            # first end if other's index less than yours,
            # second end if other's index greater than yours
            links = []
            for k in range(nworkers):
                if k != i:
                    links.append(all_links[min(i, k), max(i, k)][int(k < i)])

            proc = multiprocessing.Process(target=test_worker_conn,
                                           args=(links, result_links[i][1],
                                                 msg_sz, total_iterations))

            all_workers.append(proc)

    for worker in all_workers:
        worker.start()

    avg_times = []
    total_times = []
    for i in range(nworkers):
        result = result_links[i][0].recv()[burns:]
        total_times.append(sum(result))
        avg_times.append(total_times[i] / its)

    avg_aa_completion_time = sum(avg_times) / nworkers  # seconds
    bandwidth = msg_sz * nworkers / avg_aa_completion_time  # bytes/second

    msg_size_in_k = msg_sz / (2 ** 10)
    bw_in_mb_sec = bandwidth / (2 ** 20)
    completion_time_in_ms = avg_aa_completion_time * 1000

    for worker in all_workers:
        worker.join()

    print(f'{nworkers=} {its=} {msg_size_in_k=} {completion_time_in_ms=:.3f} {bw_in_mb_sec=:.2f}')


def main():
    args = get_args()

    if args.dragon:
        import dragon.infrastructure.parameters as dparms

        if dparms.this_process.my_puid <= 1:
            print('not running under dragon infrastructure, exiting')
            exit(-3)

        print('using dragon runtime')
        multiprocessing.set_start_method('dragon')
    else:
        print('using regular mp libs/calls')
        multiprocessing.set_start_method('spawn')

    run_test(with_queues=args.queues, nworkers=args.num_workers,
             msg_sz=2 ** args.lg_message_size,
             its=args.iterations, burns=args.burn_iterations)


if __name__ == '__main__':
    exit(main())

test/integration/aa_bench.py
DragonHPC/dragon
#!/usr/bin/env python3

"""An example target for start_single...
   Two scripts in one - a 'driver' and a 'worker'.

   The 'driver' makes a couple channels thru the interface and uses them to
   communicate with a worker it launches.
"""

import sys
import time

import dragon.channels as dchan

import dragon.globalservices.api_setup as dga
import dragon.globalservices.process as dgp
import dragon.globalservices.channel as dgch
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.parameters as dparm
import pickle


def fake_blocking_recv_bytes(handle):
    while True:
        try:
            rv = handle.recv_bytes()
            break
        except dchan.ChannelRecvError as cre:
            if cre.ex_code == dchan.ChannelRecvEx.Errors.CH_EMPTY:
                time.sleep(0.001)
            else:
                raise cre
    return rv


def driver():
    print('hello from the driver, p_uid {}'.format(dparm.this_process.my_puid))
    dga.connect_to_infrastructure()
    print('connected to the infrastructure!')

    default_pool_muid = dfacts.default_pool_muid_from_index(dparm.this_process.index)
    channel_1_desc = dgch.create(default_pool_muid)
    channel_2_desc = dgch.create(default_pool_muid)

    print('made a couple channels')
    print('channel 1: {}'.format(channel_1_desc))
    print('channel 2: {}'.format(channel_2_desc))

    channel_1 = dchan.Channel.attach(channel_1_desc.sdesc)
    channel_2 = dchan.Channel.attach(channel_2_desc.sdesc)

    print('attached to cuid: {}'.format(channel_1.cuid))
    print('attached to cuid: {}'.format(channel_2.cuid))

    # gonna send out the ids because that's more multi-nodey
    target_input_id = channel_1.cuid
    target_output_id = channel_2.cuid

    our_argdata = pickle.dumps((target_input_id, target_output_id))
    print('starting target process')
    proc_desc = dgp.create_with_argdata(exe=sys.executable,
                                        run_dir='', env={},
                                        args=[sys.argv[0], 't'],
                                        argdata=our_argdata)

    sender = channel_1.sendh()
    message = 'hyenas'.encode()
    sender.open()
    sender.send_bytes(message)
    sender.close()

    recvr = channel_2.recvh()
    recvr.open()
    rec_msg = fake_blocking_recv_bytes(recvr)
    recvr.close()

    if rec_msg != message:
        print('sent {} but got back {}'.format(message, rec_msg))
    else:
        print('messages matched')

    print('joining target process')
    dgp.join(proc_desc.p_uid)

    dgch.destroy(channel_1.cuid)
    dgch.destroy(channel_2.cuid)
    print('destroyed the channels')
    exit(0)


def target():
    print('+++hello from target p_uid {}'.format(dparm.this_process.my_puid))
    dga.connect_to_infrastructure()
    print('+++connected to infrastructure')
    args = pickle.loads(dga._ARG_PAYLOAD)
    inbound_cuid, outbound_cuid = args
    print(f'+++inbound cuid: {inbound_cuid}\t outbound cuid {outbound_cuid}')

    input_desc = dgch.query(inbound_cuid)
    output_desc = dgch.query(outbound_cuid)

    print('+++target channels queried')

    input_chan = dchan.Channel.attach(input_desc.sdesc)
    output_chan = dchan.Channel.attach(output_desc.sdesc)

    print('+++target channels attached to')

    recvr = input_chan.recvh()
    sender = output_chan.sendh()
    recvr.open()
    sender.open()
    msg = fake_blocking_recv_bytes(recvr)
    print('+++got: {}'.format(msg))
    sender.send_bytes(msg)
    sender.close()
    recvr.close()
    input_chan.detach()
    output_chan.detach()
    print('+++target channels detached from')
    exit(0)


def usage_exit():
    print('usage: [d|t]')
    exit(1)


if __name__ == "__main__":
    if 1 == len(sys.argv):
        driver()

    if 2 != len(sys.argv):
        usage_exit()

    if 'd' == sys.argv[1]:
        driver()
    elif 't' == sys.argv[1]:
        target()
    else:
        usage_exit()

test/integration/single_example.py
DragonHPC/dragon
#!/usr/bin/env python3
"""
Basic test of multiprocessing message passing through multiprocessing.Pipe.

"""
import sys

import dragon.mpbridge
import dragon.infrastructure.connection
import dragon.infrastructure.log_setup as dlog

import multiprocessing
import time
import dragon.infrastructure.util as dutil
import logging
import mp_bench.util.test_util as test_util


# could this be a deadly embrace sort of thing
# given that open() happens and something
# is being asked for out of order in full duplex world?
# I think that could have to do with this...


def external_single_echo(conn, resp_size, idx):
    conn.recv()
    conn.send(bytes(resp_size))


class SendReceiveTest:

    def __init__(self):
        parser = test_util.add_default_args('Send and receive a message of varying size using a Pipe')
        parser = test_util.add_connections_args(parser)
        parser.add_argument('--response_size', type=int, default=4, help='size of message pass back')
        parser.add_argument('--dragon', action='store_true', help='run a la single node dragon')
        self.args = test_util.process_args(parser)
        self.results = None

        if self.args.dragon:
            multiprocessing.set_start_method('dragon')  # could call connect to infrastructure here if we wanted in
            # check avail I guess.
            self.Pipe = dragon.infrastructure.connection.Pipe
            self.Process = dragon.mpbridge.DragonProcess
        else:
            multiprocessing.set_start_method('spawn')
            self.Pipe = multiprocessing.Pipe
            self.Process = multiprocessing.Process

    def run(self):

        times = test_util.iterations(self.do_send_receive, self.args)
        self.results = test_util.Results(self.args, times, "send_receive_test")
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump()

    def do_send_receive(self):

        num_workers = self.args.num_workers

        message = bytes(self.args.message_size)

        pipeconns = [self.Pipe() for _ in range(num_workers)]

        processes = []
        for idx in range(num_workers):
            processes.append(self.Process(target=external_single_echo,
                                          args=(pipeconns[idx][1],
                                                self.args.response_size,
                                                idx)))

        start_time = time.time_ns()

        for idx, handle in enumerate(processes):
            handle.start()

        for i in range(num_workers):
            pipeconns[i][0].send(message)

        for i in range(num_workers):
            pipeconns[i][0].recv()

        for idx, handle in enumerate(processes):
            handle.join()

        return time.time_ns() - start_time


if __name__ == '__main__':
    dlog.setup_logging(basename='head', the_level=logging.DEBUG)
    test = SendReceiveTest()
    print(test.run())

test/integration/smoke.py
DragonHPC/dragon
#!/usr/bin/env python3

"""An ultra simple way to start the infrastructure (shep & GS) and
   start a script under it.  The script can do what it likes.
"""
import atexit
import logging
import threading
import multiprocessing
import os
import pickle
import subprocess
import sys
import traceback

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.globalservices.api_setup as dapi
import dragon.globalservices.process as dproc
import dragon.globalservices.process_desc as pdesc
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.log_setup as log_setup
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.parameters as dparm
import dragon.utils as du
import dragon.localservices.server as sserver
import support.util as tsu


def shep_startup_channels_env(shep_stdin_queue, shep_stdout_queue, env_update, name_addition=''):
    log_setup.setup_logging(basename='shep_' + name_addition, the_level=logging.DEBUG)
    log = logging.getLogger('shep startup')
    log.info('--------------------shep started-----------------------')

    try:
        os.environ[dfacts.ENV_MODE] = dfacts.SINGLE_MODE
        os.environ[dfacts.ENV_DEFAULT_SEG_SZ] = str(dfacts.ONE_GB)
        os.environ[dfacts.ENV_INF_SEG_SZ] = str(dfacts.ONE_GB)

        shep = sserver.Shepherd(shep_stdin_queue, shep_stdout_queue, gs_start_cmd=sys.executable,
                                gs_args=['-c', 'import dragon.globalservices.server as dgs; dgs.single()'],
                                gs_env=env_update, create_channels_pools=True)

        shep.run()
        log.info('-----Shepherd exited normally')
    except Exception as e:
        log.exception('++++Shep blew up')
        print('shepherd blew up')
        print(f'{e}')
        traceback.print_exc()

    sys.stdout.flush()
    sys.stderr.flush()


class Context:
    def __init__(self):
        self.la_tag_cnt = 0
        self.mode = dfacts.SINGLE_MODE

        self.shep_stdin_rh, self.shep_stdin_wh = multiprocessing.Pipe(duplex=False)
        self.shep_stdout_rh, self.shep_stdout_wh = multiprocessing.Pipe(duplex=False)

        username = subprocess.check_output('whoami').decode().strip()
        self.pool_name = 'standalone_' + username
        self.pool_size = 2 ** 30
        self.pool_uid = dfacts.FIRST_MUID + 2 ** 20
        self.test_pool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid)

        # will need to make a 'test pool' and 'test channels' and pass
        # those to GS thru environment, by means of putting them into the shepherd dut's
        # own startup environment, then inherited.

        # these channels are ones that normally would be in the infrastructure
        # pool; they are used by the testbench to let it call the gs API
        # and get answers back in the usual way.

        # faking these up and making them really large so they won't be likely to
        # clash with anything assigned by GS in a test
        self.gs_return_cuid = 2 ** 64 - 17
        self.proc_gs_return_chan = dch.Channel(self.test_pool, self.gs_return_cuid)
        self.proc_gs_return_rh = dconn.Connection(inbound_initializer=self.proc_gs_return_chan)

        # not used in this test, but just for uniform practice.
        self.shep_return_cuid = 2 ** 64 - 18
        self.proc_shep_return_chan = dch.Channel(self.test_pool, self.shep_return_cuid)
        self.proc_shep_return_rh = dconn.Connection(inbound_initializer=self.proc_shep_return_chan)

        self.shep_dut = None
        self.inf_pool = None
        self.shep_input_chan = None
        self.shep_input_wh = None
        self.bela_input_chan = None
        self.bela_input_rh = None
        self.gs_input_chan = None
        self.gs_input_wh = None
        self.head_puid = None

    def start_infrastructure(self):
        log_name_addition = subprocess.check_output('whoami').decode().strip()

        # Make a list of channels to pass thru the environment to GS for testing purposes.
        # In reality these channels would be created by GS in advance of starting
        # the infrastructure-using process.
        test_chan_list = [(self.gs_return_cuid, self.test_pool.serialize(),
                           self.proc_gs_return_chan.serialize(), False),
                          (self.shep_return_cuid, self.test_pool.serialize(),
                           self.proc_shep_return_chan.serialize(), False)]

        val = pickle.dumps(test_chan_list)
        env_update = {dfacts.GS_TEST_CH_EV: du.B64.bytes_to_str(val),
                      dfacts.GS_LOG_BASE: 'gs_' + log_name_addition}

        shep_start_args = (self.shep_stdin_rh, self.shep_stdout_wh, env_update)

        self.shep_dut = multiprocessing.Process(target=shep_startup_channels_env,
                                                args=shep_start_args,
                                                kwargs={'name_addition': log_name_addition})

        self.shep_dut.start()

        # start shepherd startup protocol for single node
        self.shep_stdin_wh.send(dmsg.BENodeIdxSH(tag=self.la_tag_cnt, node_idx=0).serialize())
        self.la_tag_cnt += 1

        ping_be_msg = tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHPingBE)

        env2b = du.B64.str_to_bytes

        self.inf_pool = dmm.MemoryPool.attach(env2b(ping_be_msg.inf_pd))

        self.shep_input_chan = dch.Channel.attach(env2b(ping_be_msg.shep_cd))
        self.shep_input_wh = dconn.Connection(outbound_initializer=self.shep_input_chan)

        self.bela_input_chan = dch.Channel.attach(env2b(ping_be_msg.be_cd))
        self.bela_input_rh = dconn.Connection(inbound_initializer=self.bela_input_chan)

        self.gs_input_chan = dch.Channel.attach(env2b(ping_be_msg.gs_cd))
        self.gs_input_wh = dconn.Connection(outbound_initializer=self.gs_input_chan)

        # test bench is acting as though it is the GS api, and the normal module import
        # initialization won't have triggered, so we override the connections so the API
        # will work.

        dapi.test_connection_override(test_gs_input=self.gs_input_wh,
                                      test_gs_return=self.proc_gs_return_rh,
                                      test_gs_return_cuid=self.gs_return_cuid,
                                      test_shep_input=self.shep_input_wh,
                                      test_shep_return=self.proc_shep_return_rh,
                                      test_shep_return_cuid=self.shep_return_cuid)

        # complete shepherd startup.
        self.shep_input_wh.send(dmsg.BEPingSH(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.bela_input_rh, dmsg.SHChannelsUp)

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSIsUp)

    def cleanup(self):
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHalted)

        self.shep_input_wh.send(dmsg.SHTeardown(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.bela_input_rh, dmsg.SHHaltBE)

        self.shep_stdin_wh.send(dmsg.BEHalted(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHHalted)

        self.shep_stdin_rh.close()
        self.shep_stdin_wh.close()
        self.shep_stdout_rh.close()
        self.shep_stdout_wh.close()

        try:
            self.shep_input_wh.close()
        except:
            pass

        try:
            self.bela_input_rh.close()
        except:
            pass

        try:
            self.gs_input_wh.close()
        except:
            pass

        try:
            self.proc_gs_return_rh.close()
        except:
            pass

        try:
            self.proc_shep_return_rh.close()
        except:
            pass

        try:
            self.shep_input_chan.detach()
        except:
            pass

        try:
            self.bela_input_chan.detach()
        except:
            pass

        try:
            self.gs_input_chan.detach()
        except:
            pass

        try:
            self.proc_gs_return_chan.destroy()
        except:
            pass

        try:
            self.proc_shep_return_chan.destroy()
        except:
            pass

        try:
            self.inf_pool.detach()
        except:
            pass

        try:
            self.test_pool.destroy()
        except:
            pass

        if self.shep_dut is not None:
            self.shep_dut.kill()

    def start_thing(self, start_args):
        res = dproc.create(exe='python3', run_dir='',
                           args=start_args, env={},
                           options=pdesc.ProcessOptions(make_inf_channels=True))

        self.head_puid = res.p_uid

    def monitor_output(self, optional_barrier=None):
        while True:
            msg = dmsg.parse(self.bela_input_rh.recv())
            if isinstance(msg, dmsg.SHFwdOutput):
                if msg.fd_num == msg.FDNum.STDOUT.value:
                    sys.stdout.write(msg.data)
                    sys.stdout.flush()
                elif msg.fd_num == msg.FDNum.STDERR.value:
                    sys.stderr.write(msg.data)
                    sys.stderr.flush()
                if optional_barrier is not None:
                    try:
                        optional_barrier.wait()
                    except threading.BrokenBarrierError:
                        pass
            elif isinstance(msg, dmsg.GSHeadExit):
                print('+++ head proc exited, code {}'.format(msg.exit_code))
                break
            else:
                print('unexpected message: {}'.format(msg))


default_parent_stdin_filename = f'quokka_{os.getpid()}'
parent_raw_input_fh = None
coordinate_repl_prompt_barrier = threading.Barrier(2)

def get_raw_input_from_parent_via_fd(prompt='D>>'):
    'Intended to be called as substitute for raw_input() in child process.'
    global parent_raw_input_fh, default_parent_stdin_filename
    if parent_raw_input_fh is None:
        print(f'DBG {default_parent_stdin_filename=}', flush=True)
        parent_raw_input_fh = open(default_parent_stdin_filename, 'rt')
    sys.stdout.flush()
    input_text = parent_raw_input_fh.readline()
    return input_text.rstrip()

def capture_raw_input_as_parent_via_fd(
        filename=default_parent_stdin_filename,
        barrier=coordinate_repl_prompt_barrier
):
    os.mkfifo(filename)
    barrier.wait()
    print(f'DBG parent process {os.getpid()=}')
    try:
        with open(filename, 'wt') as parent_raw_input_fh:
            while True:
                print('D>> ', end='', flush=True)
                input_text = sys.stdin.readline()
                parent_raw_input_fh.write(input_text)
                parent_raw_input_fh.flush()
                try:
                    barrier.wait(timeout=1)
                except threading.BrokenBarrierError:
                    pass
    except Exception as e:
        print(f'Capture ending, reason: {e}', flush=True)

def start_repl(pipe_filename):
    import code
    g = globals()
    g['default_parent_stdin_filename'] = pipe_filename
    code.interact(readfunc=get_raw_input_from_parent_via_fd, local=g)


def main(start_args, interactive=False):
    the_ctx = Context()
    try:
        the_ctx.start_infrastructure()
    except Exception as e:
        print('Got exception {} in inf startup'.format(e))
        traceback.print_exc()
        the_ctx.cleanup()
        exit(1)

    if interactive:
        capture_raw_input_thread = threading.Thread(
            target=capture_raw_input_as_parent_via_fd,
            args=(default_parent_stdin_filename, coordinate_repl_prompt_barrier)
        )
        capture_raw_input_thread.start()

    try:
        the_ctx.start_thing(start_args)
    except Exception as e:
        print('Got exception {} in program start'.format(e))
        traceback.print_exc()
        the_ctx.cleanup()
        exit(1)

    # try this to handle ctl-C while stuff runs.
    atexit.register(the_ctx.cleanup)

    if interactive:
        monitor_output_thread = threading.Thread(
            target=Context.monitor_output,
            args=(the_ctx, coordinate_repl_prompt_barrier)
        )
        monitor_output_thread.start()
        monitor_output_thread.join()
    else:
        the_ctx.monitor_output()

    atexit.unregister(the_ctx.cleanup)
    the_ctx.cleanup()

    try:
        # Can not be done in start_repl() because still in use at that time.
        os.unlink(default_parent_stdin_filename)
    except:
        pass

    exit(0)


if __name__ == "__main__":
    if 1 == len(sys.argv):
        args = [
            '-c',
            f'import start_single; start_single.start_repl({default_parent_stdin_filename!r})'
        ]
        interactive = True
    else:
        args = sys.argv[1:]
        interactive = False

    main(args, interactive)

test/integration/start_single.py
DragonHPC/dragon
#!/usr/bin/env python3

"""An ultra simple way to start the infrastructure (LS & GS) and
   start a script under it.  The script can do what it likes.
"""
import atexit
import logging
import threading
import multiprocessing
import os
import pickle
import subprocess
import sys
import traceback

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.globalservices.api_setup as dapi
import dragon.globalservices.process as dproc
import dragon.globalservices.process_desc as pdesc
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.log_setup as log_setup
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.parameters as dparm
import dragon.localservices.local_svc as dsls
import dragon.utils as du
import support.util as tsu


def start_ls(shep_stdin_queue, shep_stdout_queue, env_update, name_addition=''):
    log_setup.setup_logging(basename='ls_' + name_addition, the_level=logging.DEBUG)
    log = logging.getLogger('ls-gs integration test channels, ls starts gs')
    log.info('--------------------ls started-----------------------')

    dparm.this_process.mode = dfacts.SINGLE_MODE
    os.environ.update(dparm.this_process.env())
    log.info('starting')
    try:
        dsls.single(test_stdin=shep_stdin_queue, test_stdout=shep_stdout_queue,
                    gs_args=[sys.executable, '-c', dfacts.GS_SINGLE_LAUNCH_CMD],
                    gs_env=env_update)
    except RuntimeError as rte:
        log.exception('runtime failure\n')
        shep_stdout_queue.send(dmsg.AbnormalTermination(tag=dsls.get_new_tag(),
                                                        err_info=f'{rte!s}').serialize())
    log.info('exited')


class Context:
    def __init__(self):
        self.la_tag_cnt = 0
        self.mode = dfacts.SINGLE_MODE

        self.shep_stdin_rh, self.shep_stdin_wh = multiprocessing.Pipe(duplex=False)
        self.shep_stdout_rh, self.shep_stdout_wh = multiprocessing.Pipe(duplex=False)

        username = subprocess.check_output('whoami').decode().strip()
        self.pool_name = 'standalone_' + username
        self.pool_size = 2 ** 30
        self.pool_uid = dfacts.FIRST_MUID + 2 ** 20
        self.test_pool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid)

        # will need to make a 'test pool' and 'test channels' and pass
        # those to GS thru environment, by means of putting them into the shepherd dut's
        # own startup environment, then inherited.

        # these channels are ones that normally would be in the infrastructure
        # pool; they are used by the testbench to let it call the gs API
        # and get answers back in the usual way.

        # faking these up and making them really large so they won't be likely to
        # clash with anything assigned by GS in a test
        self.gs_return_cuid = 2 ** 64 - 17
        self.proc_gs_return_chan = dch.Channel(self.test_pool, self.gs_return_cuid)
        self.proc_gs_return_rh = dconn.Connection(inbound_initializer=self.proc_gs_return_chan)

        # not used in this test, but just for uniform practice.
        self.shep_return_cuid = 2 ** 64 - 18
        self.proc_shep_return_chan = dch.Channel(self.test_pool, self.shep_return_cuid)
        self.proc_shep_return_rh = dconn.Connection(inbound_initializer=self.proc_shep_return_chan)

        self.shep_dut = None
        self.inf_pool = None
        self.shep_input_chan = None
        self.shep_input_wh = None
        self.bela_input_chan = None
        self.bela_input_rh = None
        self.gs_input_chan = None
        self.gs_input_wh = None
        self.head_puid = None

    def start_infrastructure(self):
        log_name_addition = subprocess.check_output('whoami').decode().strip()

        # Make a list of channels to pass thru the environment to GS for testing purposes.
        # In reality these channels would be created by GS in advance of starting
        # the infrastructure-using process.
        test_chan_list = [(self.gs_return_cuid, self.test_pool.serialize(),
                           self.proc_gs_return_chan.serialize(), False),
                          (self.shep_return_cuid, self.test_pool.serialize(),
                           self.proc_shep_return_chan.serialize(), False)]

        val = pickle.dumps(test_chan_list)
        env_update = {dfacts.GS_TEST_CH_EV: du.B64.bytes_to_str(val),
                      dfacts.GS_LOG_BASE: 'gs_' + log_name_addition}

        shep_start_args = (self.shep_stdin_rh, self.shep_stdout_wh, env_update)

        self.shep_dut = multiprocessing.Process(target=start_ls,
                                                args=shep_start_args,
                                                kwargs={'name_addition': log_name_addition})

        self.shep_dut.start()

        # start shepherd startup protocol for single node
        self.shep_stdin_wh.send(dmsg.BENodeIdxSH(tag=self.la_tag_cnt, node_idx=0).serialize())
        self.la_tag_cnt += 1

        ping_be_msg = tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHPingBE)

        env2b = du.B64.str_to_bytes

        self.inf_pool = dmm.MemoryPool.attach(env2b(ping_be_msg.inf_pd))

        self.shep_input_chan = dch.Channel.attach(env2b(ping_be_msg.shep_cd))
        self.shep_input_wh = dconn.Connection(outbound_initializer=self.shep_input_chan)

        self.bela_input_chan = dch.Channel.attach(env2b(ping_be_msg.be_cd))
        self.bela_input_rh = dconn.Connection(inbound_initializer=self.bela_input_chan)

        self.gs_input_chan = dch.Channel.attach(env2b(ping_be_msg.gs_cd))
        self.gs_input_wh = dconn.Connection(outbound_initializer=self.gs_input_chan)

        # test bench is acting as though it is the GS api, and the normal module import
        # initialization won't have triggered, so we override the connections so the API
        # will work.

        dapi.test_connection_override(test_gs_input=self.gs_input_wh,
                                      test_gs_return=self.proc_gs_return_rh,
                                      test_gs_return_cuid=self.gs_return_cuid,
                                      test_shep_input=self.shep_input_wh,
                                      test_shep_return=self.proc_shep_return_rh,
                                      test_shep_return_cuid=self.shep_return_cuid)

        # complete shepherd startup.
        self.shep_input_wh.send(dmsg.BEPingSH(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.bela_input_rh, dmsg.SHChannelsUp)

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSIsUp)

    def cleanup(self):
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHalted)

        self.shep_input_wh.send(dmsg.SHTeardown(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.bela_input_rh, dmsg.SHHaltBE)

        self.shep_stdin_wh.send(dmsg.BEHalted(tag=self.la_tag_cnt).serialize())
        self.la_tag_cnt += 1

        tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHHalted)

        self.shep_stdin_rh.close()
        self.shep_stdin_wh.close()
        self.shep_stdout_rh.close()
        self.shep_stdout_wh.close()

        try:
            self.shep_input_wh.close()
        except:
            pass

        try:
            self.bela_input_rh.close()
        except:
            pass

        try:
            self.gs_input_wh.close()
        except:
            pass

        try:
            self.proc_gs_return_rh.close()
        except:
            pass

        try:
            self.proc_shep_return_rh.close()
        except:
            pass

        try:
            self.shep_input_chan.detach()
        except:
            pass

        try:
            self.bela_input_chan.detach()
        except:
            pass

        try:
            self.gs_input_chan.detach()
        except:
            pass

        try:
            self.proc_gs_return_chan.destroy()
        except:
            pass

        try:
            self.proc_shep_return_chan.destroy()
        except:
            pass

        try:
            self.inf_pool.detach()
        except:
            pass

        try:
            self.test_pool.destroy()
        except:
            pass

        try:
            self.shep_dut.join(timeout=1)
        except TimeoutError:
            self.shep_dut.kill()

    def start_thing(self, start_args):
        res = dproc.create(exe='python3', run_dir='',
                           args=start_args, env={},
                           options=pdesc.ProcessOptions(make_inf_channels=True))

        self.head_puid = res.p_uid

    def monitor_output(self, optional_barrier=None):
        while True:
            msg = dmsg.parse(self.bela_input_rh.recv())
            if isinstance(msg, dmsg.SHFwdOutput):
                if msg.fd_num == msg.FDNum.STDOUT.value:
                    sys.stdout.write(msg.data)
                    sys.stdout.flush()
                elif msg.fd_num == msg.FDNum.STDERR.value:
                    sys.stderr.write(msg.data)
                    sys.stderr.flush()
                if optional_barrier is not None:
                    try:
                        optional_barrier.wait()
                    except threading.BrokenBarrierError:
                        pass
            elif isinstance(msg, dmsg.GSHeadExit):
                print('+++ head proc exited, code {}'.format(msg.exit_code))
                break
            else:
                print('unexpected message: {}'.format(msg))


default_parent_stdin_filename = f'quokka_{os.getpid()}'
parent_raw_input_fh = None
coordinate_repl_prompt_barrier = threading.Barrier(2)


def get_raw_input_from_parent_via_fd(prompt='D>>'):
    'Intended to be called as substitute for raw_input() in child process.'
    global parent_raw_input_fh, default_parent_stdin_filename
    if parent_raw_input_fh is None:
        print(f'DBG {default_parent_stdin_filename=}', flush=True)
        parent_raw_input_fh = open(default_parent_stdin_filename, 'rt')
    sys.stdout.flush()
    input_text = parent_raw_input_fh.readline()
    return input_text.rstrip()


def capture_raw_input_as_parent_via_fd(
        filename=default_parent_stdin_filename,
        barrier=coordinate_repl_prompt_barrier
):
    os.mkfifo(filename)
    barrier.wait()
    print(f'DBG parent process {os.getpid()=}')
    try:
        with open(filename, 'wt') as parent_raw_input_fh:
            while True:
                print('D>> ', end='', flush=True)
                input_text = sys.stdin.readline()
                parent_raw_input_fh.write(input_text)
                parent_raw_input_fh.flush()
                try:
                    barrier.wait(timeout=1)
                except threading.BrokenBarrierError:
                    pass
    except Exception as e:
        print(f'Capture ending, reason: {e}', flush=True)


def start_repl(pipe_filename):
    import code
    g = globals()
    g['default_parent_stdin_filename'] = pipe_filename
    code.interact(readfunc=get_raw_input_from_parent_via_fd, local=g)


def main(start_args, interactive=False):
    the_ctx = Context()
    try:
        the_ctx.start_infrastructure()
    except Exception as e:
        print('Got exception {} in inf startup'.format(e))
        traceback.print_exc()
        the_ctx.cleanup()
        exit(1)

    if interactive:
        capture_raw_input_thread = threading.Thread(
            target=capture_raw_input_as_parent_via_fd,
            args=(default_parent_stdin_filename, coordinate_repl_prompt_barrier)
        )
        capture_raw_input_thread.start()

    try:
        the_ctx.start_thing(start_args)
    except Exception as e:
        print('Got exception {} in program start'.format(e))
        traceback.print_exc()
        the_ctx.cleanup()
        exit(1)

    # try this to handle ctl-C while stuff runs.
    atexit.register(the_ctx.cleanup)

    if interactive:
        monitor_output_thread = threading.Thread(
            target=Context.monitor_output,
            args=(the_ctx, coordinate_repl_prompt_barrier)
        )
        monitor_output_thread.start()
        monitor_output_thread.join()
    else:
        the_ctx.monitor_output()

    atexit.unregister(the_ctx.cleanup)
    the_ctx.cleanup()

    try:
        # Can not be done in start_repl() because still in use at that time.
        os.unlink(default_parent_stdin_filename)
    except:
        pass

    exit(0)


if __name__ == "__main__":
    if 1 == len(sys.argv):
        args = [
            '-c',
            f'import start_single; start_single.start_repl({default_parent_stdin_filename!r})'
        ]
        interactive = True
    else:
        args = sys.argv[1:]
        interactive = False

    main(args, interactive)

test/integration/start_single_ls.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import logging
import json

from unittest.mock import MagicMock, patch

from dragon.channels import Channel, ChannelError
from dragon.infrastructure import facts as dfacts
from dragon.infrastructure import messages as dmsg
from dragon.infrastructure import process_desc
from dragon.infrastructure.connection import Connection, ConnectionOptions
from dragon.infrastructure.parameters import this_process, POLICY_INFRASTRUCTURE
from dragon.infrastructure.messages import AbnormalTerminationError
from dragon.launcher.util import next_tag, get_with_blocking
from dragon.launcher.network_config import NetworkConfig
from dragon.localservices.local_svc import mk_inf_resources
from dragon.managed_memory import MemoryPool, DragonPoolError, DragonMemoryError
from dragon.utils import B64


def get_args_map(network_config, frontend_sdesc, host_id, ip_addrs):
    network_prefix = "^(eth|hsn)"

    args_map = {
        "ip_addrs": ip_addrs,  # it includes the port
        "host_ids": host_id,
        "frontend_sdesc": frontend_sdesc,
        "network_prefix": f"{network_prefix}",
        "transport_test": False
    }
    return args_map


def mock_start_localservices_wrapper(func):
    def wrapper(*args, **kwargs):
        def mocked_start_localservices(*args, **kwargs):
            return MagicMock(), stdin_w_fd, stdout_r_fd

        from dragon.infrastructure.util import NewlineStreamWrapper

        stdin_r_fd, stdin_w_fd = os.pipe()
        stdout_r_fd, stdout_w_fd = os.pipe()

        ls_stdin_r = NewlineStreamWrapper(
            os.fdopen(stdin_r_fd, "rb", buffering=0), read_intent=True, write_intent=False
        )
        ls_stdout_w = NewlineStreamWrapper(
            os.fdopen(stdout_w_fd, "wb", buffering=0), read_intent=False, write_intent=True
        )

        with patch(
            "dragon.launcher.backend.LauncherBackEnd._start_localservices",
            side_effect=mocked_start_localservices
        ) as mock_localservices:
            localservices_arg_tuple = (mock_localservices, ls_stdin_r, ls_stdout_w)
            args += (localservices_arg_tuple,)
            result = func(*args, **kwargs)

        ls_stdout_w.close()
        del ls_stdout_w

        ls_stdin_r.close()
        del ls_stdin_r

        return result

    return wrapper


class LauncherBackendHelper:

    def __init__(self,
                 network_config,
                 node_index,
                 ip_addrs,
                 fe_host_id,
                 be_host_id):
        self.network_config = network_config
        self.node_index = node_index
        self.ip_addrs = ip_addrs
        self.fe_host_id = fe_host_id
        self.be_host_id = be_host_id

        self.garbage_dict = {'junk': 'abdecdf', '52': 20}

    def cleanup(self):

        try:
            del self.localservices
        except Exception:
            pass

        try:
            del self.be_overlay_tsta
        except Exception:
            pass

        try:
            self.cleanup_ls()
        except Exception:
            pass

        try:
            del self.launcher_fe
        except Exception:
            pass


    def handle_overlay_start(self, mock_overlay, mock_start_localservices):
        '''Start of mocked up overlay network'''

        log = logging.getLogger('handle_overlay_start')
        log.info("Waiting for overlay call...")
        while mock_overlay.call_args is None:
            pass

        log.info("Got overlay call...")
        overlay_args = mock_overlay.call_args.kwargs

        # Create our backend overlay TCP agent and connect to the
        # frontend's channels
        self.be_overlay_tsta = BackendOverlay(
            overlay_args["ch_in_sdesc"], overlay_args["ch_out_sdesc"]
        )

        # Let backend know its overlay agent is "up"
        self.be_overlay_tsta.send_OverlayPingBE()  # M7

        # Establish 2-way communication with backend launcher
        self.launcher_fe.recv_BEIsUp()  # M8
        self.launcher_fe.connect_to_be()
        self.launcher_fe.send_FENodeIdxBE()  # M9

        # Get generated LS stdin/stdout file descriptors. Turn them into newline stream wrappers
        _, self.ls_stdin_queue, self.ls_stdout_queue = mock_start_localservices

    def handle_network_and_frontend_start(self, mock_network_config):
        '''Define the mocked up network and frontend'''
        net = NetworkConfig.from_file(self.network_config)
        mock_network_config.return_value = net.get_network_config()[str(self.node_index)]

        self.launcher_fe = LauncherFrontEnd(self.node_index, self.fe_host_id)

        args_map = get_args_map(
            self.network_config,
            self.launcher_fe.encoded_inbound,
            self.ip_addrs,
            self.be_host_id,
        )

        return args_map

    def start_ls(self,
                 mock_overlay,
                 mock_start_localservices,
                 garble_shchannelsup=False,
                 garble_lachannelsinfo=False,
                 abort_lachannelsinfo=False):

        self.handle_overlay_start(mock_overlay, mock_start_localservices)

        # create our localservices standin class and connect to stdin/stdout queues
        self.localservices = LocalServices(
            self.ls_stdin_queue, self.ls_stdout_queue, self.node_index, self.be_host_id
        )
        self.localservices.recv_BENodeIdxSH()  # M10
        self.localservices.send_SHPingBE()  # M11

        self.localservices.recv_BEPingSH()  # M12

        if garble_shchannelsup:
            self.localservices.send_SHChannelsUP(custom_msg=json.dumps(self.garbage_dict))
            return
        else:
            self.localservices.send_SHChannelsUP()  # M13

        # Launcher Backend spawns overlay TCP agent
        # Wait for BE Launcher to initialize its Overlay network
        # We need to grab the parameters passed to the overlay mock
        # FE -> BE Send LAChannelsInfo
        self.launcher_fe.recv_SHChannelsUp()  # M14

        if garble_lachannelsinfo:
            self.launcher_fe.send_LAChannelsInfo(custom_msg=json.dumps(self.garbage_dict))
            return
        elif abort_lachannelsinfo:
            self.launcher_fe.send_LAChannelsInfo(custom_msg=dmsg.AbnormalTermination(tag=next_tag()).serialize())
            return
        else:
            self.launcher_fe.send_LAChannelsInfo()  # M15

        self.localservices.recv_LAChannelsInfo()  # M16

    def clean_startup(self, log, mock_overlay, mock_network_config, mock_start_localservices):
        '''Execute a clean bringup and exit of all backend services'''

        self.start_ls(mock_overlay, mock_start_localservices)

        # A11 Local Services spawns TA

        self.localservices.send_TAUP()  # M20
        self.launcher_fe.recv_TAUP()  # M21

        # Local Services spawns Global Services
        self.globalservices = self.localservices.launch_globalservices()  # A13
        self.globalservices.send_GSIsUp()  # M25
        self.launcher_fe.recv_GSIsUp()  # M26

    def launch_user_process(self, log):
        p_uid = 23941
        self.launcher_fe.send_GSProcessCreate()
        create_process_tag = self.globalservices.recv_GSProcessCreate()

        self.globalservices.send_GSProcessCreateResponse(
            p_uid=p_uid, create_process_tag=create_process_tag
        )
        self.launcher_fe.recv_GSProcessCreateResponse()

        self.localservices.send_SHFwdOutput_StdOut(p_uid=p_uid)
        self.localservices.send_SHFwdOutput_StdErr(p_uid=p_uid)

        self.launcher_fe.recv_SHFwdOutput()
        self.launcher_fe.recv_SHFwdOutput()

    def start_gsteardown_exit(self,
                              shteardown_abort=False):

        self.launcher_fe.send_GSTeardown()  # M5
        self.globalservices.recv_GSTeardown()  # M6

        del self.globalservices
        self.globalservices = None

        self.localservices.send_GSHalted()  # M8
        self.launcher_fe.recv_GSHalted()  # M9

        self.launcher_fe.send_SHHaltTA()  # M10
        self.localservices.recv_SHHaltTA()  # M11

        self.localservices.send_TAHalted()  # M14
        self.launcher_fe.recv_TAHalted()  # M15

        if shteardown_abort:
            self.localservices.send_AbnormalTermination()

        self.launcher_fe.send_SHTeardown()  # M16
        self.localservices.recv_SHTeardown()  # M17

        self.localservices.send_SHHaltBE()  # M18
        self.launcher_fe.recv_SHHaltBE()  # M19

        self.launcher_fe.send_BEHalted()  # M20
        self.localservices.recv_BEHalted()  # M21

        self.be_overlay_tsta.recv_BEHaltOverlay()  # M22
        self.be_overlay_tsta.send_OverlayHalted()  # M23

    def start_gsteardown_exit_no_ls_no_gs(self,
                                          log,
                                          shteardown_abort=False):

        self.launcher_fe.send_GSTeardown()  # M5
        log.debug('sent gsteardown')

        self.launcher_fe.recv_GSHalted()  # M9
        log.debug('recved gshalted')

        self.launcher_fe.send_SHHaltTA()  # M10
        log.debug('send shhaltta')
        self.launcher_fe.recv_TAHalted()  # M15

        if shteardown_abort:
            self.launcher_fe.send_AbnormalTermination()
            return
        else:
            self.launcher_fe.send_SHTeardown()  # M16
        self.localservices.recv_SHTeardown()  # M17

        self.localservices.send_SHHaltBE()  # M18
        self.launcher_fe.recv_SHHaltBE()  # M19

        self.launcher_fe.send_BEHalted()  # M20
        self.localservices.recv_BEHalted()  # M21

        self.be_overlay_tsta.recv_BEHaltOverlay()  # M22
        self.be_overlay_tsta.send_OverlayHalted()  # M23

    def start_gsteardown_exit_no_gs(self,
                                    log,
                                    shteardown_abort=False):

        self.launcher_fe.send_GSTeardown()  # M5
        log.debug('sent gsteardown')

        self.launcher_fe.recv_GSHalted()  # M9
        log.debug('recved gshalted')

        self.launcher_fe.send_SHHaltTA()  # M10
        log.debug('send shhaltta')
        self.localservices.recv_SHHaltTA()  # M11

        self.localservices.send_TAHalted()  # M14
        self.launcher_fe.recv_TAHalted()  # M15

        if shteardown_abort:
            self.localservices.send_AbnormalTermination()
            return
        else:
            self.launcher_fe.send_SHTeardown()  # M16
        self.localservices.recv_SHTeardown()  # M17

        self.localservices.send_SHHaltBE()  # M18
        self.launcher_fe.recv_SHHaltBE()  # M19

        self.launcher_fe.send_BEHalted()  # M20
        self.localservices.recv_BEHalted()  # M21

        self.be_overlay_tsta.recv_BEHaltOverlay()  # M22
        self.be_overlay_tsta.send_OverlayHalted()  # M23

    def start_abnormal_termination_exit(self,
                                        log,
                                        gs_up=True,
                                        ls_up=True,
                                        shteardown_abort=False):

        try:
            self.launcher_fe.recv_AbnormalTermination()
        except AbnormalTerminationError:
            pass

        if not gs_up and not ls_up:
            self.start_gsteardown_exit_no_ls_no_gs(log,
                                                   shteardown_abort=shteardown_abort)
        elif not gs_up:
            self.start_gsteardown_exit_no_gs(log,
                                             shteardown_abort=shteardown_abort)
        else:
            self.start_gsteardown_exit(shteardown_abort=shteardown_abort)

    def clean_shutdown(self,
                       log,
                       shteardown_abort=False):

        self.globalservices.send_GSHeadExit()  # M3
        self.launcher_fe.recv_GSHeadExit()  # M4
        self.start_gsteardown_exit(shteardown_abort=shteardown_abort)

    def cleanup_ls(self):

        self.ls_stdin_queue.stream.close()
        del self.ls_stdin_queue

        self.ls_stdout_queue.stream.close()
        del self.ls_stdout_queue


class LauncherFrontEnd:
    def __init__(self, node_index, host_id):
        self.log = logging.getLogger("_launcher_fe")
        self.node_index = node_index
        self.host_id = host_id

        self.port = dfacts.DEFAULT_TRANSPORT_PORT
        self.num_gw_channels = dfacts.DRAGON_OVERLAY_DEFAULT_NUM_GW_CHANNELS_PER_NODE

        self.create_frontend_comms()

    def __del__(self):
        self.log.debug("__del__")

        try:
            self.conn_out.close()
            del self.conn_out
        except Exception:
            pass

        try:
            self.be_ch.detach()
            del self.be_ch
        except Exception:
            pass

        try:
            self.conn_in.close()
            del self.conn_in
        except Exception:
            pass

        try:
            self.fe_inbound.destroy()
            del self.fe_inbound
        except Exception:
            pass

        try:
            self.fe_mpool.destroy()
            del self.fe_mpool
        except Exception:
            pass

    def create_frontend_comms(self):
        self.log.debug("creating frontend memory pool and channels to talk to backend")
        conn_options = ConnectionOptions(min_block_size=2**16)

        self.fe_mpool = None
        self.fe_inbound = None
        self.conn_in = None

        try:
            # Create my memory pool
            self.fe_mpool = MemoryPool(
                                int(dfacts.DEFAULT_SINGLE_DEF_SEG_SZ),
                                f"{os.getuid()}_{os.getpid()}_{self.host_id}_{dfacts.DEFAULT_POOL_SUFFIX}",
                                dfacts.FE_OVERLAY_TRANSPORT_AGENT_MUID,
                                )

            # Create my receiving channel
            self.fe_cuid = dfacts.FE_CUID

            # Channel for launcher backend to talk to me
            self.fe_inbound = Channel(self.fe_mpool, self.fe_cuid)
            self.encoded_inbound = B64(self.fe_inbound.serialize())

            self.conn_in = Connection(
                inbound_initializer=self.fe_inbound,
                options=conn_options,
                policy=POLICY_INFRASTRUCTURE,
            )
        except (
            ChannelError,
            DragonPoolError,
            DragonMemoryError,
        ) as init_err:
            if self.conn_in:
                del self.conn_in
                self.conn_in = None

            if self.fe_inbound:
                del self.fe_inbound
                self.fe_inbound = None

            if self.fe_mpool:
                self.fe_mpool.destroy()

            raise RuntimeError(
                "Overlay transport resource creation failed"
            ) from init_err

    def recv_BEIsUp(self):
        self.log.debug("waiting for BEIsUp from Launcher BE")
        self.be_is_up = get_with_blocking(self.conn_in)
        assert isinstance(self.be_is_up, dmsg.BEIsUp), "expected BSIsUp"
        self.log.debug(f"recv_BEIsUp got {self.be_is_up=}")

    def connect_to_be(self):
        self.log.debug("connecting to be")
        self.be_sdesc = B64.from_str(self.be_is_up.be_ch_desc)
        self.be_ch = Channel.attach(self.be_sdesc.decode(), mem_pool=self.fe_mpool)
        be_conn_options = ConnectionOptions(
            default_pool=self.fe_mpool, min_block_size=2**16
        )
        self.conn_out = Connection(
            outbound_initializer=self.be_ch,
            options=be_conn_options,
            policy=POLICY_INFRASTRUCTURE,
        )
        self.log.debug("connected to be")

    def send_FENodeIdxBE(self):
        fe_node_idx = dmsg.FENodeIdxBE(tag=next_tag(), node_index=self.node_index)
        self.conn_out.send(fe_node_idx.serialize())
        self.log.debug(f"send_FENodeIdxBE sent {fe_node_idx=}")

    def recv_SHChannelsUp(self):
        nnodes = 1
        self.log.debug("receiving SHChannelsUp messages")
        self.chs_up = [get_with_blocking(self.conn_in) for _ in range(nnodes)]
        self.log.debug(f"{self.chs_up=}")
        for ch_up in self.chs_up:
            assert isinstance(
                ch_up, dmsg.SHChannelsUp
            ), "la_fe received invalid channel up"

        self.log.debug("Looking for gs_cd")
        gs_cds = [ch_up.gs_cd for ch_up in self.chs_up if ch_up.gs_cd is not None]
        assert len(gs_cds) > 0, "could not find a gs_cd response"
        assert len(gs_cds) == 1, "more than one gs_cd response found"
        if len(gs_cds) == 0:
            msg = "The Global Services CD was not returned by any of the SHChannelsUp messages. Launcher Exiting."
            raise RuntimeError(msg)
        self.gs_cd = gs_cds[0]

        self.log.debug(f"recv_SHChannelsUp got {self.chs_up} with {self.gs_cd=}")

    def send_LAChannelsInfo(self, custom_msg: str = None):

        if custom_msg is None:
            la_ch_info = dmsg.LAChannelsInfo(tag=next_tag(),
                                             nodes_desc=self.chs_up,
                                             gs_cd=self.gs_cd,
                                             num_gw_channels=self.num_gw_channels,
                                             port=self.port,
                                             )
            # self.la_fe_stdout.send("A", la_ch_info.serialize())
            self.conn_out.send(la_ch_info.serialize())

            self.log.debug(f"send_LAChannelsInfo sent {la_ch_info=}")
            self.log.debug(f"{la_ch_info.nodes_desc=}")
            self.log.debug(f"{la_ch_info.gs_cd=}")
        else:
            self.conn_out.send(custom_msg)

    def recv_TAUP(self):
        ta_up = get_with_blocking(self.conn_in)
        assert isinstance(
            ta_up, dmsg.TAUp
        ), "la_fe received invalid channel up"
        self.log.debug(f"recv_TAUP got {ta_up=}")

    def recv_GSIsUp(self):
        gs_is_up = get_with_blocking(self.conn_in)
        assert isinstance(gs_is_up, dmsg.GSIsUp), "expected GSIsUP"
        self.log.debug(f"recv_GSIsUp got {gs_is_up=}")

    def send_GSProcessCreate(self):
        exe = "hello_world.py"
        args = ["--language", "espanol"]
        options = process_desc.ProcessOptions(make_inf_channels=True)
        gs_process_create = dmsg.GSProcessCreate(
            tag=next_tag(),
            p_uid=dfacts.LAUNCHER_PUID,
            r_c_uid=dfacts.BASE_BE_CUID,
            exe=exe,
            args=args,
            options=options,
        )
        self.conn_out.send(gs_process_create.serialize())
        self.log.debug(f"send_GSProcessCreate sent {gs_process_create=}")

    def recv_GSProcessCreateResponse(self):
        gs_process_create_response = get_with_blocking(self.conn_in)
        assert isinstance(
            gs_process_create_response, dmsg.GSProcessCreateResponse
        ), "expected GSIsUP"
        self.log.debug(f"recv_GSProcessCreateResponse got {gs_process_create_response=}")

    def recv_SHFwdOutput(self):
        sh_fwd_output = get_with_blocking(self.conn_in)
        assert isinstance(sh_fwd_output, dmsg.SHFwdOutput), "expected SHFwdOutput"
        self.log.debug(f"recv_SHFwdOutput got {sh_fwd_output=}")

    def recv_GSHeadExit(self):  # M4
        gs_head_exit = get_with_blocking(self.conn_in)
        assert isinstance(gs_head_exit, dmsg.GSHeadExit), "expected GSHeadExit"
        self.log.debug(f"recv_GSHeadExit got {gs_head_exit=}")

    def send_GSTeardown(self):  # M5
        gs_teardown = dmsg.GSTeardown(tag=next_tag())
        self.conn_out.send(gs_teardown.serialize())
        self.log.debug(f"send_GSTeardown sent {gs_teardown=}")

    def recv_GSHalted(self):  # M9
        gs_halted = get_with_blocking(self.conn_in)
        assert isinstance(gs_halted, dmsg.GSHalted), "expected GSHalted"
        self.log.debug(f"recv_GSHalted got {gs_halted=}")

    def send_SHHaltTA(self):  # M10
        sh_halt_ta = dmsg.SHHaltTA(tag=next_tag())
        self.conn_out.send(sh_halt_ta.serialize())
        self.log.debug(f"send_SHHaltTA sent {sh_halt_ta=}")

    def recv_TAHalted(self):  # M15
        ta_halted = get_with_blocking(self.conn_in)
        assert isinstance(ta_halted, dmsg.TAHalted), "expected TAHalted"
        self.log.debug(f"recv_TAHalted got {ta_halted=}")

    def send_SHTeardown(self):  # M16
        sh_teardown = dmsg.SHTeardown(tag=next_tag())
        self.conn_out.send(sh_teardown.serialize())
        self.log.debug(f"send_SHTeardown sent {sh_teardown=}")

    def recv_SHHaltBE(self):  # M19
        sh_halt_be = get_with_blocking(self.conn_in)
        assert isinstance(sh_halt_be, dmsg.SHHaltBE), "expected sh_halt_be"
        self.log.debug(f"recv_SHHaltBE got {sh_halt_be=}")

    def send_BEHalted(self):  # M20
        be_halted = dmsg.BEHalted(tag=next_tag())
        self.conn_out.send(be_halted.serialize())
        self.log.debug(f"send_BEHalted sent {be_halted=}")

    def recv_AbnormalTermination(self):  # M5
        abnormal_term = get_with_blocking(self.conn_in)
        self.log.debug(f"recv_AbnormalTermination got {abnormal_term=}")


class BackendOverlay:
    def __init__(self, ch_in_desc: B64, ch_out_desc: B64):
        self.log = logging.getLogger("_be_overlay_tsta")

        try:
            self.ta_ch_in = Channel.attach(ch_in_desc.decode())
            self.ta_ch_out = Channel.attach(ch_out_desc.decode())
            self.be_ta_conn = Connection(
                inbound_initializer=self.ta_ch_in,
                outbound_initializer=self.ta_ch_out,
                options=ConnectionOptions(
                    creation_policy=ConnectionOptions.CreationPolicy.EXTERNALLY_MANAGED
                ),
            )
            self.be_ta_conn.ghost = True
            self.be_ta_conn.open()

        except (ChannelError, DragonPoolError, DragonMemoryError) as init_err:
            raise RuntimeError(
                "Overlay transport resource creation failed"
            ) from init_err

    def __del__(self):
        self.log.debug("__del__")

        try:
            self.be_ta_conn.close()
            del self.be_ta_conn
        except Exception:
            pass

        try:
            self.ta_ch_out.detach()
            del self.ta_ch_out
        except Exception:
            pass

        try:
            self.ta_ch_in.detach()
            del self.ta_ch_in
        except Exception:
            pass

    def send_OverlayPingBE(self):
        overlay_ping_be = dmsg.OverlayPingBE(next_tag())
        self.be_ta_conn.send(overlay_ping_be.serialize())
        self.log.debug(f"send_OverlayPingBE sent {overlay_ping_be=}")

    def recv_BEHaltOverlay(self):  # M22
        be_halt_overlay = get_with_blocking(self.be_ta_conn)
        assert isinstance(be_halt_overlay, dmsg.BEHaltOverlay), "expected BEHaltOverlay"
        self.log.debug(f"recv_SHHaltTA got {be_halt_overlay=}")

    def send_OverlayHalted(self):  # M23
        overlay_halted = dmsg.OverlayHalted(next_tag())
        self.be_ta_conn.send(overlay_halted.serialize())
        self.log.debug(f"send_OverlayHalted sent {overlay_halted}")


class LocalServices:
    def __init__(self, ls_stdin_queue, ls_stdout_queue, node_index, host_id):
        self.log = logging.getLogger("_localservices")
        self.ls_stdin_queue = ls_stdin_queue
        self.ls_stdout_queue = ls_stdout_queue
        self.node_index = node_index
        self.host_id = host_id

        # Make LS Infrastructure channels
        (
            self.start_pools,
            self.start_channels,
            self.ls_input,
            self.la_input,
            self.ta_input_descr,
            self.ta_input,
            self.gs_input,
        ) = mk_inf_resources(node_index)

        self.gs_cd = this_process.gs_cd

    def __del__(self):
        self.log.debug("__del__")

        try:
            self.log.debug("gs_input destroy")
            self.gs_input.destroy()
            del self.gs_input
        except Exception:
            pass

        try:
            self.log.debug("ta_input close")
            self.ta_input.close()
            del self.ta_input
        except Exception:
            pass

        try:
            self.log.debug("la_input close")
            self.la_input.close()
            del self.la_input
        except Exception:
            pass

        try:
            self.log.debug("ls_input close")
            self.ls_input.close()
            del self.ls_input
        except Exception:
            pass

        try:
            self.log.debug("start_channels destroy")
            for key, chan in self.start_channels.items():
                try:
                    self.log.debug(f'destroy channel {key}')
                    chan.destroy()
                except Exception:
                    pass
            del self.start_channels
        except Exception:
            pass

        try:
            self.log.debug("start_pools destroy")
            for pool in self.start_pools.values():
                try:
                    pool.destroy()
                except Exception:
                    pass
            del self.start_pools
        except AttributeError:
            pass

    def send_SHPingBE(self):
        this_process.index = self.node_index
        be_ping = dmsg.SHPingBE(
            tag=next_tag(),
            shep_cd=this_process.local_shep_cd,
            be_cd=this_process.local_be_cd,
            gs_cd=this_process.gs_cd,
            default_pd=this_process.default_pd,
            inf_pd=this_process.inf_pd,
        )
        self.ls_stdout_queue.send(be_ping.serialize())
        self.log.debug(f"send_SHPingBE sent {be_ping=}")

    def recv_BENodeIdxSH(self):
        self.be_node_idx_sh = dmsg.parse(self.ls_stdin_queue.recv())
        assert isinstance(self.be_node_idx_sh, dmsg.BENodeIdxSH), "expected BENodeIdxSH"
        self.log.debug(f"recv_BENodeIdxSH got {self.be_node_idx_sh=}")

    def send_SHChannelsUP(self, custom_msg=None):

        if custom_msg:
            self.la_input.send(custom_msg)
        else:
            ch_up_msg = dmsg.SHChannelsUp(tag=next_tag(),
                                          host_name=self.be_node_idx_sh.host_name,
                                          host_id=self.host_id,
                                          ip_addrs=self.be_node_idx_sh.ip_addrs,
                                          shep_cd=this_process.local_shep_cd,
                                          gs_cd=self.gs_cd,
                                          idx=self.be_node_idx_sh.node_idx)

            self.la_input.send(ch_up_msg.serialize())
            self.log.debug(f"send_SHChannelsUP sent {ch_up_msg=}")

    def recv_BEPingSH(self):
        be_ping_sh = dmsg.parse(self.ls_input.recv())
        assert isinstance(be_ping_sh, dmsg.BEPingSH), "expected BSPingSH"
        self.log.debug(f"recv_BEPingSH got {be_ping_sh=}")

    def recv_LAChannelsInfo(self):
        # Recv LAChannelsInfo Broadcast
        la_channels_info = dmsg.parse(self.ls_input.recv())
        assert isinstance(
            la_channels_info, dmsg.LAChannelsInfo
        ), "expected LAChannelsInfo"
        self.gs_cd = la_channels_info.gs_cd
        self.log.debug(f"recv_LAChannelsInfo got {la_channels_info=}")
        self.log.debug(f"{la_channels_info.nodes_desc=}")
        self.log.debug(f"{la_channels_info.gs_cd=}")

    def send_TAUP(self):
        ch_list = []
        ta_up = dmsg.TAUp(
            tag=next_tag(), idx=self.node_index, test_channels=ch_list
        )
        self.la_input.send(ta_up.serialize())
        self.log.debug(f"send_TAUP sent {ta_up=}")

    def launch_globalservices(self):
        return GlobalServices(self.gs_cd, self.la_input)

    def send_SHFwdOutput_StdOut(self, p_uid):
        sh_fwd_output = dmsg.SHFwdOutput(
            tag=next_tag(),
            p_uid=p_uid,
            idx=self.node_index,
            fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
            data="Hola! This is STDOUT",
        )
        self.la_input.send(sh_fwd_output.serialize())
        self.log.debug(f"send_SHFwdOutput_StdErr sent via STDOUT {sh_fwd_output=}")

    def send_SHFwdOutput_StdErr(self, p_uid):
        sh_fwd_output = dmsg.SHFwdOutput(
            tag=next_tag(),
            p_uid=p_uid,
            idx=self.node_index,
            fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
            data="Hola! This is STDERR",
        )
        self.la_input.send(sh_fwd_output.serialize())
        self.log.debug(f"send_SHFwdOutput_StdErr sent via STDERR {sh_fwd_output=}")

    def send_GSHalted(self):  # M8
        gs_halted = dmsg.GSHalted(tag=next_tag())
        self.la_input.send(gs_halted.serialize())
        self.log.debug(f"send_GSHalted sent {gs_halted=}")

    def send_AbnormalTermination(self):
        abnormal_term = dmsg.AbnormalTermination(tag=next_tag())
        self.la_input.send(abnormal_term.serialize())
        self.log.debug(f"send_AbnormalTermination sent {abnormal_term=}")

    def recv_SHHaltTA(self):  # M11
        sh_halt_ta = dmsg.parse(self.ls_input.recv())
        assert isinstance(sh_halt_ta, dmsg.SHHaltTA), "expected SHHaltTA"
        self.log.debug(f"recv_SHHaltTA got {sh_halt_ta=}")

    def send_TAHalted(self):  # M14
        ta_halted = dmsg.TAHalted(tag=next_tag())
        self.la_input.send(ta_halted.serialize())
        self.log.debug(f"send_TAHalted sent {ta_halted=}")

    def recv_SHTeardown(self):  # M17
        sh_teardown = dmsg.parse(self.ls_input.recv())
        assert isinstance(sh_teardown, dmsg.SHTeardown), "expected SHTeardown"
        self.log.debug(f"recv_SHTeardown got {sh_teardown=}")

    def send_SHHaltBE(self):  # M18
        sh_halt_be = dmsg.SHHaltBE(tag=next_tag())
        self.la_input.send(sh_halt_be.serialize())
        self.log.debug(f"send_SHHaltBE sent {sh_halt_be=}")

    def recv_BEHalted(self):
        be_halted = dmsg.parse(self.ls_stdin_queue.recv())
        assert isinstance(be_halted, dmsg.BEHalted), "expected BEHalted"
        self.log.debug(f"recv_BEHalted got {be_halted=}")


class GlobalServices:
    def __init__(self, gs_cd, la_input):
        self.log = logging.getLogger("_globalservices")
        self.gs_cd = gs_cd
        self.la_input = la_input
        self.la_input.ghost = True

        self.gs_ch_in = Channel.attach(B64.from_str(gs_cd).decode())
        self.gs_queue = Connection(
            inbound_initializer=self.gs_ch_in, policy=POLICY_INFRASTRUCTURE
        )

    def __del__(self):
        self.log.debug("__del__")

        try:
            self.gs_queue.close()
            del self.gs_queue
        except Exception:
            pass

        try:
            self.gs_ch_in.detach()
            del self.gs_ch_in
        except Exception:
            pass

        del self.la_input

    def send_GSIsUp(self):
        gs_is_up = dmsg.GSIsUp(tag=next_tag())
        self.la_input.send(gs_is_up.serialize())
        self.log.debug(f"send_GSIsUp sent {gs_is_up=}")

    def recv_GSProcessCreate(self):
        gs_process_create = get_with_blocking(self.gs_queue)
        assert isinstance(
            gs_process_create, dmsg.GSProcessCreate
        ), "expected GSProcessCreate"
        self.log.debug(f"recv_GSProcessCreate got {gs_process_create=}")
        return gs_process_create.tag

    def send_GSProcessCreateResponse(self, p_uid, create_process_tag):
        gs_process_create_response = dmsg.GSProcessCreateResponse(
            tag=next_tag(),
            ref=create_process_tag,
            err=dmsg.GSProcessCreateResponse.Errors.SUCCESS,
            desc=process_desc.ProcessDescriptor(
                p_uid=p_uid, p_p_uid=47622, name="hello_world", node=0
            ),
        )
        self.la_input.send(gs_process_create_response.serialize())
        self.log.debug(f"send_GSProcessCreateResponse sent {gs_process_create_response=}")

    def send_GSHeadExit(self):
        gs_head_exit = dmsg.GSHeadExit(tag=next_tag(), exit_code=0)
        self.la_input.send(gs_head_exit.serialize())
        self.log.debug(f"send_GSHeadExit sent {gs_head_exit=}")

    def recv_GSTeardown(self):
        gs_teardown = get_with_blocking(self.gs_queue)
        assert isinstance(gs_teardown, dmsg.GSTeardown), "expected GSTeardown"
        self.log.debug(f"recv_GSTeardown got {gs_teardown=}")

test/launcher/backend_testing_mocks.py
DragonHPC/dragon
import os
import logging

from dragon.launcher.frontend import LauncherFrontEnd
from dragon.launcher.util import next_tag
from dragon.launcher.network_config import NetworkConfig

from dragon.infrastructure.process_desc import ProcessDescriptor
from dragon.infrastructure.connection import Connection, ConnectionOptions
from dragon.infrastructure.parameters import POLICY_INFRASTRUCTURE
from dragon.infrastructure import facts as dfacts
from dragon.infrastructure import messages as dmsg

from dragon.channels import Channel, ChannelError
from dragon.managed_memory import MemoryPool, DragonPoolError, DragonMemoryError
from dragon.utils import set_host_id, B64
from dragon.dlogging.util import setup_FE_logging


def run_frontend(args_map):

    # Before doing anything set my host ID
    set_host_id(dfacts.FRONTEND_HOSTID)
    setup_FE_logging(log_device_level_map=args_map['log_device_level_map'],
                     basename='dragon', basedir=os.getcwd())
    log = logging.getLogger('run_frontend')
    for key, value in args_map.items():
        if value is not None:
            log.info(f'args_map: {key}: {value}')

    with LauncherFrontEnd(args_map=args_map) as fe_server:
        fe_server.run_startup()
        fe_server.run_app()
        fe_server.run_msg_server()


def open_overlay_comms(ch_in_desc: B64,
                       ch_out_desc: B64):
    '''Attach to Frontend's overlay network channels and open a connection'''
    try:
        ta_ch_in = Channel.attach(ch_in_desc.decode())
        ta_ch_out = Channel.attach(ch_out_desc.decode())
        fe_ta_conn = Connection(inbound_initializer=ta_ch_in,
                                outbound_initializer=ta_ch_out,
                                options=ConnectionOptions(creation_policy=ConnectionOptions.CreationPolicy.EXTERNALLY_MANAGED))
        fe_ta_conn.ghost = True
        fe_ta_conn.open()

    except (ChannelError, DragonPoolError, DragonMemoryError) as init_err:
        raise RuntimeError('Overlay transport resource creation failed') from init_err

    return ta_ch_in, ta_ch_out, fe_ta_conn


def open_backend_comms(frontend_sdesc: str,
                       network_config: str):
    '''Attach to frontend channel and create channels for backend'''
    be_mpool = None
    try:
        conn_options = ConnectionOptions(min_block_size=2 ** 16)
        conn_policy = POLICY_INFRASTRUCTURE

        be_mpool = MemoryPool(int(dfacts.DEFAULT_BE_OVERLAY_TRANSPORT_SEG_SZ),
                              f'{os.getuid()}_{os.getpid()}_{2}' + dfacts.DEFAULT_POOL_SUFFIX,
                              dfacts.be_pool_muid_from_hostid(2))

        be_ch_out = Channel.attach(B64.from_str(frontend_sdesc).decode(),
                                   mem_pool=be_mpool)

        net = NetworkConfig.from_file(network_config)
        net_conf = net.get_network_config()

        be_nodes = {}
        for node in net_conf.values():
            be_cuid = dfacts.be_fe_cuid_from_hostid(node.host_id)

            # Add this to the object as it lets teardown know it needs to clean this up
            be_ch_in = Channel(be_mpool, be_cuid)

            overlay_inout = Connection(inbound_initializer=be_ch_in,
                                     outbound_initializer=be_ch_out,
                                     options=conn_options,
                                     policy=conn_policy)
            overlay_inout.ghost = True
            be_nodes[node.host_id] = {'conn': overlay_inout,
                                      'ch_in': be_ch_in,
                                      'hostname': node.name,
                                      'ip_addrs': node.ip_addrs}

    except (ChannelError, DragonPoolError, DragonMemoryError) as init_err:
        # Try to clean up the pool
        if be_mpool is not None:
            be_mpool.destroy()
        raise RuntimeError('Overlay transport resource creation failed') from init_err

    return be_mpool, be_ch_out, be_ch_in, be_nodes, overlay_inout


def send_beisup(nodes):
    '''Send valid beisup messages'''

    for host_id, node in nodes.items():
        be_up_msg = dmsg.BEIsUp(tag=next_tag(),
                                be_ch_desc=str(B64(node['ch_in'].serialize())),
                                host_id=host_id)
        node['conn'].send(be_up_msg.serialize())


def recv_fenodeidx(nodes):
    '''recv FENoe4deIdxBE and finish filling out node dictionary'''
    log = logging.getLogger('recv_fe_nodeidx')
    for node in nodes.values():
        fe_node_idx_msg = dmsg.parse(node['conn'].recv())
        assert isinstance(fe_node_idx_msg, dmsg.FENodeIdxBE), 'la_be node_index from fe expected'

        log.info(f'got FENodeIdxBE for index {fe_node_idx_msg.node_index}')
        node['node_index'] = fe_node_idx_msg.node_index
        if node['node_index'] < 0:
            raise RuntimeError("frontend giving bad node indices")
        node['is_primary'] = node['node_index'] == 0
        if node['is_primary']:
            primary_conn = node['conn']
        log.info(f'constructed be node: {node}')

    return primary_conn


def send_shchannelsup(nodes, mpool):

    log = logging.getLogger('send_shchannelsup')
    for host_id, node in nodes.items():
        ls_cuid = dfacts.shepherd_cuid_from_index(node['node_index'])
        ls_ch = Channel(mpool, ls_cuid)
        node['ls_ch'] = ls_ch

        if node['is_primary']:
            node['gs_ch'] = Channel(mpool, dfacts.GS_INPUT_CUID)
            gs_cd = B64.bytes_to_str(node['gs_ch'].serialize())
        else:
            node['gs_ch'] = None
            gs_cd = None

        ch_up_msg = dmsg.SHChannelsUp(tag=next_tag(),
                                      host_name=node['hostname'],
                                      host_id=host_id,
                                      ip_addrs=node['ip_addrs'],
                                      shep_cd=B64.bytes_to_str(node['ls_ch'].serialize()),
                                      gs_cd=gs_cd,
                                      idx=node['node_index'])
        log.info(f'construct SHChannelsUp: {ch_up_msg}')
        node['conn'].send(ch_up_msg.serialize())
        log.info(f'sent SHChannelsUp for {node["node_index"]}')

    log.info('sent all SHChannelsUp')


def recv_lachannelsinfo(nodes):
    '''Loop to recv all LAChannelsInfo messages in frontend bcast'''
    for host_id, node in nodes.items():
        la_channels_info_msg = dmsg.parse(node['conn'].recv())
        assert isinstance(la_channels_info_msg, dmsg.LAChannelsInfo), 'la_be expected all ls channels info from la_fe'
    return la_channels_info_msg


def send_taup(nodes):
    for host_id, node in nodes.items():
        ta_up = dmsg.TAUp(tag=next_tag(), idx=node['node_index'])
        node['conn'].send(ta_up.serialize())


def send_abnormal_term(conn):
    abnorm = dmsg.AbnormalTermination(tag=next_tag())
    conn.send(abnorm.serialize())


def handle_gsprocesscreate(primary_conn):
    '''Manage a valid response to GSProcessCreate'''

    log = logging.getLogger('handle_gsprocesscreate')
    proc_create = dmsg.parse(primary_conn.recv())
    log.info('presumably got GSProcessCreate')
    assert isinstance(proc_create, dmsg.GSProcessCreate)
    log.info('recvd GSProccessCreate')

    # Send response
    gs_desc = ProcessDescriptor(p_uid=5000,  # Just a dummy value
                                name=proc_create.user_name,
                                node=0,
                                p_p_uid=proc_create.p_uid)
    response = dmsg.GSProcessCreateResponse(tag=next_tag(),
                                            ref=proc_create.tag,
                                            err=dmsg.GSProcessCreateResponse.Errors.SUCCESS,
                                            desc=gs_desc)
    primary_conn.send(response.serialize())


def stand_up_backend(mock_overlay, mock_launch, network_config):

    log = logging.getLogger('mock_backend_standup')
    # Get the mock's input args to the
    while mock_overlay.call_args is None:
        pass
    overlay_args = mock_overlay.call_args.kwargs

    overlay = {}

    # Connect to overlay comms to talk to fronteend
    overlay['ta_ch_in'], overlay['ta_ch_out'], overlay['fe_ta_conn'] = open_overlay_comms(overlay_args['ch_in_sdesc'],
                                                                                    overlay_args['ch_out_sdesc'])

    # Let frontend know the overlay is "up"
    overlay['fe_ta_conn'].send(dmsg.OverlayPingLA(next_tag()).serialize())

    # Grab the frontend channel descriptor for the launched backend and
    # send it mine
    while mock_launch.call_args is None:
        pass
    launch_be_args = mock_launch.call_args.kwargs
    log.info(f'got be args: {launch_be_args}')

    # Connect to backend comms for frontend-to-backend and back comms
    overlay['be_mpool'], overlay['be_ch_out'], overlay['be_ch_in'], overlay['be_nodes'], overlay['overlay_inout'] = open_backend_comms(launch_be_args['frontend_sdesc'],
                                                                                                                           network_config)
    log.info('got backend up')

    return overlay


def handle_bringup(mock_overlay, mock_launch, network_config):

    log = logging.getLogger('mock_fulL_bringup')

    overlay = stand_up_backend(mock_overlay, mock_launch, network_config)

    # Send BEIsUp
    send_beisup(overlay['be_nodes'])
    log.info('send BEIsUp messages')

    # Recv FENodeIdxBE
    overlay['primary_conn'] = recv_fenodeidx(overlay['be_nodes'])
    log.info('got all the FENodeIdxBE messages')

    # Fudge some SHChannelsUp messages
    send_shchannelsup(overlay['be_nodes'], overlay['be_mpool'])
    log.info(f'sent shchannelsup: {[node["gs_ch"].serialize() for node in overlay["be_nodes"].values() if node["gs_ch"] is not None]}')

    # Receive LAChannelsInfo
    la_info = recv_lachannelsinfo(overlay['be_nodes'])
    log.info('la_be received LAChannelsInfo')

    # Send TAUp
    send_taup(overlay['be_nodes'])
    log.info('sent TAUp messages')

    # Send gs is up from primary
    overlay['primary_conn'].send(dmsg.GSIsUp(tag=next_tag()).serialize())
    log.info('send GSIsUp')

    return overlay, la_info


def handle_overlay_teardown(overlay_conn):
    '''complete teardown of frontend overlay process'''
    halt_on = dmsg.parse(overlay_conn.recv())
    assert isinstance(halt_on, dmsg.LAHaltOverlay)
    overlay_conn.send(dmsg.OverlayHalted(tag=next_tag()).serialize())


def handle_teardown(nodes, primary_conn, overlay_conn,
                    timeout_backend=False,
                    timeout_overlay=False,
                    gs_head_exit=True,
                    abort_shteardown=None):
    '''Do full teardown from perspective of backend'''

    if gs_head_exit:
        primary_conn.send(dmsg.GSHeadExit(exit_code=0, tag=next_tag()).serialize())

    # Recv GSTeardown
    gs_teardown = dmsg.parse(primary_conn.recv())
    assert isinstance(gs_teardown, dmsg.GSTeardown)
    log = logging.getLogger('handle_teardown')
    log.debug('got gstteardown')

    primary_conn.send(dmsg.GSHalted(tag=next_tag()).serialize())
    log.debug('sent gshalted')

    # Recv SHHaltTA for everyone and send back TAHalted
    for node in nodes.values():
        ta_halt_msg = dmsg.parse(node['conn'].recv())
        assert isinstance(ta_halt_msg, dmsg.SHHaltTA), 'SHHaltTA from fe expected'
        node['conn'].send(dmsg.TAHalted(tag=next_tag()).serialize())
    log.debug('received all SHHaltTA from frontend')

    if timeout_backend:
        log.info('returning early during teardown to test frontend teardown timeout')
        return

    # Recv SHTeardown and send shhaltbe
    for index, node in enumerate(nodes.values()):
        sh_teardown_msg = dmsg.parse(node['conn'].recv())
        assert isinstance(sh_teardown_msg, dmsg.SHTeardown), 'SHTeardown from fe expected'
        if abort_shteardown is not None and index == abort_shteardown:
            node['conn'].send(dmsg.AbnormalTermination(tag=next_tag()).serialize())
        node['conn'].send(dmsg.SHHaltBE(tag=next_tag()).serialize())

    # Recv BEHalted
    for node in nodes.values():
        be_halted_msg = dmsg.parse(node['conn'].recv())
        assert isinstance(be_halted_msg, dmsg.BEHalted), 'BEHalted from fe expected'

    # Recv FE's LAHaltOverlay for overlay and then tell it we've shut down our mocked up overlay tree
    if not timeout_overlay:
        handle_overlay_teardown(overlay_conn)

test/launcher/frontend_testing_mocks.py
DragonHPC/dragon
import threading
from functools import wraps


def catch_thread_exceptions(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        exceptions_caught_in_threads = {}

        def custom_excepthook(args):
            thread_name = args.thread.name
            exceptions_caught_in_threads[thread_name] = {
                'thread': args.thread,
                'exception': {
                    'type': args.exc_type,
                    'value': args.exc_value,
                    'traceback': args.exc_traceback
                }
            }
        # Registering our custom excepthook to catch the exception in the threads
        old_excepthook = threading.excepthook
        threading.excepthook = custom_excepthook

        result = func(*args + (exceptions_caught_in_threads, ), **kwargs)

        threading.excepthook = old_excepthook
        return result

    return wrapper

test/launcher/launcher_testing_utils.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import logging
import threading
import string
import random

import unittest
from unittest.mock import patch

from dragon.infrastructure.facts import PROCNAME_LA_BE
from dragon.infrastructure.messages import AbnormalTerminationError

from dragon.launcher.backend import LauncherBackEnd
from .launcher_testing_utils import catch_thread_exceptions
from .backend_testing_mocks import LauncherBackendHelper, mock_start_localservices_wrapper


def start_backend(args):

    # Before doing anything set my host ID
    from dragon.utils import set_procname
    from dragon.dlogging.util import DragonLoggingServices as dls
    from dragon.dlogging.util import setup_BE_logging, LOGGING_OUTPUT_DEVICE_ACTOR_FILE

    # Enable debug logging in BE service
    os.environ[
        f"DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_ACTOR_FILE.upper()}"
    ] = "DEBUG"

    set_procname(PROCNAME_LA_BE)

    # Set up the logging level. If debug, start logging immediately
    level, fname = setup_BE_logging(dls.LA_BE)

    log = logging.getLogger("start_backend")
    for key, value in args.items():
        if value is not None:
            log.info(f"args: {key}: {value}")
    with LauncherBackEnd(args["transport_test"], args["network_prefix"]) as be_server:
        be_server.run_startup(args["ip_addrs"], args["host_ids"], args["frontend_sdesc"], level, fname)
        be_server.run_msg_server()

    del os.environ[
        f"DRAGON_LOG_DEVICE_{LOGGING_OUTPUT_DEVICE_ACTOR_FILE.upper()}"
    ]


class BackendBringUpTeardownTest(unittest.TestCase):
    def setUp(self):

        self.test_dir = os.path.dirname(os.path.realpath(__file__))
        self.network_config = os.path.join(self.test_dir, "slurm_primary.yaml")

        self.node_index = 0  # so localservices will create GS_CD
        self.ip_addrs = "127.0.0.1:6570"
        self.fe_host_id = "fe"
        self.be_host_id = "be"

        self.be_helper = LauncherBackendHelper(self.network_config,
                                               self.node_index,
                                               self.ip_addrs,
                                               self.fe_host_id,
                                               self.be_host_id)

    def tearDown(self):

        # Make sure we don't leave any logging handers up
        log = logging.getLogger()
        for handler in log.handlers:
            log.removeHandler(handler)
            handler.close()

    def start_backend_thread(self, args_map):

        # get startup going in another thread. Note: need to do threads
        # in order to use all our mocks
        self.be_thread = threading.Thread(
            name="Backend Server", target=start_backend, args=(args_map,), daemon=False
        )
        self.be_thread.start()

    @mock_start_localservices_wrapper
    @patch("dragon.launcher.backend.start_overlay_network")
    @patch("dragon.launcher.backend.NodeDescriptor.get_local_node_network_conf")
    def test_clean_exit(self, mock_localservices_tuple, mock_network_config, mock_overlay):
        """Test that we can come up and exit cleanly"""

        log = logging.getLogger("clean_be_startup_teardown")

        args_map = self.be_helper.handle_network_and_frontend_start(mock_network_config)
        self.start_backend_thread(args_map)

        self.be_helper.clean_startup(log, mock_overlay, mock_network_config, mock_localservices_tuple)
        self.be_helper.launch_user_process(log)
        self.be_helper.clean_shutdown(log)
        self.be_helper.cleanup()

        # join on the be thread
        self.be_thread.join()

    @catch_thread_exceptions
    @mock_start_localservices_wrapper
    @patch("dragon.launcher.backend.start_overlay_network")
    @patch("dragon.launcher.backend.NodeDescriptor.get_local_node_network_conf")
    def test_local_services_launch_exception(self, exceptions_caught_in_threads,
                                             mock_localservices_tuple, mock_network_config, mock_overlay):
        '''Test unable to Popen backend launcher'''

        # Set up our breaking of local services start
        all_chars = string.ascii_letters + string.digits + string.punctuation
        exception_text = ''.join(random.choices(all_chars, k=64))
        mock_start_localservices, _, _ = mock_localservices_tuple
        mock_start_localservices.side_effect = RuntimeError(exception_text)

        # Get our frontend crap going so we can get the overlay up
        args_map = self.be_helper.handle_network_and_frontend_start(mock_network_config)
        self.start_backend_thread(args_map)

        # Start up the overlay
        self.be_helper.handle_overlay_start(mock_overlay, mock_localservices_tuple)

        # At this point, our LS mock should throw its side-effect:
        self.be_thread.join()

        # Clean up local services
        self.be_helper.cleanup()

        log = logging.getLogger('test_local_services_launch_exception')
        log.info(f'exceptions: {exceptions_caught_in_threads}')
        assert mock_start_localservices.call_count == 1
        assert 'Backend Server' in exceptions_caught_in_threads
        assert exceptions_caught_in_threads['Backend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Backend Server']['exception']['value']) == exception_text

    @catch_thread_exceptions
    @mock_start_localservices_wrapper
    @patch("dragon.launcher.backend.start_overlay_network")
    @patch("dragon.launcher.backend.NodeDescriptor.get_local_node_network_conf")
    def test_garbled_lachannelsinfo(self, exceptions_caught_in_threads,
                                    mock_localservices_tuple, mock_network_config, mock_overlay):
        '''Test that we exit correctly if the LAChannelsInfo we receive is garbled'''

        log = logging.getLogger('test_garbled_shchannelsup')

        # Get our frontend crap going so we can get the overlay up
        args_map = self.be_helper.handle_network_and_frontend_start(mock_network_config)
        self.start_backend_thread(args_map)

        # Cleanly start everything to LS returning its SHChannelsUp
        self.be_helper.start_ls(mock_overlay, mock_localservices_tuple, garble_lachannelsinfo=True)

        # Confirm we send out an abnormal term to the frontend so it can
        # issue a GSTeardown to the other nodes
        self.be_helper.start_abnormal_termination_exit(log, gs_up=False)

        # Clean up local services
        self.be_helper.cleanup()

        # At this point, our LS mock should throw its side-effect:
        self.be_thread.join()
        log.debug(f'exceptions: {exceptions_caught_in_threads}')
        assert 'Backend Server' in exceptions_caught_in_threads
        assert exceptions_caught_in_threads['Backend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Backend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @mock_start_localservices_wrapper
    @patch("dragon.launcher.backend.start_overlay_network")
    @patch("dragon.launcher.backend.NodeDescriptor.get_local_node_network_conf")
    def test_garbled_shchannelsup(self, exceptions_caught_in_threads,
                                  mock_localservices_tuple, mock_network_config, mock_overlay):
        '''Test that we exit correctly if the SHChannelsUp from LS is garbled'''

        log = logging.getLogger('test_garbled_shchannelsup')

        # Get our frontend crap going so we can get the overlay up
        args_map = self.be_helper.handle_network_and_frontend_start(mock_network_config)
        self.start_backend_thread(args_map)

        # Cleanly start everything up to SHChannelsUp
        self.be_helper.start_ls(mock_overlay, mock_localservices_tuple, garble_shchannelsup=True)

        # Confirm we send out an abnormal term to the frontend so it can
        # issue a GSTeardown to the other nodes
        self.be_helper.start_abnormal_termination_exit(log, gs_up=False, ls_up=False)

        # Clean up local services
        self.be_helper.cleanup()

        # At this point, our LS mock should throw its side-effect:
        self.be_thread.join()

        log.info(f'exceptions: {exceptions_caught_in_threads}')
        assert 'Backend Server' in exceptions_caught_in_threads
        assert exceptions_caught_in_threads['Backend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Backend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @mock_start_localservices_wrapper
    @patch("dragon.launcher.backend.start_overlay_network")
    @patch("dragon.launcher.backend.NodeDescriptor.get_local_node_network_conf")
    def test_app_exec_abnormal_term_from_frontend(self, exceptions_caught_in_threads,
                                                  mock_localservices_tuple, mock_network_config, mock_overlay):
        '''Test that we exit if we get an AbnormalTermination from the frontend during app execution'''

        log = logging.getLogger('test_app_exec_abnormal_term_from_frontend')

        # Get our frontend crap going so we can get the overlay up
        args_map = self.be_helper.handle_network_and_frontend_start(mock_network_config)
        self.start_backend_thread(args_map)

        # Cleanly start everything up to SHChannelsUp
        self.be_helper.clean_startup(log, mock_overlay, mock_network_config, mock_localservices_tuple)
        self.be_helper.launch_user_process(log)

        self.be_helper.localservices.send_AbnormalTermination()

        self.be_helper.start_abnormal_termination_exit(log)

        # Clean up local services
        self.be_helper.cleanup()
        log.debug('cleaned up backend helpers')

        # Should be able to join here
        self.be_thread.join()
        log.debug('cleaned up backend')

        log.info(f'exceptions: {exceptions_caught_in_threads}')
        assert 'Backend Server' in exceptions_caught_in_threads
        assert exceptions_caught_in_threads['Backend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Backend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @mock_start_localservices_wrapper
    @patch("dragon.launcher.backend.start_overlay_network")
    @patch("dragon.launcher.backend.NodeDescriptor.get_local_node_network_conf")
    def test_early_abnormal_term_from_frontend(self, exceptions_caught_in_threads,
                                               mock_localservices_tuple, mock_network_config, mock_overlay):
        '''Test that we exit if we get an AbnormalTermination from the frontend during app execution'''

        log = logging.getLogger('test_early_abnormal_term_from_frontend')

        # Get our frontend crap going so we can get the overlay up
        args_map = self.be_helper.handle_network_and_frontend_start(mock_network_config)
        self.start_backend_thread(args_map)

        # Cleanly start everything to LS returning its SHChannelsUp
        self.be_helper.start_ls(mock_overlay, mock_localservices_tuple, abort_lachannelsinfo=True)

        # Confirm we send out an abnormal term to the frontend so it can
        # issue a GSTeardown to the other nodes
        self.be_helper.start_abnormal_termination_exit(log, gs_up=False)

        # Clean up local services
        self.be_helper.cleanup()

        # At this point, our LS mock should throw its side-effect:
        self.be_thread.join()

        log.debug(f'exceptions: {exceptions_caught_in_threads}')
        assert 'Backend Server' in exceptions_caught_in_threads
        assert exceptions_caught_in_threads['Backend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Backend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @mock_start_localservices_wrapper
    @patch("dragon.launcher.backend.start_overlay_network")
    @patch("dragon.launcher.backend.NodeDescriptor.get_local_node_network_conf")
    def test_late_abnormal_term_from_frontend(self, exceptions_caught_in_threads,
                                              mock_localservices_tuple, mock_network_config, mock_overlay):
        '''Test that we exit if we get an AbnormalTermination from the frontend during app execution'''

        log = logging.getLogger('test_late_abnormal_term_from_frontend')

        # Get our frontend crap going so we can get the overlay up
        args_map = self.be_helper.handle_network_and_frontend_start(mock_network_config)
        self.start_backend_thread(args_map)

        # Cleanly start everything up to SHChannelsUp
        self.be_helper.clean_startup(log, mock_overlay, mock_network_config, mock_localservices_tuple)
        self.be_helper.launch_user_process(log)
        self.be_helper.clean_shutdown(log, shteardown_abort=True)

        # Clean up local services
        self.be_helper.cleanup()

        # Should be able to join here
        self.be_thread.join()

        log.info(f'exceptions: {exceptions_caught_in_threads}')
        assert 'Backend Server' in exceptions_caught_in_threads
        assert exceptions_caught_in_threads['Backend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Backend Server']['exception']['value']) == 'Abnormal exit detected'


if __name__ == "__main__":
    # Use the fork method for these so I can do mocking inside the mp process.
    unittest.main()

test/launcher/test_backend_bringup.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import logging
import random
import string
import unittest
import threading
import json
from subprocess import TimeoutExpired
from unittest.mock import patch

from dragon.launcher.util import next_tag
from dragon.launcher.launchargs import get_parser
from dragon.launcher.network_config import NetworkConfig

from dragon.infrastructure import messages as dmsg
from dragon.launcher.wlm import WLM
from dragon.channels import ChannelError
from dragon.managed_memory import DragonMemoryError
from dragon.utils import B64
from dragon.infrastructure.facts import DEFAULT_TRANSPORT_NETIF, DEFAULT_OVERLAY_NETWORK_PORT, TransportAgentOptions

from .launcher_testing_utils import catch_thread_exceptions

from .frontend_testing_mocks import run_frontend, open_backend_comms, open_overlay_comms
from .frontend_testing_mocks import send_beisup, handle_teardown, recv_fenodeidx, send_shchannelsup, recv_lachannelsinfo
from .frontend_testing_mocks import handle_gsprocesscreate, handle_bringup, stand_up_backend, handle_overlay_teardown
from .frontend_testing_mocks import send_abnormal_term


def get_args_map(network_config, **kwargs):

    parser = get_parser()
    arg_list = ['--wlm', 'slurm',
                '--network-config', f'{network_config}',
                '--network-prefix', '^(eth|hsn)']
    for val in kwargs.values():
        arg_list = arg_list + val

    arg_list.append('hello_world.py')

    args = parser.parse_args(args=arg_list)
    if args.basic_label or args.verbose_label:
        args.no_label = False
    args_map = {key: value for key, value in vars(args).items() if value is not None}

    return args_map


class FrontendBringUpTeardownTest(unittest.TestCase):

    def setUp(self):
        self.test_dir = os.path.dirname(os.path.realpath(__file__))
        self.network_config = os.path.join(self.test_dir, 'slurm_primary.yaml')
        self.bad_network_config = os.path.join(self.test_dir, 'slurm_bad.yaml')

        self.be_mpool = None
        self.be_ch_out = None
        self.be_ch_in = None
        self.overlay_inout = None

        self.ls_ch = None

        self.ta_ch_in = None
        self.ta_ch_out = None
        self.fe_ta_conn = None

    def tearDown(self):

        try:
            self.fe_ta_conn.close()
        except (ConnectionError, AttributeError):
            pass

        try:
            self.be_ch_out.detach()
        except (ChannelError, AttributeError):
            pass

        try:
            for node in self.be_nodes.values():
                node['conn'].close()
                node['ch_in'].destroy()
        except (AttributeError, ChannelError):
            pass

        try:
            for node in self.be_nodes.values():
                node['ls_ch'].destroy()
                if node['is_primary']:
                    node['gs_ch'].destroy()
        except (AttributeError, ChannelError, KeyError):
            pass

        try:
            self.ta_ch_out.detach()
        except (AttributeError, ChannelError):
            pass

        try:
            self.ta_ch_in.detach()
        except (AttributeError, ChannelError):
            pass

        try:
            self.be_mpool.destroy()
            del self.be_mpool
        except (AttributeError, DragonMemoryError):
            pass

    def do_bringup(self, mock_overlay, mock_launch):

        overlay, la_info = handle_bringup(mock_overlay, mock_launch, self.network_config)
        self.ta_ch_in = overlay['ta_ch_in']
        self.ta_ch_out = overlay['ta_ch_out']
        self.fe_ta_conn = overlay['fe_ta_conn']
        self.be_mpool = overlay['be_mpool']
        self.be_ch_out = overlay['be_ch_out']
        self.be_ch_in = overlay['be_ch_in']
        self.be_nodes = overlay['be_nodes']
        self.overlay_inout = overlay['overlay_inout']
        self.primary_conn = overlay['primary_conn']

        return la_info

    def get_backend_up(self, mock_overlay, mock_launch):

        overlay = stand_up_backend(mock_overlay, mock_launch, self.network_config)
        self.ta_ch_in = overlay['ta_ch_in']
        self.ta_ch_out = overlay['ta_ch_out']
        self.fe_ta_conn = overlay['fe_ta_conn']
        self.be_mpool = overlay['be_mpool']
        self.be_ch_out = overlay['be_ch_out']
        self.be_ch_in = overlay['be_ch_in']
        self.be_nodes = overlay['be_nodes']
        self.overlay_inout = overlay['overlay_inout']

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_clean_exit(self, mock_overlay, mock_launch):
        '''Test a clean bring-up and teardown'''

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        la_info = self.do_bringup(mock_overlay, mock_launch)

        # Check we launched the backend with default transport
        self.assertEqual(la_info.transport, TransportAgentOptions.TCP)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send GSHalted
        handle_teardown(self.be_nodes, self.primary_conn, self.fe_ta_conn)

        # Join on the frontend thread
        fe_proc.join()

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_clean_exit_with_tcp_launch(self, mock_overlay, mock_launch):
        '''Test a clean bring-up and teardown using TCP agent'''

        args_map = get_args_map(self.network_config,
                                arg1=['-t', 'tcp'])

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        la_info = self.do_bringup(mock_overlay, mock_launch)

        # Check we launched the backend with default transport
        self.assertEqual(la_info.transport, TransportAgentOptions.TCP)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send GSHalted
        handle_teardown(self.be_nodes, self.primary_conn, self.fe_ta_conn)

        # Join on the frontend thread
        fe_proc.join()

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_default_transport(self, mock_overlay, mock_launch):
        '''Test a clean bring-up and teardown using default transport (TCP)'''

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        self.do_bringup(mock_overlay, mock_launch)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send GSHalted
        handle_teardown(self.be_nodes, self.primary_conn, self.fe_ta_conn)

        # Join on the frontend thread
        fe_proc.join()

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_bad_network_config(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test with an invalid network config'''

        args_map = get_args_map(self.bad_network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)

        fe_proc.start()
        fe_proc.join()

        logging.info(f"exception: {exceptions_caught_in_threads}")
        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Unable to acquire backend network configuration from input file.'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_launch_backend_exception(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test unable to Popen backend launcher'''

        all_chars = string.ascii_letters + string.digits + string.punctuation
        exception_text = ''.join(random.choices(all_chars, k=64))
        mock_launch.side_effect = RuntimeError(exception_text)

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)

        fe_proc.start()

        # Get the mock's input args to the
        while mock_overlay.call_args is None:
            pass
        overlay_args = mock_overlay.call_args.kwargs

        # Connect to overlay comms to talk to fronteend
        self.ta_ch_in, self.ta_ch_out, self.fe_ta_conn = open_overlay_comms(overlay_args['ch_in_sdesc'],
                                                                            overlay_args['ch_out_sdesc'])
        # Let frontend know the overlay is "up"
        self.fe_ta_conn.send(dmsg.OverlayPingLA(next_tag()).serialize())

        handle_overlay_teardown(self.fe_ta_conn)

        fe_proc.join()

        assert mock_launch.call_count == 1
        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == "Backend launch failed from launcher frontend"

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_launch_overlay_exception(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test exception during launch of overlay network'''

        all_chars = string.ascii_letters + string.digits + string.punctuation
        exception_text = ''.join(random.choices(all_chars, k=64))
        mock_overlay.side_effect = RuntimeError(exception_text)

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)

        fe_proc.start()
        fe_proc.join()

        assert mock_overlay.call_count == 1
        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) \
            == 'Overlay transport agent launch failed on launcher frontend'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_message_out_of_order(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test message out of order'''

        log = logging.getLogger('msg_out_of_order')
        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get the mock's input args to the
        while mock_overlay.call_args is None:
            pass
        overlay_args = mock_overlay.call_args.kwargs

        # Connect to overlay comms to talk to fronteend
        self.ta_ch_in, self.ta_ch_out, self.fe_ta_conn = open_overlay_comms(overlay_args['ch_in_sdesc'],
                                                                            overlay_args['ch_out_sdesc'])

        # Let frontend know the overlay is "up"
        self.fe_ta_conn.send(dmsg.OverlayPingLA(next_tag()).serialize())

        # Grab the frontend channel descriptor for the launched backend and
        # send it mine
        while mock_launch.call_args is None:
            pass
        launch_be_args = mock_launch.call_args.kwargs
        log.info(f'got be args: {launch_be_args}')

        # Connect to backend comms for frontend-to-backend and back comms
        self.be_mpool, self.be_ch_out, self.be_ch_in, self.be_nodes, self.overlay_inout = open_backend_comms(launch_be_args['frontend_sdesc'],
                                                                                                           self.network_config)
        log.info('got backend up')

        # Send BEIsUp
        send_beisup(self.be_nodes)
        log.info('send BEIsUp messages')

        # Recv FENodeIdxBE
        recv_fenodeidx(self.be_nodes)
        log.info('got all the FENodeIdxBE messages')

        # The Dragon FrontEnd is expecting us to some SHChannelsUp messages.
        # Instead, we're going to send a message that the FE is not expecting.

        # Send BEIsUp
        send_beisup(self.be_nodes)
        log.info('send BEIsUp messages')

        handle_overlay_teardown(self.fe_ta_conn)

        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_garbled_beisup(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test receiving a garbage message in middle of loop'''

        log = logging.getLogger('garbled_beisup')
        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        self.get_backend_up(mock_overlay, mock_launch)
        log.info('got backend up')

        # Send BEIsUp
        for i, (host_id, node) in enumerate(self.be_nodes.items()):
            be_up_msg = dmsg.BEIsUp(tag=next_tag(),
                                    be_ch_desc=str(B64(node['ch_in'].serialize())),
                                    host_id=host_id)
            if i == 1:
                garbage_dict = {'junk': 'abdecdf', '52': 20}
                node['conn'].send(json.dumps(garbage_dict))
            else:
                node['conn'].send(be_up_msg.serialize())

        handle_overlay_teardown(self.fe_ta_conn)

        log.debug('sitting on frontend join')
        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._kill_backend')
    @patch('dragon.launcher.frontend.LauncherFrontEnd._wait_on_wlm_proc_exit')
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sigkill_to_wlm(self, exceptions_caught_in_threads, mock_overlay, mock_launch,
                            mock_wait, mock_kill):
        '''Test that we transmit a SIGKILL for a non-responsive backend'''
        mock_wait.side_effect = TimeoutExpired(None, 'raising a fake timeout')

        log = logging.getLogger('sigkill_to_wlm')
        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        self.get_backend_up(mock_overlay, mock_launch)
        log.info('got backend up')

        # Send BEIsUp
        for i, (host_id, node) in enumerate(self.be_nodes.items()):
            be_up_msg = dmsg.BEIsUp(tag=next_tag(),
                                    be_ch_desc=str(B64(node['ch_in'].serialize())),
                                    host_id=host_id)
            if i == 1:
                garbage_dict = {'junk': 'abdecdf', '52': 20}
                node['conn'].send(json.dumps(garbage_dict))
            else:
                node['conn'].send(be_up_msg.serialize())

        handle_overlay_teardown(self.fe_ta_conn)
        log.debug('sitting on frontend join')
        fe_proc.join()

        log.info(f'exception: {exceptions_caught_in_threads}')
        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'
        mock_wait.assert_called_once()
        mock_kill.assert_called_once()

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_garbled_taup(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test receiving garbage TAUp in middle of loop'''

        log = logging.getLogger('garbled_taup')
        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        self.get_backend_up(mock_overlay, mock_launch)
        log.info('got backend up')

        # Send BEIsUp
        send_beisup(self.be_nodes)
        log.info('send BEIsUp messages')

        # Recv FENodeIdxBE
        self.primary_conn = recv_fenodeidx(self.be_nodes)
        log.info('got all the FENodeIdxBE messages')

        # Fudge some SHChannelsUp messages
        send_shchannelsup(self.be_nodes, self.be_mpool)
        log.info(f'sent shchannelsup: {[node["gs_ch"].serialize() for node in self.be_nodes.values() if node["gs_ch"] is not None]}')

        # Receive LAChannelsInfo
        recv_lachannelsinfo(self.be_nodes)
        log.info('la_be received LAChannelsInfo')

        # Send TAUp
        for i, (host_id, node) in enumerate(self.be_nodes.items()):
            if i == 2:
                garbage_msg = {'fizz': 'abdecdf', 'buzz': 20}
                node['conn'].send(json.dumps(garbage_msg))
            else:
                ta_up = dmsg.TAUp(tag=next_tag(), idx=node['node_index'])
                node['conn'].send(ta_up.serialize())

        handle_overlay_teardown(self.fe_ta_conn)

        log.debug('sitting on frontend join')
        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_abnormal_term_in_app_exec(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test the ability of frontend to exit from an Abnormal Termination during app execution'''

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        self.do_bringup(mock_overlay, mock_launch)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send an abormal termination rather than proceeding with teardown
        for i, (host_id, node) in enumerate(self.be_nodes.items()):
            if i == 2:
                send_abnormal_term(node['conn'])

        # Proceed with teardown
        handle_teardown(self.be_nodes, self.primary_conn, self.fe_ta_conn, gs_head_exit=False)

        # Join on the frontend thread
        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_backend_timeout_in_abnormal_teardown(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test the frontend is able to recover from a hung backend'''

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        self.do_bringup(mock_overlay, mock_launch)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send an abormal termination rather than proceeding with teardown
        for i, (host_id, node) in enumerate(self.be_nodes.items()):
            if i == 2:
                send_abnormal_term(node['conn'])

        # Proceed with teardown
        handle_teardown(self.be_nodes,
                        self.primary_conn,
                        self.fe_ta_conn,
                        gs_head_exit=False,
                        timeout_backend=True)

        # Join on the frontend thread
        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_overlay_timeout_in_abnormal_teardown(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test the frontend is able to recover from a hung overlay transport agent'''

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        self.do_bringup(mock_overlay, mock_launch)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send an abormal termination rather than proceeding with teardown
        for i, (host_id, node) in enumerate(self.be_nodes.items()):
            if i == 2:
                send_abnormal_term(node['conn'])

        # Proceed with teardown
        handle_teardown(self.be_nodes,
                        self.primary_conn,
                        self.fe_ta_conn,
                        gs_head_exit=False,
                        timeout_overlay=True)

        # Join on the frontend thread
        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_abnormal_term_in_teardown(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        '''Test the ability of frontend to exit from an Abnormal Termination during teardown'''

        args_map = get_args_map(self.network_config)

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        self.do_bringup(mock_overlay, mock_launch)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        handle_teardown(self.be_nodes, self.primary_conn, self.fe_ta_conn,
                        gs_head_exit=True,
                        abort_shteardown=3)

        fe_proc.join()

        log = logging.getLogger('test_abnormal_term_in_teardown')
        log.info(f'exception = {exceptions_caught_in_threads}')
        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) == 'Abnormal exit detected'

    @patch('dragon.launcher.network_config.NetworkConfig.from_wlm')
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_wlm_slurm_launch_selection(self, mock_overlay, mock_launch, mock_config):
        """Test that we honor user's slurm WLM selection rather than auto detect"""
        parser = get_parser()
        arg_list = ['--wlm', 'slurm',
                    '--network-prefix', '^(eth|hsn)',
                    'hello_world.py']

        args = parser.parse_args(args=arg_list)
        if args.basic_label or args.verbose_label:
            args.no_label = False
        args_map = {key: value for key, value in vars(args).items() if value is not None}

        # Give the net conf a correct object so that rest of this test can complete
        net_config_obj = NetworkConfig.from_file(self.network_config)
        mock_config.side_effect = [net_config_obj]

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        self.do_bringup(mock_overlay, mock_launch)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send GSHalted
        handle_teardown(self.be_nodes, self.primary_conn, self.fe_ta_conn)

        # Join on the frontend thread
        fe_proc.join()

        # Check the config mock actually use slurm
        assert mock_config.called_with(workload_manager=WLM.SLURM,
                                       port=DEFAULT_OVERLAY_NETWORK_PORT,
                                       network_prefix=DEFAULT_TRANSPORT_NETIF,
                                       sigint_trigger=None)

    @catch_thread_exceptions
    @patch('dragon.launcher.network_config.NetworkConfig.from_wlm')
    def test_wlm_ssh_auto_detect(self, exceptions_caught_in_threads, mock_config):
        """Test we throw an error if a user doesn't specify a WLM, and we only detect ssh"""
        parser = get_parser()
        arg_list = ['-l', 'dragon_file=DEBUG',
                    '--network-prefix', '^(eth|hsn)',
                    'hello_world.py']
        args = parser.parse_args(args=arg_list)
        if args.basic_label or args.verbose_label:
            args.no_label = False
        args_map = {key: value for key, value in vars(args).items() if value is not None}

        # Give the net conf a correct object so that rest of this test can complete
        net_config_obj = NetworkConfig.from_file(self.network_config)
        mock_config.side_effect = [net_config_obj]

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()
        fe_proc.join()

        log = logging.getLogger('auto_ssh')
        log.error(f'exception: {exceptions_caught_in_threads}')
        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == RuntimeError
        assert 'SSH was only supported launcher found' in str(exceptions_caught_in_threads['Frontend Server']['exception']['value'])

    @patch('dragon.launcher.network_config.NetworkConfig.from_wlm')
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_pbs_slurm_launch_selection(self, mock_overlay, mock_launch, mock_config):
        """Test that we honor user's pbs+pals WLM selection rather than auto detect"""
        args_map = get_args_map(self.network_config,
                                arg1=['--wlm', 'slurm'])

        parser = get_parser()
        arg_list = ['--wlm', 'pbs+pals',
                    '--network-prefix', '^(eth|hsn)',
                    'hello_world.py']

        args = parser.parse_args(args=arg_list)
        if args.basic_label or args.verbose_label:
            args.no_label = False
        args_map = {key: value for key, value in vars(args).items() if value is not None}

        # Give the net conf a correct object so that rest of this test can complete
        net_config_obj = NetworkConfig.from_file(self.network_config)
        mock_config.side_effect = [net_config_obj]

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()

        # Get backend up
        self.do_bringup(mock_overlay, mock_launch)

        # Receive GSProcessCreate
        handle_gsprocesscreate(self.primary_conn)

        # Send GSHalted
        handle_teardown(self.be_nodes, self.primary_conn, self.fe_ta_conn)

        # Join on the frontend thread
        fe_proc.join()

        # Check the config mock actually use slurm
        assert mock_config.called_with(workload_manager=WLM.SLURM,
                                       port=DEFAULT_OVERLAY_NETWORK_PORT,
                                       network_prefix=DEFAULT_TRANSPORT_NETIF,
                                       sigint_trigger=None)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    unittest.main()

test/launcher/test_frontend_bringup.py
DragonHPC/dragon
#!/usr/bin/env python3
import unittest
import logging
import threading
from os import environ, path
from io import StringIO
from unittest.mock import patch
from dragon.launcher.launchargs import get_parser
from dragon.launcher import launch_selector as dls

from .launcher_testing_utils import catch_thread_exceptions
from .frontend_testing_mocks import run_frontend


def get_args_map(network_config, network_prefix='^(eth|hsn)', **kwargs):

    parser = get_parser()
    arg_list = ['--wlm', 'slurm',
                '--network-config', f'{network_config}',
                '--network-prefix', network_prefix]
    for val in kwargs.values():
        arg_list = arg_list + val

    arg_list.append('hello_world.py')
    args = parser.parse_args(args=arg_list)
    if args.basic_label or args.verbose_label:
        args.no_label = False
    args_map = {key: value for key, value in vars(args).items() if value is not None}

    return args_map


class LaunchSelectionTest(unittest.TestCase):

    @patch('sys.argv', ['dragon', 'dummy.py'])
    @patch('shutil.which', return_value=None)
    def test_single_automatic_selection(self, mock_which):
        """Route to the single node launcher"""
        multi_mode = dls.determine_environment()
        self.assertFalse(multi_mode)

    @patch('sys.argv', ['dragon', 'dummy.py'])
    @patch('shutil.which')
    @patch.dict(environ, {"SLURM_JOB_ID": "3195234.0"})
    def test_multinode_automatic_selection_slurm(self, mock_which):
        """Route to the multinode launcher"""
        mock_which.side_effect = ['/opt/slurm/bin/qstat', '/usr/bin/srun', None]
        multi_mode = dls.determine_environment()
        self.assertTrue(multi_mode)

    @patch('sys.argv', ['dragon', 'dummy.py'])
    @patch('shutil.which')
    @patch.dict(environ, {"SLURM_JOB_ID": ""})
    def test_multinode_auto_without_allocation_slurm(self, mock_which):
        """Raise exception due to slurm being present with no job allocation"""
        mock_which.side_effect = ['/opt/slurm/bin/qstat', '/usr/bin/srun', None]
        with self.assertRaises(RuntimeError) as e:
            dls.determine_environment()
        self.assertTrue('Executing in a Slurm environment, but with no job allocation' in str(e.exception))

    @patch('sys.argv', ['dragon', 'dummy.py'])
    @patch('shutil.which')
    @patch.dict(environ, {"PBS_NODEFILE": ""})
    def test_multinode_auto_without_allocation_pbs_bad_mpiexec(self, mock_which):
        """Raise exception due to PBS  being present with no job allocation and a bad mpiexec in path"""
        mock_which.side_effect = ['/opt/pbs/bin/qstat', None, '/usr/bin/mpiexec']
        with self.assertRaises(RuntimeError) as e:
            dls.determine_environment()
        print(f'{str(e.exception)}')
        self.assertTrue('PBS has been detected on the system. However, Dragon is only' in str(e.exception))

    @patch('sys.argv', ['dragon', 'dummy.py'])
    @patch('shutil.which')
    @patch.dict(environ, {"PBS_NODEFILE": "/a/presumably/valid/path"})
    def test_multinode_auto_with_allocation_pbs_bad_mpiexec(self, mock_which):
        """Raise exception due to PBS being present with job allocation but a bad mpiexec in path"""
        mock_which.side_effect = ['/opt/pbs/bin/qstat', None, '/usr/bin/mpiexec']
        with self.assertRaises(RuntimeError) as e:
            dls.determine_environment()
        print(f'{str(e.exception)}')
        self.assertTrue('PBS has been detected on the system. However, Dragon is only' in str(e.exception))

    @patch('sys.argv', ['dragon', 'dummy.py'])
    @patch('shutil.which')
    @patch.dict(environ, {"PBS_NODEFILE": "/a/presumably/valid/path"})
    def test_multinode_auto_with_allocation_pbs_good_mpiexec(self, mock_which):
        """Correctly identify we're doing PBS and PALS"""
        mock_which.side_effect = ['/opt/pbs/bin/qstat', None, '/opt/cray/pe/pals/default/bin/mpiexec']
        multi_mode = dls.determine_environment()
        self.assertTrue(multi_mode)

    @patch('sys.argv', ['dragon', 'dummy.py'])
    @patch('shutil.which')
    @patch.dict(environ, {"PBS_NODEFILE": ""})
    def test_multinode_auto_without_allocation_pbs_good_mpiexec(self, mock_which):
        """Correctly identify we're doing PBS and PALS"""
        mock_which.side_effect = ['/opt/pbs/bin/qstat', None, '/opt/cray/pe/pals/default/bin/mpiexec']
        with self.assertRaises(RuntimeError) as e:
            dls.determine_environment()
        self.assertTrue('Using a supported PALS with PBS config. However, no active jobs allocation' in str(e.exception))

    @patch('sys.argv', ['dragon', '--single-node-override', 'dummy.py'])
    def test_single_node_override(self):
        """Test go to single node if requested"""
        self.assertFalse(dls.determine_environment())

    @patch('sys.argv', ['dragon', '--multi-node-override', 'dummy.py'])
    def test_multi_node_override(self):
        """Test go to multi node if requested"""
        self.assertTrue(dls.determine_environment())

    @patch('sys.stderr', new_callable=StringIO)
    @patch('sys.argv', ['dragon', '--single-node-override', '--multi-node-override', 'dummy.py'])
    def test_bad_override(self, mock_stderr):
        """Test erroneous request of multi and single node"""
        with self.assertRaises(SystemExit):
            dls.determine_environment()
        self.assertRegexpMatches(mock_stderr.getvalue(), r" not allowed with argument")


class LaunchOptionsTest(unittest.TestCase):

    def setUp(self):
        self.test_dir = path.dirname(path.realpath(__file__))
        self.network_config = path.join(self.test_dir, 'slurm_primary.yaml')

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_too_many_nodes_requested(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        """What happens when a user requests more nodes than they have in an allocation"""

        args_map = get_args_map(self.network_config,
                                arg1=['--nodes', '10010'])

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()
        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == ValueError
        assert str(exceptions_caught_in_threads['Frontend Server']['exception']['value']) \
            == 'Not enough backend nodes allocated to match requested'

    @patch('sys.stderr', new_callable=StringIO)
    def test_nonint_nodes_requested(self, mock_stderr):
        """what happens when users passes a non-int for # nodes"""

        with self.assertRaises(SystemExit):
            get_args_map(self.network_config,
                         arg1=['--nodes', 'abcdef'])
        self.assertTrue("invalid non_negative_int value: 'abcdef'" in mock_stderr.getvalue())

    @patch('sys.stderr', new_callable=StringIO)
    def test_bad_port_requested(self, mock_stderr):
        """what happens when user passes a bad value to port"""

        with self.assertRaises(SystemExit):
            get_args_map(self.network_config,
                         arg1=['--port', '23'])
        self.assertTrue("must be in port range 1024-65535" in mock_stderr.getvalue())

        with self.assertRaises(SystemExit):
            get_args_map(self.network_config,
                         arg1=['--port', '799999'])
        self.assertTrue("must be in port range 1024-65535" in mock_stderr.getvalue())

        with self.assertRaises(SystemExit):
            get_args_map(self.network_config,
                         arg1=['--port', 'abcdef'])
        self.assertTrue("invalid valid_port_int value" in mock_stderr.getvalue())

    @catch_thread_exceptions
    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_bad_network_prefix(self, exceptions_caught_in_threads, mock_overlay, mock_launch):
        """What happens when user requests a network prefix that doesn't exist"""
        args_map = get_args_map(self.network_config,
                                network_prefix='garbage-prefix')

        # get startup going in another thread. Note: need to do threads in order to use
        # all our mocks
        fe_proc = threading.Thread(name='Frontend Server',
                                   target=run_frontend,
                                   args=(args_map,),
                                   daemon=False)
        fe_proc.start()
        fe_proc.join()

        assert 'Frontend Server' in exceptions_caught_in_threads  # there was an exception in thread  1
        assert exceptions_caught_in_threads['Frontend Server']['exception']['type'] == ValueError
        self.assertTrue('No IP addresses found for' in
                        str(exceptions_caught_in_threads['Frontend Server']['exception']['value']))
        self.assertTrue('matching regex pattern: garbage-prefix' in
                        str(exceptions_caught_in_threads['Frontend Server']['exception']['value']))


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    unittest.main()

test/launcher/test_launch_options.py
DragonHPC/dragon
#!/usr/bin/env python3
import unittest
import logging
import subprocess
import json
from os import environ, path
import random
import string


from dragon.launcher.network_config import WLM, ConfigOutputType, NetworkConfig
from dragon.infrastructure.node_desc import NodeDescriptor
from dragon.infrastructure.facts import DEFAULT_TRANSPORT_NETIF, DEFAULT_OVERLAY_NETWORK_PORT


class FileNetworkConfigTest(unittest.TestCase):

    def setUp(self):
        '''Read a network config from file'''
        base_abs_path = path.dirname(__file__)
        self.bad_yaml = path.join(base_abs_path, "slurm_bad.yaml")
        self.good_yaml = path.join(base_abs_path, "slurm.yaml")
        self.good_primary_yaml = path.join(base_abs_path, "slurm_primary.yaml")
        self.bad_json = path.join(base_abs_path, "slurm_bad.json")
        self.good_json = path.join(base_abs_path, "slurm.json")

    def test_good_input_yaml_parser(self):
        net_conf = NetworkConfig.from_file(self.good_yaml, ConfigOutputType.YAML)
        self.assertIsInstance(net_conf, NetworkConfig)

        config = net_conf.get_network_config()
        for key, value in config.items():
            self.assertIsInstance(key, str)
            self.assertIsInstance(value, NodeDescriptor)

    def test_good_input_json_parser(self):
        net_conf = NetworkConfig.from_file(self.good_json, ConfigOutputType.JSON)
        self.assertIsInstance(net_conf, NetworkConfig)

        config = net_conf.get_network_config()
        for key, value in config.items():
            self.assertIsInstance(key, str)
            self.assertIsInstance(value, NodeDescriptor)

    def test_bad_input_yaml_parser(self):
        with self.assertRaises(TypeError) as e:
            NetworkConfig.from_file(self.bad_yaml, ConfigOutputType.YAML)
        self.assertTrue("got an unexpected keyword argument 'junk-key'" in str(e.exception))

    def test_bad_input_json_parser(self):
        with self.assertRaises(TypeError) as e:
            NetworkConfig.from_file(self.bad_json, ConfigOutputType.JSON)
        self.assertTrue("got an unexpected keyword argument 'junk-key'" in str(e.exception))

    def test_primary_from_file_selection(self):
        '''Confirm we behave if primary is given in file'''
        net = NetworkConfig.from_file(self.good_primary_yaml, ConfigOutputType.YAML)
        self.assertIsInstance(net, NetworkConfig)

        conf = net.get_network_config()
        primary_keys = [k for k, v in conf.items() if conf[k].is_primary]

        # Should only be 1
        self.assertEqual(len(primary_keys), 1)

        # YAML has node index 1 selected
        self.assertEqual(primary_keys[0], '1')

    def test_no_primary_from_file_selection(self):
        '''Confirm selection by host id works correctly'''
        net = NetworkConfig.from_file(self.good_yaml, ConfigOutputType.YAML)
        self.assertIsInstance(net, NetworkConfig)

        conf = net.get_network_config()
        primary_keys = [k for k, v in conf.items() if conf[k].is_primary]

        # Should only be 1
        self.assertEqual(len(primary_keys), 1)

        # '0' has the smallest host ID and should be selected as primary
        self.assertEqual(primary_keys[0], '0')

    def test_no_primary_with_input_from_file_selection(self):
        '''Confirm we can override default behavior with specific hostname'''
        prime_node = 'nid00006'
        net = NetworkConfig.from_file(self.good_yaml, ConfigOutputType.YAML,
                                      primary_hostname=prime_node)
        self.assertIsInstance(net, NetworkConfig)

        conf = net.get_network_config()
        primary_keys = [k for k, v in conf.items() if conf[k].is_primary]

        # Should only be 1
        self.assertEqual(len(primary_keys), 1)

        # Every primary node should be at '0', but we can
        # confirm the correct name
        self.assertEqual(conf['0'].name, prime_node)

    def test_primary_with_input_from_file_selection(self):
        '''Confirm we can override always ignore primary hostname if on in config file specific hostname'''
        with self.assertRaises(RuntimeError) as e:
            NetworkConfig.from_file(self.good_primary_yaml, ConfigOutputType.YAML,
                                    primary_hostname='nid00006')
        self.assertTrue('Primary hostname input by user but is already set in network config file' in str(e.exception))

        with self.assertRaises(RuntimeError) as e:
            NetworkConfig.from_file(self.good_primary_yaml, ConfigOutputType.YAML,
                                    primary_hostname='nid0000654')
        self.assertTrue('Primary hostname input by user but is already set in network config file' in str(e.exception))

    def test_nonexistent_primary_node(self):
        '''Confirm we can raise an error if requested primary doesn't exist and one isn't already selected'''
        with self.assertRaises(RuntimeError) as e:
            NetworkConfig.from_file(self.good_yaml, ConfigOutputType.YAML,
                                    primary_hostname='nid0000654')
        self.assertTrue('Input hostname does not match any available' in str(e.exception))


class SlurmNetworkConfigTest(unittest.TestCase):

    def setUp(self):
        '''Set values specific to slurm'''
        self.wlm = WLM.SLURM
        self.config_files = {}

    def test_slurm(self):

        # Only run this test if we have a slurm allocation
        if not environ.get("SLURM_JOB_ID"):
            logging.info("Slurm test is checking for error")
            with self.assertRaises(RuntimeError):
                net = NetworkConfig.from_wlm(workload_manager=WLM.SLURM,
                                             port=DEFAULT_OVERLAY_NETWORK_PORT,
                                             network_prefix=DEFAULT_TRANSPORT_NETIF)
        else:
            logging.info("slurm test launched backend config jobs")
            net = NetworkConfig.from_wlm(workload_manager=WLM.SLURM,
                                         port=DEFAULT_OVERLAY_NETWORK_PORT,
                                         network_prefix=DEFAULT_TRANSPORT_NETIF)
            self.assertIsInstance(net, NetworkConfig)

            config = net.get_network_config()
            for key, value in config.items():
                self.assertIsInstance(key, str)
                self.assertIsInstance(value, NodeDescriptor)

    def test_slurm_stdout(self):
        args = ['python3',
                '-m',
                'dragon.launcher.network_config',
                '--wlm',
                'slurm']
        if not environ.get("SLURM_JOB_ID"):
            logging.info("Slurm test is checking for error")
            try:
                proc = subprocess.run(args=args, capture_output=True)
            except Exception as e:
                isinstance(e, RuntimeError)
        else:
            logging.info("slurm test launched backend config jobs")
            proc = subprocess.run(args=args, timeout=30, capture_output=True, text=True)
            net = NetworkConfig.from_sdict(json.loads(proc.stdout))
            self.assertIsInstance(net, NetworkConfig)

            config = net.get_network_config()
            for key, value in config.items():
                self.assertIsInstance(key, str)
                self.assertIsInstance(value, NodeDescriptor)

    def test_slurm_bad_network_prefix(self):
        args = ['python3',
                '-m',
                'dragon.launcher.network_config',
                '--wlm', 'slurm',
                '--network-prefix', 'garbage']

        if environ.get("SLURM_JOB_ID"):
            proc = subprocess.run(args=args, capture_output=True, text=True)
            self.assertNotEqual(proc.returncode, 0)
            self.assertTrue('ValueError: No IP addresses found for' in proc.stderr)
            self.assertTrue(' matching regex pattern: garbage' in proc.stderr)


class PbsPalsNetworkConfigTest(unittest.TestCase):

    def setUp(self):
        '''Set values specific to PBS+PALS'''
        self.wlm = WLM.PBS_PALS
        self.config_files = {}

    def test_pbs_pals(self):

        # Only run this test if we have a slurm allocation
        if not environ.get("PBS_NODEFILE"):
            logging.info("PBS+PALS test is checking for error")
            with self.assertRaises(RuntimeError):
                net = NetworkConfig.from_wlm(workload_manager=WLM.PBS_PALS,
                                             port=DEFAULT_OVERLAY_NETWORK_PORT,
                                             network_prefix=DEFAULT_TRANSPORT_NETIF)
        else:
            logging.info("PBS+PALS test launched backend config jobs")
            net = NetworkConfig.from_wlm(workload_manager=WLM.PBS_PALS,
                                         port=DEFAULT_OVERLAY_NETWORK_PORT,
                                         network_prefix=DEFAULT_TRANSPORT_NETIF)
            self.assertIsInstance(net, NetworkConfig)

            config = net.get_network_config()
            for key, value in config.items():
                self.assertIsInstance(key, str)
                self.assertIsInstance(value, NodeDescriptor)

    def test_pbs_pals_stdout(self):
        args = ['python3',
                '-m',
                'dragon.launcher.network_config',
                '--wlm',
                'pbs+pals']
        if not environ.get("PBS_NODEFILE"):
            logging.info("PBS+PALS test is checking for error")
            try:
                proc = subprocess.run(args=args, capture_output=True)
            except Exception as e:
                isinstance(e, RuntimeError)
        else:
            logging.info("PBS+PALS test launched backend config jobs")
            proc = subprocess.run(args=args, timeout=30, capture_output=True, text=True)
            net = NetworkConfig.from_sdict(json.loads(proc.stdout))
            self.assertIsInstance(net, NetworkConfig)

            config = net.get_network_config()
            for key, value in config.items():
                self.assertIsInstance(key, str)
                self.assertIsInstance(value, NodeDescriptor)

    def test_pbs_pals_bad_network_prefix(self):
        args = ['python3',
                '-m',
                'dragon.launcher.network_config',
                '--wlm', 'pbs+pals',
                '--network-prefix', 'garbage']

        if environ.get("PBS_NODEFILE"):
            proc = subprocess.run(args=args, capture_output=True, text=True)
            self.assertNotEqual(proc.returncode, 0)
            self.assertTrue('ValueError: No IP addresses found for' in proc.stderr)
            self.assertTrue(' matching regex pattern: garbag' in proc.stderr)


class SSHNetworkConfigTest(unittest.TestCase):

    def setUp(self):
        '''Set values specific to slurm'''

        base_abs_path = path.dirname(__file__)
        self.bad_hostfile = path.join(base_abs_path, "bad_hostfile.txt")
        self.good_hostfile = path.join(base_abs_path, "good_hostfile.txt")

        self.wlm = WLM.SSH
        self.config_files = {}

    def test_host_parsing(self):
        '''test the config tool parses input arguments correctly'''

        from dragon.launcher.network_config import get_args
        # Test that we error out if specify hostfile and hostlist
        input_args = ['--wlm', 'ssh',
                      '--hostfile', self.good_hostfile,
                      '--hostlist', 'host1,host2']
        with self.assertRaises(SystemExit):
            get_args(inputs=input_args)

        # Test we error when set with a different type of wlm
        input_args = ['--wlm', 'slurm',
                      '--hostfile', self.good_hostfile]
        with self.assertRaises(NotImplementedError) as e:
            get_args(inputs=input_args)
        self.assertTrue('hostlist and hostfile arguments are only supported in the SSH case' in str(e.exception))

        # test that we throw an error when hostfile is not formatted correctly
        input_args = ['--wlm', 'ssh',
                      '--hostfile', self.bad_hostfile]
        with self.assertRaises(ValueError) as e:
            get_args(inputs=input_args)
        self.assertTrue('Hostname is invalid:' in str(e.exception))

        # test that we like a good hostfile
        input_args = ['--wlm', 'ssh',
                      '--hostfile', self.good_hostfile]

        wlm, args = get_args(inputs=input_args)
        self.assertEqual(args.hostlist[0], 'host-1')
        self.assertEqual(args.hostlist[1], 'host-2')

        # test that we handle good hostlists
        input_args = ['--wlm', 'ssh',
                      '--hostlist', 'host-1,host-2']
        wlm, args = get_args(inputs=input_args)
        self.assertEqual(args.hostlist[0], 'host-1')
        self.assertEqual(args.hostlist[1], 'host-2')

        # Test that white space after comma is okay
        input_args = ['--wlm', 'ssh',
                      '--hostlist', 'host-1,  host-2']
        wlm, args = get_args(inputs=input_args)
        self.assertEqual(args.hostlist[0], 'host-1')
        self.assertEqual(args.hostlist[1], 'host-2')

        # Test that a psychotic tab is okay after comma is okay
        input_args = ['--wlm', 'ssh',
                      '--hostlist', 'host-1,\thost-2']
        wlm, args = get_args(inputs=input_args)
        self.assertEqual(args.hostlist[0], 'host-1')
        self.assertEqual(args.hostlist[1], 'host-2')

        # Test that we invalidate too long hostname
        input_args = ['--wlm', 'ssh',
                      '--hostlist', ''.join(random.choices(string.ascii_letters, k=300))]
        with self.assertRaises(ValueError) as e:
            wlm, args = get_args(inputs=input_args)
        self.assertTrue('Hostname is invalid:' in str(e.exception))

        too_long_part = ''.join(random.choices(string.ascii_letters, k=64)) + '.' + ''.join(random.choices(string.ascii_letters, k=20))
        input_args = ['--wlm', 'ssh',
                      '--hostlist', too_long_part]
        with self.assertRaises(ValueError) as e:
            wlm, args = get_args(inputs=input_args)
        self.assertTrue('Hostname is invalid:' in str(e.exception))


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    unittest.main()

test/launcher/test_network_config.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import logging
import unittest
from unittest.mock import patch
import threading
from time import time

from dragon.launcher.launch_selector import determine_environment
from dragon.launcher.frontend import LauncherFrontEnd
from dragon.launcher.util import next_tag, SRQueue, get_with_timeout
from dragon.launcher.launchargs import get_args

from dragon.infrastructure.process_desc import ProcessDescriptor
from dragon.infrastructure import facts as dfacts
from dragon.infrastructure import messages as dmsg
import dragon.utils as du

from dragon.managed_memory import DragonPoolError, DragonMemoryError
from dragon.channels import Channel, ChannelError
from dragon.utils import set_host_id, B64
from dragon.dlogging.util import setup_FE_logging

from .frontend_testing_mocks import open_overlay_comms, open_backend_comms


def send_beisup(nodes):
    '''Send valid beisup messages'''

    for host_id, node in nodes.items():
        be_up_msg = dmsg.BEIsUp(tag=next_tag(),
                                be_ch_desc=B64.bytes_to_str(node['ch_in'].serialize()),
                                host_id=host_id)
        node['conn'].send(be_up_msg.serialize())


def recv_fenodeidx(nodes):
    '''recv FENoe4deIdxBE and finish filling out node dictionary'''
    log = logging.getLogger('recv_fe_nodeidx')
    for node in nodes.values():
        fe_node_idx_msg = dmsg.parse(node['conn'].recv())
        assert isinstance(fe_node_idx_msg, dmsg.FENodeIdxBE), 'la_be node_index from fe expected'

        log.info(f'got FENodeIdxBE for index {fe_node_idx_msg.node_index}')
        node['node_index'] = fe_node_idx_msg.node_index
        if node['node_index'] < 0:
            raise RuntimeError("frontend giving bad node indices")
        node['is_primary'] = node['node_index'] == 0
        if node['is_primary']:
            primary_conn = node['conn']
        log.info(f'constructed be node: {node}')

    return primary_conn


def send_shchannelsup(nodes, mpool):

    log = logging.getLogger('send_shchannelsup')
    for host_id, node in nodes.items():
        ls_cuid = dfacts.shepherd_cuid_from_index(node['node_index'])
        ls_ch = Channel(mpool, ls_cuid)
        node['ls_ch'] = ls_ch

        if node['is_primary']:
            node['gs_ch'] = Channel(mpool, dfacts.GS_INPUT_CUID)
            gs_cd = du.B64.bytes_to_str(node['gs_ch'].serialize())
        else:
            node['gs_ch'] = None
            gs_cd = None

        ch_up_msg = dmsg.SHChannelsUp(tag=next_tag(),
                                      host_name=node['hostname'],
                                      host_id=host_id,
                                      ip_addrs=node['ip_addrs'],
                                      shep_cd=B64.bytes_to_str(node['ls_ch'].serialize()),
                                      gs_cd=gs_cd,
                                      idx=node['node_index'])
        log.info(f'construct SHChannelsUp: {ch_up_msg}')
        node['conn'].send(ch_up_msg.serialize())
        log.info(f'sent SHChannelsUp for {node["node_index"]}')

    log.info('sent all SHChannelsUp')


def recv_lachannelsinfo(nodes):
    '''Loop to recv all LAChannelsInfo messages in frontend bcast'''
    for host_id, node in nodes.items():
        la_channels_info_msg = dmsg.parse(node['conn'].recv())
        assert isinstance(la_channels_info_msg, dmsg.LAChannelsInfo), 'la_be expected all ls channels info from la_fe'


def send_taup(nodes):
    for host_id, node in nodes.items():
        ta_up = dmsg.TAUp(tag=next_tag(), idx=node['node_index'])
        node['conn'].send(ta_up.serialize())


def handle_gsprocesscreate(primary_conn, proc_create):
    '''Manage a valid response to GSProcessCreate'''

    # Send response
    gs_desc = ProcessDescriptor(p_uid=5000,  # Just a dummy value
                                name=proc_create.user_name,
                                node=0,
                                p_p_uid=proc_create.p_uid)
    response = dmsg.GSProcessCreateResponse(tag=next_tag(),
                                            ref=proc_create.tag,
                                            err=dmsg.GSProcessCreateResponse.Errors.SUCCESS,
                                            desc=gs_desc)
    primary_conn.send(response.serialize())

    # And go ahead and immediately exit because we're not actually doing anything
    primary_conn.send(dmsg.GSHeadExit(exit_code=0, tag=next_tag()).serialize())


def get_args_map(network_config, from_wlm=False):

    if from_wlm:
        arg_list = ['-t', 'tcp',
                    'launcher/helloworld.py']
    else:
        arg_list = ['--wlm', 'slurm',
                    '--network-config', f'{network_config}',
                    '--network-prefix', "",
                    'helloworld.py']
    args_map = get_args(arg_list)
    return args_map


def cleanup_mocks(fe_ta_conn,
                  be_ch_out, be_ch_in, be_nodes,
                  ls_ch,
                  ta_ch_out, ta_ch_in,
                  be_mpool):

    try:
        fe_ta_conn.close()
    except (AttributeError, ChannelError, ConnectionError):
        pass

    try:
        be_ch_out.detach()
    except (AttributeError, ChannelError):
        pass

    try:
        for node in be_nodes.values():
            try:
                node['conn'].close()
            except (KeyError, AttributeError, ConnectionError):
                pass

            try:
                node['ch_in'].destroy()
            except (KeyError, AttributeError, ChannelError):
                pass
    except AttributeError:
        pass

    try:
        for node in be_nodes.values():
            try:
                node['ls_ch'].destroy()
                if node['is_primary']:
                    node['gs_ch'].destroy()
            except (KeyError, AttributeError, ChannelError):
                pass
    except AttributeError:
        pass

    try:
        ta_ch_out.detach()
    except (AttributeError, ChannelError):
        pass

    try:
        ta_ch_in.detach()
    except (AttributeError, ChannelError):
        pass

    try:
        be_mpool.destroy()
    except (AttributeError, DragonMemoryError, DragonPoolError):
        pass


def run_frontend_supporting_mocks(mock_overlay=None,
                                  mock_launch=None,
                                  args_map=None,
                                  overlay_only=False,
                                  no_proc_create=False,
                                  hang_backend=False,
                                  hang_overlay=False,
                                  exit_at_fenodeidx=False,
                                  exit_queue=None):

    log = logging.getLogger('frontend mocks')

    be_mpool = None
    be_ch_out = None
    be_ch_in = None
    overlay_inout = None

    ls_ch = None

    ta_ch_in = None
    ta_ch_out = None
    fe_ta_conn = None
    be_nodes = None

    log.info("inside frontend mocks")
    # Get the mock's input args to the
    try:

        while mock_overlay.call_args is None:
            pass
        log.info("getting mock overlay args")
        overlay_args = mock_overlay.call_args.kwargs

        # Connect to overlay comms to talk to fronteend
        log.info("opening overlay comms overlay args")
        ta_ch_in, ta_ch_out, fe_ta_conn = open_overlay_comms(overlay_args['ch_in_sdesc'],
                                                             overlay_args['ch_out_sdesc'])

        # Let frontend know the overlay is "up"
        log.info("sending tapingsh")
        fe_ta_conn.send(dmsg.OverlayPingLA(next_tag()).serialize())

        if overlay_only:
            halt_on = dmsg.parse(fe_ta_conn.recv())
            assert isinstance(halt_on, dmsg.LAHaltOverlay)
            fe_ta_conn.send(dmsg.OverlayHalted(tag=next_tag()).serialize())

            raise Exception('exiting due to return_after_overlay_up being True')

        # Grab the frontend channel descriptor for the launched backend and
        # send it mine
        while mock_launch.call_args is None:
            pass
        launch_be_args = mock_launch.call_args.kwargs
        log.info(f'got be args: {launch_be_args}')

        mock_launch.wait.return_value = None
        log.info("set srun launch wait to None")

        # Connect to backend comms for frontend-to-backend and back comms
        be_mpool, be_ch_out, be_ch_in, be_nodes, overlay_inout = open_backend_comms(launch_be_args['frontend_sdesc'],
                                                                                  args_map['network_config'])
        log.info('got backend up')

        # Send BEIsUp
        send_beisup(be_nodes)
        log.info('send BEIsUp messages')

        # Recv FENodeIdxBE
        if not exit_at_fenodeidx:
            primary_conn = recv_fenodeidx(be_nodes)
            log.info('got all the FENodeIdxBE messages')
        else:
            log.info('told to return at FENodeIdx recv')
            cleanup_mocks(fe_ta_conn, be_ch_out, be_ch_in, be_nodes,
                          ls_ch, ta_ch_out, ta_ch_in, be_mpool)
            return

        # Fudge some SHChannelsUp messages
        send_shchannelsup(be_nodes, be_mpool)
        log.info(f'sent shchannelsup: {[node["gs_ch"].serialize() for node in be_nodes.values() if node["gs_ch"] is not None]}')

        # Receive LAChannelsInfo
        recv_lachannelsinfo(be_nodes)
        log.info('la_be received LAChannelsInfo')

        # Send TAUp
        send_taup(be_nodes)
        log.info('sent TAUp messages')

        # Send gs is up from primary
        primary_conn.send(dmsg.GSIsUp(tag=next_tag()).serialize())
        log.info('send GSIsUp')

        while True:
            gs_msg = None
            try:
                log.debug('doing gs get')
                gs_msg = get_with_timeout(primary_conn, timeout=0.01)
            except TimeoutError:
                log.debug('gs timeout')
                pass

            try:
                log.debug('doing exit get')
                _ = get_with_timeout(exit_queue, timeout=0.01)
                log.debug('cleaning up mocks')
                cleanup_mocks(fe_ta_conn, be_ch_out, be_ch_in, be_nodes,
                              ls_ch, ta_ch_out, ta_ch_in, be_mpool)
                return
            except TimeoutError:
                log.debug('exit timeout')
                pass

            if isinstance(gs_msg, dmsg.GSProcessCreate):
                log.info("handling GSProcessCreate on la_be")
                handle_gsprocesscreate(primary_conn, gs_msg)
            elif isinstance(gs_msg, dmsg.GSTeardown):
                log.info('la_be received GSTeardown')
                primary_conn.send(dmsg.GSHalted(tag=next_tag()).serialize())
                break

        nhalted = 0
        goal_halted = len(be_nodes)
        while True:

            for node in be_nodes.values():
                lmsg = None
                log.debug('posting mock recv')
                try:
                    lmsg = get_with_timeout(node['conn'], timeout=0.01)
                except TimeoutError:
                    pass

                try:
                    _ = get_with_timeout(exit_queue, timeout=0.01)
                    cleanup_mocks(fe_ta_conn, be_ch_out, be_ch_in, be_nodes,
                                  ls_ch, ta_ch_out, ta_ch_in, be_mpool)
                    return
                except TimeoutError:
                    pass

                if isinstance(lmsg, dmsg.SHHaltTA):
                    log.info('recvd SHHaltTA')
                    node['conn'].send(dmsg.TAHalted(tag=next_tag()).serialize())
                elif isinstance(lmsg, dmsg.SHTeardown):
                    log.info('recvd SHTeardown')
                    if hang_backend:
                        # Wait to tear everything down until I know I won't make everything go
                        # crazy
                        log.debug('waiting on exit queue signal')
                        while True:
                            log.debug('posting exit queue recv')
                            la_exit = exit_queue.recv(timeout=1)
                            if la_exit is not None:
                                break
                        log.debug('exit queue signal received')
                        cleanup_mocks(fe_ta_conn, be_ch_out, be_ch_in, be_nodes,
                                      ls_ch, ta_ch_out, ta_ch_in, be_mpool)
                        mock_launch.wait.return_value = 2
                        log.debug('returning early from mocks')
                        return
                    else:
                        node['conn'].send(dmsg.SHHaltBE(tag=next_tag()).serialize())

                elif isinstance(lmsg, dmsg.BEHalted):
                    log.info("recvd BEHalted")
                    nhalted += 1
                    log.info(f'inner nhalted ({nhalted}) v. goal_halted ({goal_halted})')
                    if nhalted == goal_halted:
                        log.info('breaking out of first BEHalted')
                        break

            log.info(f'nhalted ({nhalted}) v. goal_halted ({goal_halted})')
            if nhalted == goal_halted:
                log.info('breaking out of second BEHalted')
                break

        # Handle the final TAHalted for overlay
        log.info('posting recv for FE overlay')
        halt_on = dmsg.parse(fe_ta_conn.recv())
        assert isinstance(halt_on, dmsg.LAHaltOverlay), 'was excpecting an SHHaltTA'
        if hang_overlay:
            # Wait to tear everything down until I know I won't make everything go
            # crazy
            log.debug('waiting on exit queue signal')
            _ = exit_queue.recv()
            log.debug('exit queue signal received')
            cleanup_mocks(fe_ta_conn, be_ch_out, be_ch_in, be_nodes,
                          ls_ch, ta_ch_out, ta_ch_in, be_mpool)
            mock_launch.wait.return_value = 2
            log.debug('breaking out of mocks before halting overlay')
            return
        else:
            fe_ta_conn.send(dmsg.OverlayHalted(tag=next_tag()).serialize())
    except Exception as e:
        log.info(f'hit frontend mocks exception: {e}')
        pass

        if mock_launch is not None:
            mock_launch.wait.return_value = None

    cleanup_mocks(fe_ta_conn, be_ch_out, be_ch_in, be_nodes,
                  ls_ch, ta_ch_out, ta_ch_in, be_mpool)

    # Let the test know we're done:
    if mock_launch is not None:
        mock_launch.wait.return_value = 0

    log.info('returning from frontends mock')


class SigIntTest(unittest.TestCase):

    def setUp(self):

        self.test_dir = os.path.dirname(os.path.realpath(__file__))
        self.network_config = os.path.join(self.test_dir, 'slurm_primary.yaml')

        self.be_mpool = None
        self.be_ch_out = None
        self.be_ch_in = None
        self.overlay_inout = None

        self.ls_ch = None

        self.ta_ch_in = None
        self.ta_ch_out = None
        self.fe_ta_conn = None

        self.args_map = get_args_map(self.network_config, from_wlm=False)
        self.wlm_args_map = get_args_map(self.network_config, from_wlm=True)

        self.exit_queue = SRQueue()

        # Some tests will only work in a multinode environment with
        # an activate allocation
        try:
            self.multi_mode = determine_environment(self.wlm_args_map)
        except RuntimeError:
            self.multi_mode = False
            pass

        # Before doing anything set my host ID
        set_host_id(dfacts.FRONTEND_HOSTID)
        setup_FE_logging(log_device_level_map=self.args_map['log_device_level_map'],
                         basename='dragon', basedir=os.getcwd())
        log = logging.getLogger('sigint setUp')
        for key, value in self.args_map.items():
            if value is not None:
                log.info(f'args_map: {key}: {value}')

    def tearDown(self):

        if self.fe_ta_conn is not None:
            self.fe_ta_conn.close()

        if self.be_ch_out is not None:
            self.be_ch_out.detach()

        if self.be_ch_in is not None:
            for node in self.be_nodes.values():
                node['conn'].close()
                node['ch_in'].destroy()

        if self.ls_ch is not None:
            for node in self.be_nodes.values():
                node['ls_ch'].destroy()

                if node['is_primary']:
                    node['gs_ch'].destroy()

        if self.ta_ch_out is not None:
            self.ta_ch_out.detach()

        if self.ta_ch_in is not None:
            self.ta_ch_in.detach()

        if self.be_mpool is not None:
            self.be_mpool.destroy()

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_clean_exit(self, mock_overlay, mock_be_launch):
        '''Test a clean bring-up and teardown'''

        # Start my mocks support the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'no_proc_create': False,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        fe_server = LauncherFrontEnd(args_map=self.args_map)
        fe_server.run_startup()
        fe_server.run_app()
        fe_server.run_msg_server()

        mock_procs.join()

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_0(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 0 in launcher'''

        sigint_trigger = 0
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_1(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 1 in launcher'''

        sigint_trigger = 1
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

    def test_sig_1_no_mock(self):
        '''Test SIGINT with full runtime being raised at case 1 in launcher'''

        if self.multi_mode:
            sigint_trigger = 1
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    def test_sig_minus_1(self):
        '''Test SIGINT with full runtime being raised at case -1 in launcher'''

        if self.multi_mode:

            sigint_trigger = -1
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)
            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()

        else:
            print("Unable to run. Requires WLM job allocation")

    def test_sig_minus_2(self):
        '''Test SIGINT with full runtime being raised at case -2 in launcher'''
        if self.multi_mode:
            sigint_trigger = -2
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()

        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_2(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 2 in launcher'''

        sigint_trigger = 2
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()

    def test_sig_2_no_mock(self):
        '''Test SIGINT with full runtime being raised at case 2 in launcher'''
        if self.multi_mode:
            sigint_trigger = 2
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_3(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 3 in launcher'''

        # Start my mocks support the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': True,
                     'no_proc_create': False,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 3
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()

        mock_procs.join()

    def test_sig_3_no_mock(self):
        '''Test SIGINT with full runtime being raised at case 3 in launcher'''
        if self.multi_mode:
            sigint_trigger = 3
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_4(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 4 in launcher'''

        # Start my mocks support the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': True,
                     'no_proc_create': False,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 4
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

        mock_procs.join()

    def test_sig_4_no_mock(self):
        '''Test SIGINT with full runtime being raised at case 4 in launcher'''
        if self.multi_mode:
            sigint_trigger = 4
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_5(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 5 in launcher'''

        # Start my mocks support the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'no_proc_create': True,
                     'exit_at_fenodeidx': True,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 5
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

        self.exit_queue.send(dmsg.LAExit(tag=next_tag()).serialize())
        mock_procs.join()

    def test_sig_5_no_mock(self):
        '''Test SIGINT with full runtime being raised at case 5 in launcher'''
        if self.multi_mode:
            sigint_trigger = 5
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_6(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 6 in launcher'''

        # Start my mocks support the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'no_proc_create': False,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 6
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()
        log = logging.getLogger('test_sig_5')
        self.exit_queue.send(dmsg.LAExit(tag=next_tag()).serialize())
        mock_procs.join()

    def test_sig_6_no_mock(self):
        '''Test SIGINT with full runtime being raised at case 6 in launcher'''
        if self.multi_mode:
            sigint_trigger = 6
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sig_7(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 7 in launcher'''

        # Start my mocks support the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'no_proc_create': False,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 7
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        log = logging.getLogger('test_sig_7')
        log.debug('entering main test')
        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

        self.exit_queue.send(dmsg.LAExit(tag=next_tag()).serialize())
        mock_procs.join()

    def test_sig_7_no_mock(self):
        '''Test SIGINT with full runtime being raised at case 7 in launcher'''
        if self.multi_mode:
            sigint_trigger = 7
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sigint_hung_backend(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 7 in launcher with a hanging backend'''

        # Start my mocks supporting the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'no_proc_create': False,
                     'hang_backend': True,
                     'hang_overlay': False,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 7
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

        # Let the mock thread know to cleanup:
        log = logging.getLogger('test_sigint_hung_backend')
        log.debug('sending exit signal')
        self.exit_queue.send(dmsg.LAExit(tag=next_tag()).serialize())

        mock_procs.join()

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_sigint_hung_overlay(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 7 in launcher with a hanging overlay'''

        # Start my mocks supporting the frontend in a background thread
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'no_proc_create': False,
                     'hang_backend': False,
                     'hang_overlay': True,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 7
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

        # Let the mock thread know to cleanup:
        self.exit_queue.send(dmsg.LAExit(tag=next_tag()).serialize())

        mock_procs.join()

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_rapid_sigint(self, mock_overlay, mock_be_launch):
        '''Test SIGINT being raised at case 8 -- handles 2 SIGINT signals transmitted'''

        # Start my mocks support the frontend in a background thread()
        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'no_proc_create': True,
                     'exit_at_fenodeidx': True,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 8
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        start = time()
        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

        # this test should exit much more quickly since we're triggering the
        # quick exit via 2 SIGINTs and should be faster than the 5 second
        # default timeout
        self.assertLessEqual(time() - start, 5.0)

        self.exit_queue.send(dmsg.LAExit(tag=next_tag()).serialize())
        mock_procs.join()

    def test_rapid_sigint_no_mock(self):
        '''Test SIGINT being raised at case 8 -- handles 2 SIGINT signals transmitted. No mocks'''

        if self.multi_mode:
            sigint_trigger = 8
            fe_server = LauncherFrontEnd(args_map=self.wlm_args_map,
                                         sigint_trigger=sigint_trigger)

            with self.assertRaises(KeyboardInterrupt):
                fe_server.run_startup()
                fe_server.run_app()
                fe_server.run_msg_server()
        else:
            print("Unable to run. Requires WLM job allocation")

    @patch('dragon.launcher.frontend.LauncherFrontEnd._launch_backend')
    @patch('dragon.launcher.frontend.start_overlay_network')
    def test_teardown_with_hung_backend_sigint(self, mock_overlay, mock_be_launch):
        '''Test SIGINT during a clean teardown to trigger a quick exit'''

        mock_args = {'mock_overlay': mock_overlay,
                     'mock_launch': mock_be_launch,
                     'args_map': self.args_map,
                     'overlay_only': False,
                     'hang_backend': True,
                     'exit_queue': self.exit_queue}

        mock_procs = threading.Thread(name='Frontend Supporting Mocks',
                                      target=run_frontend_supporting_mocks,
                                      kwargs=mock_args,
                                      daemon=False)

        mock_procs.start()

        sigint_trigger = 9
        fe_server = LauncherFrontEnd(args_map=self.args_map,
                                     sigint_trigger=sigint_trigger)

        start = time()
        with self.assertRaises(KeyboardInterrupt):
            fe_server.run_startup()
            fe_server.run_app()
            fe_server.run_msg_server()

        self.assertLessEqual(time() - start, 5.0)

        # Let the mock thread know to cleanup:
        self.exit_queue.send(dmsg.LAExit(tag=next_tag()).serialize())

        mock_procs.join()


if __name__ == "__main__":
    unittest.main()

test/launcher/test_signal_handling.py
DragonHPC/dragon
import os
import dragon
import multiprocessing as mp

def runit():
    print("Hello from Worker", os.uname().nodename)

def main():
    mp.set_start_method('dragon')
    print("Hello Dragon", os.uname().nodename)
    p = mp.Process(target=runit, args=())
    p.start()
    p.join()

if __name__ == "__main__":
    main()
test/launcher_multi/helloworld.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import socket
import time


def main():
    """Open a file, write, sleep, and exit

    Just a simple test to confirm local services is
    capable of bringing up a serial process on a given compute node.

    """
    _user = os.environ.get('USER', str(os.getuid()))
    _hostname = socket.gethostname()
    try:
        filename = os.path.join(os.environ.get('DRAGON_LA_LOG_DIR'),
                                f'simple_compute_{_user}.txt')
    except TypeError:
        filename = f'/tmp/simple_compute_{_user}.txt'

    if os.path.exists(filename):
        os.remove(filename)

    for i in range(5):
        with open(filename, 'a') as f:
            f.write(f'hello from {_hostname}\n')
        time.sleep(5)


if __name__ == "__main__":
    main()

test/launcher_multi/serial_compute.py
DragonHPC/dragon
import subprocess
import json
from pprint import pprint


def get_ip_addresses():

    try:
        local_args = ['python3',
                      '-m',
                      'dragon.transport.ifaddrs',
                      '--ip',
                      '--no-loopback',
                      '--up',
                      '--running']

        srun_args = ['srun',
                     '--nodes=1']

        local_proc = subprocess.run(local_args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        lout = json.loads(local_proc.stdout)
        print(f"Type: {type(lout)}")
        pprint(lout)
        remote_proc = subprocess.run(srun_args + local_args,
                                     stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        rout = json.loads(remote_proc.stdout)
        print(f"Type: {type(rout)}")
        pprint(rout)

    except Exception as e:
        raise RuntimeError(e)


if __name__ == '__main__':

    get_ip_addresses()

test/launcher_multi/test_tcp_bringup/test_ifaddr.py
DragonHPC/dragon
import logging
import os
import argparse

from dragon.managed_memory import MemoryPool
from dragon.dlogging.util import setup_logging
from dragon.channels import Channel, register_gateways_from_env
from dragon.infrastructure import facts as dfacts
from dragon.infrastructure import messages as dmsg

from dragon.infrastructure.connection import Connection, ConnectionOptions
from dragon.launcher.channel_tree import start_channel_tree
from dragon.utils import B64

setup_logging(basename='backend', level=logging.DEBUG)
LOGGER = logging.getLogger('la_be')


def get_args():

    parser = argparse.ArgumentParser(prog='launcher-be',
                                     description='Arguments for standing up backend infrastructure tree for Dragon')
    parser.add_argument('--ip-addrs', metavar='FRONTEND_IP', dest='ip_addrs', nargs='+',
                        type=str)
    parser.add_argument('--local-sdesc', dest='local_sdesc', type=str)
    parser.add_argument('--host-ids', dest='host_ids', type=str, nargs='+')

    parser.add_argument('-n', '--node-index', dest='node_index', type=int)
    parser.add_argument('--frontend-sdesc', dest='frontend_sdesc')

    args = parser.parse_args()
    return args


if __name__ == '__main__':

    nnodes = 1

    args = get_args()
    LOGGER.info(f"args: {args}")
    # Create my memory pool
    be_mpool = MemoryPool(int(dfacts.DEFAULT_SINGLE_DEF_SEG_SZ),
                          f'{os.getuid()}_{os.getpid()}_{args.node_index}' + dfacts.DEFAULT_POOL_SUFFIX,
                          dfacts.default_pool_muid_from_index(args.node_index),
                          )

    # Create my receiving channel:
    be_cuid = 1545
    local_cuid = 1600
    gw_cuid = 6001

    fe_sdesc = B64.from_str(args.frontend_sdesc)
    be_inbound = Channel(be_mpool, be_cuid)
    local_ch = Channel(be_mpool, local_cuid)
    local_in = Connection(inbound_initializer=local_ch)
    overlay_in = Connection(inbound_initializer=be_inbound)

    # Create a gateway channel for myself
    gw_ch = Channel(be_mpool, gw_cuid)

    # Set gateway in my environment and register them
    encoded_ser_gw = B64(gw_ch.serialize())
    encoded_ser_gw_str = str(encoded_ser_gw)
    os.environ[dfacts.GW_ENV_PREFIX + '1'] = encoded_ser_gw_str
    register_gateways_from_env()

    LOGGER.info('Memory pools and channels created')

    # start my transport agent
    LOGGER.debug(f'standing up tcp agent with gw: {encoded_ser_gw_str}')
    tree_proc = start_channel_tree(args.node_index,
                                   str(B64(local_ch.serialize())),
                                   encoded_ser_gw_str,
                                   args.host_ids,
                                   args.ip_addrs)

    # Wait on a message telling me it's good to go.
    LOGGER.info('Channel tree initializing....')

    ping_back = dmsg.parse(local_in.recv())
    LOGGER.debug(f'recvd msg type {type(ping_back)}')

    # Connect to frontend, provide a mem pool since it will try to use
    # the default pool that doesn't exist yet since there aren't any overlay transport channels yet
    be_outbound = Channel.attach(fe_sdesc.decode(), mem_pool=be_mpool)
    conn_options = ConnectionOptions(default_pool=be_mpool)
    overlay_out = Connection(outbound_initializer=be_outbound, options=conn_options)

    # Send my serialized descriptor to the frontend
    overlay_out.send(str(B64(be_inbound.serialize())))
    LOGGER.info(f'sent my serialized descriptor to the frontend: {str(B64(be_inbound.serialize()))}')
    # Try getting a backend channel descriptor
    handshake_msg = overlay_in.recv()
    LOGGER.info(f'backend recvd: {handshake_msg}')

    ############### DO MY TESTING

    # TEARDOWN
    LOGGER.info('Tearing down')

    LOGGER.info('Received all TAHalted')

    overlay_out.close()
    overlay_in.close()

    tree_proc.terminate()
    tree_proc.wait(timeout=3)
    tree_proc.kill()

    gw_ch.destroy()
    local_in.close()
    local_ch.destroy()
    be_inbound.destroy()
    be_mpool.destroy()

    LOGGER.debug('Teardown completed')

test/launcher_multi/test_tcp_bringup/test_launcher_be.py
DragonHPC/dragon
import time
import logging
import os
import subprocess

from dragon.managed_memory import MemoryPool
from dragon.dlogging.util import setup_logging
from dragon.channels import Channel, register_gateways_from_env
from dragon.infrastructure import facts as dfacts
from dragon.infrastructure import messages as dmsg
from dragon.infrastructure.connection import Connection, ConnectionOptions
from dragon.infrastructure.node_desc import NodeDescriptor

from dragon.launcher.channel_tree import start_channel_tree
from dragon.launcher.network_config import NetworkConfig, WLM

from dragon.utils import B64

setup_logging(basename='frontend', level=logging.DEBUG)
LOGGER = logging.getLogger('la_fe')
NUM_NODES = 1
ENV = None


if __name__ == '__main__':

    # Get the node config for the backend
    LOGGER.debug('Setting up')
    net = NetworkConfig.from_wlm(workload_manager=WLM.SLURM,
                                 port=7779,
                                 network_prefix=dfacts.DEFAULT_TRANSPORT_NETIF)
    net_config = net.get_network_config()

    # Add the frontend config
    net_config['f'] = NodeDescriptor.get_local_node_network_conf(port_range=7779)

    LOGGER.debug(f'network config: {net_config}')

    # Create my memory pool
    fe_mpool = MemoryPool(int(dfacts.DEFAULT_SINGLE_DEF_SEG_SZ),
                          f'{os.getuid()}_{os.getpid()}_0' + dfacts.DEFAULT_POOL_SUFFIX,
                          dfacts.default_pool_muid_from_index(1),
                          )

    # Create my receiving channel:
    fe_cuid = 1544
    local_cuid = 1654
    gw_cuid = 6000

    # Channel for backend to come to
    fe_inbound = Channel(fe_mpool, fe_cuid)
    encoded_inbound = B64(fe_inbound.serialize())
    encoded_inbound_str = str(encoded_inbound)
    conn_in = Connection(inbound_initializer=fe_inbound)

    # Channel for tcp to tell me it's up
    local_ch = Channel(fe_mpool, local_cuid)
    local_in = Connection(inbound_initializer=local_ch)

    # Create a gateway channel for my tcp agent
    gw_ch = Channel(fe_mpool, gw_cuid)
    LOGGER.info('Memory pools and channels created')

    # Set gateway in my environment and register them
    encoded_ser_gw = B64(gw_ch.serialize())
    encoded_ser_gw_str = str(encoded_ser_gw)
    os.environ[dfacts.GW_ENV_PREFIX + '1'] = encoded_ser_gw_str
    register_gateways_from_env()

    # start my transport agent
    nnodes = len(net_config) - 1  # Exclude the frontend node from this
    host_ids = [str(node_desc.host_id) for node_desc in net_config.values()]
    ip_addrs = [node_desc.ip_addrs[0] for node_desc in net_config.values()]
    LOGGER.debug(f'standing up tcp agent with gw: {encoded_ser_gw_str}')
    tree_proc = start_channel_tree(None, str(B64(local_ch.serialize())),
                                   encoded_ser_gw_str, host_ids, ip_addrs)

    # Wait on a started message
    LOGGER.info('Channel tree initializing...')
    ping_back = dmsg.parse(local_in.recv())
    LOGGER.debug(f'recvd msg type {type(ping_back)}')    # Start my backend

    LOGGER.debug("standing up backend")
    backend_args = ['srun', f'--nodes={nnodes}', 'python3', './test_launcher_be.py', '--ip-addrs']
    backend_args = backend_args + ip_addrs + ['--host-ids'] + host_ids + ['--frontend-sdesc', encoded_inbound_str, '--node-index', '0']
    srun_procs = subprocess.Popen(args=backend_args)
    LOGGER.debug("backend started")

    # Try getting a backend channel descriptor
    be_ch_desc = conn_in.recv()
    LOGGER.debug(f'received descriptor: {be_ch_desc}')

    be_sdesc = B64.from_str(be_ch_desc)
    be_ch = Channel.attach(be_sdesc.decode(), mem_pool=fe_mpool)
    conn_options = ConnectionOptions(default_pool=fe_mpool)
    conn_out = Connection(outbound_initializer=be_ch, options=conn_options)
    msg_out = "hi from the frontend"
    LOGGER.debug(f'sending to be: {msg_out}')
    conn_out.send("hi from the frontend")
    time.sleep(5)
    ############### DO MY TESTING

    # TEARDOWN
    LOGGER.info('Tearing down')

    conn_in.close()
    conn_out.close()

    tree_proc.terminate()
    tree_proc.wait(timeout=3)
    tree_proc.kill()

    gw_ch.destroy()
    local_in.close()
    local_ch.destroy()
    fe_inbound.destroy()
    fe_mpool.destroy()

    LOGGER.debug('Teardown completed')



test/launcher_multi/test_tcp_bringup/test_launcher_fe.py
DragonHPC/dragon
import random
import string
import sys
import traceback
import time
import pickle
import multiprocessing as mp

from dragon.channels import Channel, register_gateways_from_env, POLLIN
from dragon.utils import B64

import dragon.infrastructure.parameters as dparms
from dragon.infrastructure.connection import Connection


def send_to(ch_descr, message):
    send_ch = Channel.attach(ch_descr)
    send_h = send_ch.sendh()
    send_h.open()
    send_h.send_bytes(message)
    send_h.close()
    send_ch.detach()


def recv_from(ch_descr):
    recv_ch = Channel.attach(ch_descr)
    recv_h = recv_ch.recvh()
    recv_h.open()
    message = recv_h.recv_bytes()
    recv_h.close()
    recv_ch.detach()
    return message


def send_message_around_ring_of_nodes(this_node_index, send_to_node_index, recv_from_node_index, send_node, recv_node):
    """Send a random secret message around the ring of allocated nodes. Each node has an allocated channel that
       can be sent to/received from.
       Depending on the values of send_node and recv_node, this function can send the message or clockwise,
       counter clockwise. Also depending on the values of send_node and recv_node, either the send or recv will
       be remote.
        Clockwise Local Send / Remote Recv    │    Clockwise Remote Send / Local Recv
                    ┌────────┐                │                ┌────────┐
                    │        │                │                │        │
            ┌─recv──┤ Node 0 ◄─recv─┐         │        ┌─send──► Node 0 ├─send─┐
            │       │        │      │         │        │       │        │      │
            │       └──▲───┬─┘     msg        │        │       └──▲───┬─┘     msg
            │          └─msg        │         │        │          └───┘        │
        ┌───▼────┐     send     ┌───┴────┐    │    ┌───┴────┐     recv     ┌───▼────┐
        │        ├─┐          ┌─►        │    │    │        ├─┐          ┌─►        │
        │ Node 4 │ │send  send│ │ Node   │    │    │ Node 4 │ │recv  recv│ │ Node   │
        │        ◄─┘          └─┤        │    │    │        ◄─┘         msg┤        │
        └┬───────┘              └───────▲┘    │    └▲───────┘              └───────┬┘
         │    send              send    │     │     │    recv              recv    │
       recv   ┌───┐             ┌───┐  recv   │   send   ┌───┐             ┌───┐  send
         │ ┌──▼───┴─┐        ┌──▼───┴─┐ │     │     │ ┌──▼───┴─┐        ┌──▼───┴─┐ │
         │ │        │        │        │ │     │     │ │        │        │        │ │
         └─► Node 3 ├──recv──► Node 2 ├─┘     │     └─┤ Node 3 ◄──send──┤ Node 2 ◄─┘
           │        │        │        │       │       │        │        │        │
           └────────┘        └────────┘       │       └────────┘        └────────┘
    """
    try:
        if this_node_index == 0:
            my_secret_message = ''.join(random.choice(string.ascii_letters) for _ in range(64)).encode()

            print(f'{this_node_index}: Sending node {send_to_node_index} my secret message {my_secret_message}', flush=True)
            send_to(send_node, my_secret_message)

            print(f'{this_node_index}: Waiting to receive my message from node {recv_from_node_index}', flush=True)
            check_message = recv_from(recv_node)

            print(f'{this_node_index}: Received {check_message}', flush=True)
            if check_message == my_secret_message:
                print(f'{this_node_index}: SUCCESS. I received the same message that I sent.', flush=True)
            else:
                print(f'{this_node_index}: ERROR. I did not receive the same message that I sent.', flush=True)
        else:
            print(f'{this_node_index}: Waiting to receive a secret message from node {recv_from_node_index}', flush=True)
            the_secret_message = recv_from(recv_node)

            print(f'{this_node_index}: Received {the_secret_message}', flush=True)
            print(f'{this_node_index}: Forwarding to node {send_to_node_index}', flush=True)

            send_to(send_node, the_secret_message)

        print(f'{this_node_index}: Done', flush=True)
    except Exception:
        print("!!!!There was an EXCEPTION!!!!!")
        ex_type, ex_value, ex_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
        print(tb_str)


def poll_ch(ch_descr):
    _ch = Channel.attach(ch_descr)
    rc = _ch.poll(event_mask=POLLIN)
    return rc


def test_remote_polling(this_node_index, node_to_poll_index, node_to_poll):
    """Test remote polling via the transport agent.
    From node 0, attach to the first channel from the channel list on node 1.
    From node 0, call poll with the POLLIN constant. This is defined in pydragon_channels.pyx.
    On node 0, if poll returns without raising an exception and return True, then the test succeeded.
    On node 1, sleep for five seconds, then attach to the first channel on its own node.
    On node 1, send a message into the channel you just attached to.
    On node 1, sleep for five more seconds to make sure that node 0 gets to successful completion.
    If the node index is not 0 or 1, then the program should sleep for 10 seconds and then exit.
    """
    try:
        rc = True
        if this_node_index == 0:
            print(f'{this_node_index}: Ready to poll on node {node_to_poll_index}', flush=True)
            rc = poll_ch(node_to_poll)
            if rc:
                print(f'{this_node_index}: SUCCESS. The poll completed successfully.', flush=True)
                recv_from(node_to_poll) # receive the message to clean the channel
            else:
                print(f'{this_node_index}: ERROR. Something went wrong with poll.', flush=True)
        elif this_node_index == 1:
            time.sleep(5)
            my_secret_message = ''.join(random.choice(string.ascii_letters) for _ in range(64)).encode()
            print(f'{this_node_index}: Sending node {node_to_poll_index} my secret message {my_secret_message}', flush=True)
            send_to(node_to_poll, my_secret_message)
        else:
            print(f'{this_node_index}: I am doing nothing. Moving on.', flush=True)

        print(f'{this_node_index}: Done with test_remote_polling', flush=True)
        return rc
    except Exception:
        print("!!!!There was an EXCEPTION on test_remote_polling!!!!!", flush=True)
        ex_type, ex_value, ex_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
        print(tb_str)


def _connection_comm(this_node_index, nodes_to_channel_map, some_msg, isconnin, isconnout, isbytes):
    if this_node_index == 0:
        if isconnout:
            ch_descr = B64.from_str(nodes_to_channel_map[0][0]).decode()
            _ch = Channel.attach(ch_descr)
            connout = Connection(outbound_initializer=_ch)

            if isbytes:
                connout.send_bytes(some_msg)
            else:
                connout.send(some_msg)

            connout.ghost_close()
        if isconnin:
            ch_descr = B64.from_str(nodes_to_channel_map[0][1]).decode()
            _ch = Channel.attach(ch_descr)
            connin = Connection(inbound_initializer=_ch)

            print(f'{this_node_index}: Waiting to receive message from node 1', flush=True)
            if isbytes:
                msg = connin.recv_bytes()
            else:
                msg = connin.recv()
            assert some_msg == msg

            connin.ghost_close()

    elif this_node_index == 1:
        if isconnin:
            ch_descr = B64.from_str(nodes_to_channel_map[0][0]).decode()
            _ch = Channel.attach(ch_descr)
            connin = Connection(inbound_initializer=_ch)

            print(f'{this_node_index}: Waiting to receive a message from node 0', flush=True)
            if isbytes:
                msg = connin.recv_bytes()
            else:
                msg = connin.recv()
            assert some_msg == msg

            connin.ghost_close()
        if isconnout:
            ch_descr = B64.from_str(nodes_to_channel_map[0][1]).decode()
            _ch = Channel.attach(ch_descr)
            connout = Connection(outbound_initializer=_ch)

            if isbytes:
                connout.send_bytes(some_msg)
            else:
                connout.send(some_msg)

            connout.ghost_close()

def test_connection_basic(nodes_to_channel_map):
    """Basic communication between nodes 0 and 1, through Connection objects.
    Node 0 sends some data to node 1 and then waits to receive it back.
    Node 1 is doing remote operations, both send and recv,
    whereas Node 0 is doing local operations only.
    :param nodes_to_channel_map: map of pre-created channels on all the nodes of the allocation
    :type nodes_to_channel_map: dict, with key as the node index and value a list of channels
    """
    this_node_index = dparms.this_process.index
    some_msg = ['a', 'b', 'c', 1, 2, 3]
    try:
        if this_node_index == 0 or this_node_index == 1:
            _connection_comm(this_node_index, nodes_to_channel_map, some_msg, True, True, False)
        else:
            print(f'{this_node_index}: I am doing nothing for test_connection_basic.', flush=True)
    except Exception:
        print("!!!!There was an EXCEPTION on test_connection_basic!!!!!")
        ex_type, ex_value, ex_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
        print(tb_str)

def test_connection_different_msg_sizes(nodes_to_channel_map):
    """Remote receive of different message sizes.
    :param nodes_to_channel_map: map of pre-created channels on all the nodes of the allocation
    :type nodes_to_channel_map: dict, with key as the node index and value a list of channels
    """
    this_node_index = dparms.this_process.index
    for lg_size in range(25):
        some_msg = bytes(2 ** lg_size)
        try:
            if this_node_index == 0:
                _connection_comm(0, nodes_to_channel_map, some_msg, isconnin=False, isconnout=True, isbytes=False)
            elif this_node_index == 1:
                _connection_comm(1, nodes_to_channel_map, some_msg, isconnin=True, isconnout=False, isbytes=False)
            else:
                print(f'{this_node_index}: I am doing nothing for test_connection_different_msg_sizes.', flush=True)
        except Exception:
            print("!!!!There was an EXCEPTION on test_connection_different_msg_sizes!!!!!")
            ex_type, ex_value, ex_tb = sys.exc_info()
            tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
            print(tb_str)

def test_connection_different_objs(nodes_to_channel_map):
    """Remote receive of different object types.
    :param nodes_to_channel_map: map of pre-created channels on all the nodes of the allocation
    :type nodes_to_channel_map: dict, with key as the node index and value a list of channels
    """
    this_node_index = dparms.this_process.index
    objs = ['hyena', {'hyena': 'hyena'}, bytes(100000), bytearray(200000), bytes('this is a test', 'utf-8'), b'\x01\x02\x03\x04\x05'*1000]
    for some_msg in objs:
        try:
            if this_node_index == 0:
                _connection_comm(0, nodes_to_channel_map, some_msg, isconnin=False, isconnout=True, isbytes=False)
            elif this_node_index == 1:
                _connection_comm(1, nodes_to_channel_map, some_msg, isconnin=True, isconnout=False, isbytes=False)
            else:
                print(f'{this_node_index}: I am doing nothing for test_connection_different_objs.', flush=True)
        except Exception:
            print("!!!!There was an EXCEPTION on test_connection_different_objs!!!!!")
            ex_type, ex_value, ex_tb = sys.exc_info()
            tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
            print(tb_str)

def test_connection_poll(nodes_to_channel_map):
    """Remote Connection poll.
    :param nodes_to_channel_map: map of pre-created channels on all the nodes of the allocation
    :type nodes_to_channel_map: dict, with key as the node index and value a list of channels
    """
    this_node_index = dparms.this_process.index
    some_msg = b'\x01\x02\x03\x04\x05'
    try:
        if this_node_index == 0:
            ch_descr = B64.from_str(nodes_to_channel_map[0][0]).decode()
            _ch = Channel.attach(ch_descr)
            connout = Connection(outbound_initializer=_ch)
            time.sleep(3)
            connout.send(some_msg)
            connout.ghost_close()
        elif this_node_index == 1:
            ch_descr = B64.from_str(nodes_to_channel_map[0][0]).decode()
            _ch = Channel.attach(ch_descr)
            connin = Connection(inbound_initializer=_ch)
            rc = connin.poll(timeout=20.0)
            assert rc == True
            msg = connin.recv()
            assert msg == some_msg
            connin.ghost_close()
        else:
            print(f'{this_node_index}: I am doing nothing for test_connection_poll.', flush=True)
    except Exception:
        print("!!!!There was an EXCEPTION on test_connection_poll!!!!!")
        ex_type, ex_value, ex_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
        print(tb_str)

def test_connection_remote_recv_bytes_from_send(nodes_to_channel_map):
    """Local send and remote recv_bytes. Edge case.
    :param nodes_to_channel_map: map of pre-created channels on all the nodes of the allocation
    :type nodes_to_channel_map: dict, with key as the node index and value a list of channels
    """
    this_node_index = dparms.this_process.index
    some_msg = {'hyenas': 'are awesome'}
    try:
        if this_node_index == 0:
            ch_descr = B64.from_str(nodes_to_channel_map[0][0]).decode()
            _ch = Channel.attach(ch_descr)
            connout = Connection(outbound_initializer=_ch)
            connout.send(some_msg)
            connout.ghost_close()
        elif this_node_index == 1:
            ch_descr = B64.from_str(nodes_to_channel_map[0][0]).decode()
            _ch = Channel.attach(ch_descr)
            connin = Connection(inbound_initializer=_ch)
            ser_msg = connin.recv_bytes()
            msg = pickle.loads(ser_msg)
            assert msg == some_msg
            connin.ghost_close()
        else:
            print(f'{this_node_index}: I am doing nothing for test_connection_remote_recv_bytes_from_send.', flush=True)
    except Exception:
        print("!!!!There was an EXCEPTION on test_connection_remote_recv_bytes_from_send!!!!!")
        ex_type, ex_value, ex_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
        print(tb_str)

def test_connection_recv_from_remote_send_bytes(nodes_to_channel_map):
    """Remote send_bytes and local recv. Edge case.
    :param nodes_to_channel_map: map of pre-created channels on all the nodes of the allocation
    :type nodes_to_channel_map: dict, with key as the node index and value a list of channels
    """
    this_node_index = dparms.this_process.index
    some_msg = {'hyenas': 'are awesome'}
    try:
        if this_node_index == 0:
            ch_descr = B64.from_str(nodes_to_channel_map[1][0]).decode()
            _ch = Channel.attach(ch_descr)
            connout = Connection(outbound_initializer=_ch)
            ser_msg = pickle.dumps(some_msg)
            connout.send_bytes(ser_msg)
            connout.ghost_close()
        elif this_node_index == 1:
            ch_descr = B64.from_str(nodes_to_channel_map[1][0]).decode()
            _ch = Channel.attach(ch_descr)
            connin = Connection(inbound_initializer=_ch)
            msg = connin.recv()
            assert msg == some_msg
            connin.ghost_close()
        else:
            print(f'{this_node_index}: I am doing nothing for test_connection_recv_from_remote_send_bytes.', flush=True)
    except Exception:
        print("!!!!There was an EXCEPTION on test_connection_recv_from_remote_send_bytes!!!!!")
        ex_type, ex_value, ex_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
        print(tb_str)

def send_stuff(nnodes, this_node_index, all_connouts, msg):
    register_gateways_from_env()
    for i in range(nnodes):
        if i is not this_node_index:
            # print(f"Node {this_node_index} sending to node {i}")
            all_connouts[this_node_index, i].send(msg)

def recv_stuff(nnodes, connin, msg):
    register_gateways_from_env()
    for i in range(nnodes - 1):
        # print(f"Node {this_node_index} receiving msg {i}")
        msg_recv = connin.recv()
        assert msg_recv == msg

def test_connection_alltoall(nodes_to_channel_map):
    this_node_index = dparms.this_process.index
    nnodes = len(nodes_to_channel_map)

    all_connouts = {}
    for i in range(nnodes):
        if i is not this_node_index:
            # connection for remote sending on node i
            ch_descr = B64.from_str(nodes_to_channel_map[i][0]).decode()
            _ch = Channel.attach(ch_descr)
            all_connouts[this_node_index, i] = Connection(outbound_initializer=_ch)

    # connection object for local receiving on this node
    ch_descr = B64.from_str(nodes_to_channel_map[this_node_index][0]).decode()
    _ch = Channel.attach(ch_descr)
    connin = Connection(inbound_initializer=_ch)

    msg = b'\x01\x02\x03\x04\x05'

    try:
        send_proc = mp.Process(target=send_stuff, args=(nnodes, this_node_index, all_connouts, msg))
        recv_proc = mp.Process(target=recv_stuff, args=(nnodes, connin, msg))
        send_proc.start()
        recv_proc.start()
        send_proc.join()
        recv_proc.join()
    except Exception:
        print("!!!!There was an EXCEPTION on test_connection_alltoall!!!!!")
        ex_type, ex_value, ex_tb = sys.exc_info()
        tb_str = "".join(traceback.format_exception(ex_type, ex_value, ex_tb))
        print(tb_str)


def main():
    mp.set_start_method('spawn')

    register_gateways_from_env()
    nodes_to_channel_map = eval(input())
    this_node_index = dparms.this_process.index

    next_node_index = (this_node_index + 1) % len(nodes_to_channel_map)
    prev_node_index = (this_node_index - 1) % len(nodes_to_channel_map)
    print(f'{this_node_index}: prev_node_index={prev_node_index} next_node_index={next_node_index}', flush=True)
    print(f'{this_node_index}: Got node to channel map={nodes_to_channel_map}', flush=True)

    next_node = B64.from_str(nodes_to_channel_map[next_node_index][0]).decode()
    this_node = B64.from_str(nodes_to_channel_map[this_node_index][0]).decode()
    prev_node = B64.from_str(nodes_to_channel_map[prev_node_index][0]).decode()

    print('Initiating clockwise remote send/local recv ring test:', flush=True)
    send_message_around_ring_of_nodes(this_node_index,
                                      send_to_node_index=next_node_index,
                                      recv_from_node_index=this_node_index,
                                      send_node=next_node, recv_node=this_node)

    print('Initiating clockwise local send/remote recv ring test:', flush=True)
    send_message_around_ring_of_nodes(this_node_index,
                                      send_to_node_index=this_node_index,
                                      recv_from_node_index=prev_node_index,
                                      send_node=this_node, recv_node=prev_node)

    # this is a connection test including all nodes
    print("\ntest_connection_alltoall", flush=True)
    test_connection_alltoall(nodes_to_channel_map)

    # Beginning of tests including only 2 nodes
    # this is a channel test including only two nodes
    node_to_poll_index = 1
    sync_index = 0
    sync_channel = B64.from_str(nodes_to_channel_map[sync_index][0]).decode()
    node_to_poll = B64.from_str(nodes_to_channel_map[node_to_poll_index][0]).decode()
    print("Initiating remote polling:", flush=True)
    rc = test_remote_polling(this_node_index,
                             node_to_poll_index=node_to_poll_index,
                             node_to_poll=node_to_poll)
    if rc:
        print(f"The poll test was a SUCCESS on node {this_node_index}!", flush=True)
    else:
        print(f"The poll test failed on node {this_node_index}!", flush=True)

    # these are connection tests including only 2 nodes
    print("\ntest_connection_basic", flush=True)
    test_connection_basic(nodes_to_channel_map)
    print("\ntest_connection_different_msg_sizes", flush=True)
    test_connection_different_msg_sizes(nodes_to_channel_map)
    print("\ntest_connection_different_objs", flush=True)
    test_connection_different_objs(nodes_to_channel_map)
    print("\ntest_connection_recv_bytes_from_send", flush=True)
    test_connection_remote_recv_bytes_from_send(nodes_to_channel_map)
    print("\ntest_connection_recv_from_remote_send_bytes", flush=True)
    test_connection_recv_from_remote_send_bytes(nodes_to_channel_map)
    print("\ntest_connection_poll", flush=True)
    test_connection_poll(nodes_to_channel_map)


if __name__ == "__main__":
    main()

test/launcher_multi/transport_test.py
DragonHPC/dragon
import dragon.localservices.local_svc as dsls
import dragon.infrastructure.log_setup as dlog
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.facts as dfacts
import logging
import os
import threading

dlog.setup_logging(basename='try_ls', the_level=logging.DEBUG)
log = logging.getLogger('driver')

print(f'{os.getpid()=}, {os.getpgid(0)=}')
os.setpgid(0, 0)
msg = f'now {os.getpgid(0)=}'
print(msg)
log.info(msg)

ls = dsls.LocalServer()

dying = dmsg.SHProcessCreate(tag=0, p_uid=dfacts.GS_PUID, r_c_uid=dfacts.GS_INPUT_CUID,
                             t_p_uid=99, exe='/bin/sleep', args=['2'])

ls.create_process(dying)
print('created dying process')

channel_id, msg = ls.responses.get(timeout=1)
assert isinstance(msg, dmsg.SHProcessCreateResponse)
proc = ls.new_procs.get(timeout=1)

print(f'got response: {proc.pid=}')
print(f'{channel_id=}, {msg=!r}')

death = threading.Thread(target=ls.watch_death)
death.start()

print('waiting for death')
channel_id, msg = ls.responses.get(timeout=3)
assert isinstance(msg, dmsg.SHProcessExit)
print(f'{channel_id=}, {msg=!r}')

print('killing death')
ls.set_shutdown('test')
death.join(timeout=1)
print('victory over death')

# now one that exercises some i/o.

ls.shutdown_sig.clear()

chatty = dmsg.SHProcessCreate(tag=1, p_uid=dfacts.GS_PUID, r_c_uid=dfacts.GS_INPUT_CUID,
                              t_p_uid=100, exe='/bin/echo', args=['hyenas'])

iowatcher = threading.Thread(target=ls.watch_output)
iowatcher.start()
ls.create_process(chatty)

print('chatty process created')

channel_id, msg = ls.responses.get(timeout=1)
print(f'{channel_id=}, {msg=!r}')
channel_id, msg = ls.responses.get(timeout=1)
print(f'{channel_id=}, {msg=!r}')

ls.set_shutdown('test 2')
iowatcher.join(timeout=1)
print('iowatcher left well')

test/localservices/smoke.py
DragonHPC/dragon
import dragon.infrastructure.minconnection as dimc
import dragon.channels as dch
import dragon.managed_memory as dmm

# a basic example of how to use MinConnection

pool_name = "hyenas"
pool_size = 2 ** 30
pool_prealloc_blocks = None
pool_uid = 17
print('pool creating')
mpool = dmm.MemoryPool(pool_size, pool_name, pool_uid, pool_prealloc_blocks)
print('pool created')

channel_uid = 42
chan = dch.Channel(mpool, channel_uid)
print('channel created')

reading_mc = dimc.MinConnection(mpool, chan, True)
writing_mc = dimc.MinConnection(mpool, chan, False)

sendstr = 'hyenas!'
writing_mc.send(sendstr)
gotstr = reading_mc.recv()
assert (sendstr.strip() == gotstr.strip())

print('polling empty')
pollres = reading_mc.poll(2)
print(f'poll res is {pollres}')

writing_mc.send(sendstr)
print('polling nonempty')
pollres = reading_mc.poll(2)
print(f'poll res is {pollres}')
print('polling nonempty again')
pollres = reading_mc.poll(2)
print(f'poll res is {pollres}')
print('reading')
anotherstr = reading_mc.recv()
assert (sendstr.strip() == anotherstr.strip())

# this will hang, but it shows
# the default queue depth
# for i in range(1500):
#     print(f'pushing {i}')
#     writing_mc.send(sendstr)


# this closes the handles, but the
# underlying channel object is still there
# the connection object does not take ownership
# of this.
reading_mc.close()
writing_mc.close()

chan.destroy()
print('channel destroyed')

mpool.destroy()
print('pool destroyed')
print('goodbye')

test/minconnection/basic_example.py
DragonHPC/dragon
#!/usr/bin/env python3

"""Test script for basic MinConnection object"""

import multiprocessing
import subprocess
import unittest

import dragon.infrastructure.minconnection as dimc
import dragon.infrastructure.messages as dmsg
import dragon.channels as dch
import dragon.managed_memory as dmm


# tolerate double-detach on a channel
def soft_detach_chan(the_chan):
    try:
        the_chan.detach()
    except dch.ChannelError as dche:
        if dche.Errors.FAIL == dche.ex_code:
            if not dche.args[0].startswith('Could not detach'):
                raise dche


# tolerate double-detach on a pool
def soft_detach_pool(the_pool):
    try:
        the_pool.detach()
    except dmm.DragonPoolDetachFail:
        # Expect DetachFail, don't raise
        pass


def parse_echo(the_msgstr):
    msg = dmsg.parse(the_msgstr)
    return msg.serialize()


def echo_end_last(input_sc, input_sp, output_sc, output_sp, last_str, mutator=lambda x: x):
    """Echos messages coming from the input to the output.

        Exits when last_str is seen (and echoed)
    """

    last_str = last_str.strip()

    input_chan = dch.Channel.attach(input_sc)
    output_chan = dch.Channel.attach(output_sc)
    input_pool = dmm.MemoryPool.attach(input_sp)
    output_pool = dmm.MemoryPool.attach(output_sp)

    input_mc = dimc.MinConnection(input_pool, input_chan, reading=True)
    output_mc = dimc.MinConnection(output_pool, output_chan, reading=False)

    rec_str = None
    while rec_str != last_str:
        rec_str = input_mc.recv().strip()

        if rec_str != last_str:
            out_str = mutator(rec_str)
        else:
            out_str = rec_str

        output_mc.send(out_str)

    input_mc.close()
    output_mc.close()

    soft_detach_chan(input_chan)
    soft_detach_chan(output_chan)
    soft_detach_pool(input_pool)
    soft_detach_pool(output_pool)


def worker_pickling_echo_end_last(writer_mc, reader_mc, last_str, mutator=lambda x: x):
    last_str = last_str.strip()

    rec_str = None
    while rec_str != last_str:
        rec_str = reader_mc.recv().strip()

        if rec_str != last_str:
            out_str = mutator(rec_str)
        else:
            out_str = rec_str

        writer_mc.send(out_str)

    writer_mc.close()
    reader_mc.close()

    soft_detach_pool(writer_mc.the_channel)
    soft_detach_pool(reader_mc.the_channel)
    soft_detach_pool(writer_mc.the_pool)
    soft_detach_pool(reader_mc.the_pool)


class MinConnectionTest(unittest.TestCase):
    def setUp(self):
        username = subprocess.check_output('whoami').decode().strip()
        self.pool_name = 'mctest_' + username
        self.pool_size = 2 ** 30
        self.pool_uid = 17
        self.mpool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid)
        self.mpool_ser = self.mpool.serialize()

    def tearDown(self) -> None:
        self.mpool.destroy()

    def make_in_and_out_channels(self):
        in_channel_uid = 42
        in_chan = dch.Channel(self.mpool, in_channel_uid)
        in_chan_ser = in_chan.serialize()

        out_channel_uid = 43
        out_chan = dch.Channel(self.mpool, out_channel_uid)
        out_chan_ser = out_chan.serialize()

        return in_chan, in_chan_ser, out_chan, out_chan_ser

    def test_nothing(self):
        in_chan, _, out_chan, _ = self.make_in_and_out_channels()
        in_chan.destroy()
        out_chan.destroy()

    def test_echo(self):
        in_chan, in_chan_ser, out_chan, out_chan_ser = self.make_in_and_out_channels()

        inbound_mc = dimc.MinConnection(self.mpool, in_chan, reading=True)
        outbound_mc = dimc.MinConnection(self.mpool, out_chan, reading=False)

        words = 'hyena hyena hyena hyena'
        lastword = 'nomorehyena'

        echo_process = multiprocessing.Process(target=echo_end_last, args=(out_chan_ser, self.mpool_ser,
                                                                           in_chan_ser, self.mpool_ser,
                                                                           lastword))
        echo_process.start()

        for k in range(10):
            outbound_mc.send(words)
            rec = inbound_mc.recv()
            self.assertEqual(words.strip(), rec.strip())

        outbound_mc.send(lastword)
        rec = inbound_mc.recv()
        self.assertEqual(lastword.strip(), rec.strip())

        inbound_mc.close()
        outbound_mc.close()

        echo_process.join()
        in_chan.destroy()
        out_chan.destroy()

    def test_pickled_path_echo(self):
        in_chan, in_chan_ser, out_chan, out_chan_ser = self.make_in_and_out_channels()

        their_reader_mc = dimc.MinConnection(self.mpool, in_chan, reading=True)
        my_writer_mc = dimc.MinConnection(self.mpool, in_chan, reading=False)
        their_writer_mc = dimc.MinConnection(self.mpool, out_chan, reading=False)
        my_reader_mc = dimc.MinConnection(self.mpool, out_chan, reading=True)

        words = 'hyena Hyena hyena Hyena'
        lastword = 'no more Hyena'

        echo_process = multiprocessing.Process(target=worker_pickling_echo_end_last,
                                               args=(their_writer_mc, their_reader_mc, lastword))
        echo_process.start()

        for k in range(10):
            my_writer_mc.send(words)
            rec = my_reader_mc.recv()
            self.assertEqual(words.strip(), rec.strip())

        my_writer_mc.send(lastword)
        rec = my_reader_mc.recv()
        self.assertEqual(lastword.strip(), rec.strip())

        their_writer_mc.close()
        their_reader_mc.close()
        my_writer_mc.close()
        my_reader_mc.close()

        echo_process.join()

        in_chan.destroy()
        out_chan.destroy()

    def test_msg_echo(self):
        in_chan, in_chan_ser, out_chan, out_chan_ser = self.make_in_and_out_channels()

        inbound_mc = dimc.MinConnection(self.mpool, in_chan, reading=True)
        outbound_mc = dimc.MinConnection(self.mpool, out_chan, reading=False)

        lastword = 'nomorehyena'

        echo_process = multiprocessing.Process(target=echo_end_last, args=(out_chan_ser, self.mpool_ser,
                                                                           in_chan_ser, self.mpool_ser,
                                                                           lastword, parse_echo))
        echo_process.start()

        for msg in (dmsg.GSIsUp(1), dmsg.BEHalted(2), dmsg.GSHalted(3)):
            outbound_mc.send(msg.serialize())
            rec = inbound_mc.recv()
            return_msg = dmsg.parse(rec)
            self.assertEqual(type(msg), type(return_msg))
            self.assertEqual(msg.get_sdict(), return_msg.get_sdict())

        outbound_mc.send(lastword)
        rec = inbound_mc.recv()
        self.assertEqual(lastword.strip(), rec.strip())

        inbound_mc.close()
        outbound_mc.close()

        echo_process.join()
        in_chan.destroy()
        out_chan.destroy()


if __name__ == "__main__":
    unittest.main()

test/minconnection/minconnection_obj_tests.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest
from minconnection.minconnection_obj_tests import MinConnectionTest

if __name__ == "__main__":
    unittest.main()

test/minconnection/test_minconnection.py
DragonHPC/dragon
import multiprocessing as mp
import os
import test_util
import datetime
import time
import json

class BasicParameterStudy:
    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Run a basic parameter study')
        parser = test_util.add_connections_args(parser)
        parser.add_argument('--response_size', type=int, default=4, 
                            help='size of the message to pass back')
        parser.add_argument('--work_items', type=int, default=8,
                            help='number of work items in the parameter study')
        self.args = test_util.process_args(parser)
        mp.set_start_method('spawn')
        
    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_parameter_study, self.args)
        self.results = ParameterResults(self.args, times)
        if self.args.json:
            return self.results.dump()
        return self.results

    def dummy_func(self, x):
        return bytes(self.args.response_size)

    def generate_work_items(self):
        for item in range(self.args.work_items):
            yield bytes(self.args.message_size)

    def do_parameter_study(self):
        inputlist = [bytes(self.args.message_size) for i in range(self.args.work_items)]
        
        start_time_setup = time.time_ns()
        with mp.Pool(self.args.num_workers) as the_pool:
            start_time_no_setup = time.time_ns()
            the_pool.map(self.dummy_func, inputlist)
            end_time_no_setup = time.time_ns()
        end_time_setup = time.time_ns()

        return (end_time_setup - start_time_setup, end_time_no_setup - start_time_no_setup)

class ParameterResults():
    def __init__(self, args, times):
        self.args = args
        self.data_dict = {}
        self.times = times
        self.data_dict['num_workers'] = args.num_workers
        self.data_dict['response_size'] = args.response_size
        self.data_dict['message_size'] = args.message_size
        self.data_dict['work_items'] = args.work_items
        self.data_dict['hostname'] = test_util.get_hostname().strip()
        self.data_dict['date'] = datetime.datetime.now()
        self.__process_data()

    def __process_data(self):
        sum_size = (self.args.message_size + self.args.response_size) * self.args.work_items
        avg_latency = 0
        avg_compute_time=0
        for time in self.times:
            avg_latency = avg_latency + time[0]
            avg_compute_time = avg_compute_time + time[1]
        avg_latency = avg_latency / len(self.times)
        avg_compute_time = avg_compute_time / len(self.times)
        avg_setup_time = avg_latency - avg_compute_time
        avg_setup_time_pw = avg_setup_time / self.args.num_workers
        avg_compute_time_pwi = avg_compute_time / self.args.work_items
        avg_bandwidth = sum_size / avg_compute_time

        self.data_dict['avg_latency'] = avg_latency / 1000
        self.data_dict['avg_compute_time'] = avg_compute_time / 1000
        self.data_dict['avg_setup_time'] = avg_setup_time / 1000
        self.data_dict['avg_setup_time_per_worker'] = avg_setup_time_pw / 1000
        self.data_dict['avg_compute_time_per_work_item'] = avg_compute_time_pwi / 1000
        self.data_dict['avg_bandwidth'] = avg_bandwidth * (1000000000 / (2**20))

    def __str__(self):
        return (
                f'Test: Basic Parameter Study\n'
                f'Hostname: {self.data_dict["hostname"]}\n'
                f'Date: {self.data_dict["date"]}\n'
                f'Number of workers: {self.data_dict["num_workers"]}\n'
                f'Number of work items: {self.data_dict["work_items"]}\n'
                f'Message size (bytes): {self.data_dict["message_size"]}\n'
                f'Response size (bytes): {self.data_dict["response_size"]}\n'
                f'Average latency (\xb5s): {self.data_dict["avg_latency"]}\n'
                f'Average compute time (\xb5s): {self.data_dict["avg_compute_time"]}\n'
                f'Average setup time (\xb5s): {self.data_dict["avg_setup_time"]}\n'
                f'Average setup time per worker (\xb5s): {self.data_dict["avg_setup_time_per_worker"]}\n'
                f'Average compute time per work item (\xb5s): {self.data_dict["avg_compute_time_per_work_item"]}\n'
                f'Average bandwidth (mB/s): {self.data_dict["avg_bandwidth"]}'
               )
    def dump(self):
        return json.dumps(self.data_dict, indent=4, default=str)

    def python_dict(self):
        return self.data_dict



if __name__ == '__main__':
    test = BasicParameterStudy()
    print(test.run())

test/mp_bench/basic_parameter_study.py
DragonHPC/dragon
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
# import os
# import sys
# sys.path.insert(0, os.path.abspath('.'))


# -- Project information -----------------------------------------------------

project = 'Python Multiprocessing Benchmarks'
copyright = '2020-2023 HPE Dragon'
author = 'HPE Dragon'


# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']


# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
html_theme = 'alabaster'

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']
test/mp_bench/doc/conf.py
DragonHPC/dragon
import subprocess
import resource
import json
import csv
import os
import datetime
import argparse
import sys
import multiprocessing

class TestAutomation:
    def __init__(self, open_file_desc_limit):
        self.open_file_desc_limit = open_file_desc_limit
        self.directory_name = ('results'+ 
                              str((subprocess.check_output('hostname').decode()).strip())+ 
                              str(datetime.datetime.now()))
        self.cpus = multiprocessing.cpu_count()
        self.args = None
        self.__setup()

    def __setup(self):
        os.mkdir(self.directory_name)
        self.args = self.__get_args()
        if (not len(sys.argv) > 1) or ((len(sys.argv) == 3) and (self.args.num_iterations is not None)):
            self.run_all_tests()
        else:
            self.__run_subset_tests()

    def __get_args(self):
        parser = argparse.ArgumentParser()
        parser.add_argument('--basic_process_test', action='store_true',
                            help='run suite of benchmarks for basic_process_test.py')
        parser.add_argument('--basic_pool_test', action='store_true',
                            help='run suite of benchmarks for basic_pool_test.py')
        parser.add_argument('--pool_invocation_test', action='store_true',
                            help='run suite of benchmarks for pool_invocation_test.py')      
        parser.add_argument('--send_receive_test', action='store_true',
                            help='run suite of benchmarks for send_receive_test.py')
        parser.add_argument('--chain_token_test', action='store_true',
                            help='run suite of benchmarks for chain_token_test.py')
        parser.add_argument('--chain_token_pipe_test', action='store_true',
                            help='run suite of benchmarks for chain_token_pipe_test.py')
        parser.add_argument('--chain_token_test_simple', action='store_true',
                            help='run suite of benchmarks for chain_token_test_simple.py')
        parser.add_argument('--basic_managed_dict_test', action='store_true',
                            help='run suite of benchmarks for basic_managed_dict_test.py')
        parser.add_argument('--access_managed_dict_test', action='store_true',
                            help='run suite of benchmarks for access_managed_dict_test.py') 
        parser.add_argument('--num_iterations', type=int, default=None,
                            help='number of iteration for each instance of run tests')
        parser.add_argument('--spins', type=int, default=None,
                            help='number of spins for each instance of a test with spins.'+
                            ' can lead to duplicated tests.')
        parser.add_argument('--max_workers', type=int, default=self.cpus,
                            help='maximum number of workers a test will run with')
        return parser.parse_args() 
        
        
    def run_all_tests(self):
        self.run_basic_pool_test()
        self.run_basic_process_test()
        self.run_pool_invocation_test()
        self.run_send_receive_test()
        self.run_chain_token_test()
        self.run_chain_token_pipe_test()
        self.run_chain_token_test_simple()
        self.run_basic_managed_dict_test()
        self.run_access_managed_dict_test()

    def __run_subset_tests(self):
        test_run = False
        if self.args.basic_pool_test:
            self.run_basic_pool_test()
            test_run = True
        if self.args.basic_process_test:
            self.run_basic_process_test()
            test_run = True
        if self.args.pool_invocation_test:
            self.run_pool_invocation_test()
            test_run = True
        if self.args.send_receive_test:
            self.run_send_receive_test()
            test_run = True
        if self.args.chain_token_test:
            self.run_chain_token_test()
            test_run = True
        if self.args.chain_token_pipe_test:
            self.run_chain_token_pipe_test()
            test_run = True
        if self.args.chain_token_test_simple:
            self.run_chain_token_test_simple()
            test_run = True
        if self.args.basic_managed_dict_test:
            self.run_basic_managed_dict_test()
            test_run = True
        if self.args.access_managed_dict_test:
            self.run_access_managed_dict_test()
            test_run = True
        if not test_run:
            self.run_all_tests()

    def run_basic_pool_test(self):
        input_dict = self.__process_csv_input('automation_input/basic_pool_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/basic_pool_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/basic_pool_test.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/basic_pool_test_output.csv')

    def run_basic_process_test(self):
        input_dict = self.__process_csv_input('automation_input/basic_process_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/basic_process_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/basic_process_test.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/basic_process_test_output.csv')

    def run_pool_invocation_test(self):
        input_dict = self.__process_csv_input('automation_input/pool_invocation_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/pool_invocation_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/pool_invocation_test.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/pool_invocation_test_output.csv')

    def run_send_receive_test(self):
        input_dict = self.__process_csv_input('automation_input/send_receive_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/send_receive_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--message_size', input_dict['message_size'][i],
                                                       '--response_size', input_dict['response_size'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/send_receive_test.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/send_receive_test_output.csv')

    def run_chain_token_test(self):
        input_dict = self.__process_csv_input('automation_input/chain_token_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/chain_token_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--message_size', input_dict['message_size'][i],
                                                       '--spins', input_dict['spins'][i],
                                                       '--num_tokens', input_dict['num_tokens'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/chain_token_test.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/chain_token_test_output.csv')

    def run_chain_token_pipe_test(self):
        input_dict = self.__process_csv_input('automation_input/chain_token_pipe_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/chain_token_pipe_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--message_size', input_dict['message_size'][i],
                                                       '--spins', input_dict['spins'][i],
                                                       '--num_tokens', input_dict['num_tokens'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/chain_token_pipe_test.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/chain_token_pipe_test_output.csv')  

    def run_chain_token_test_simple(self):
        input_dict = self.__process_csv_input('automation_input/chain_token_test_simple_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/chain_token_test_simple.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--message_size', input_dict['message_size'][i],
                                                       '--spins', input_dict['spins'][i],
                                                       '--num_tokens', input_dict['num_tokens'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/chain_token_test_simple.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/chain_token_test_simple_output.csv')

    def run_basic_managed_dict_test(self):
        input_dict = self.__process_csv_input('automation_input/basic_managed_dict_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/basic_managed_dict_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--message_size', input_dict['message_size'][i],
                                                       '--spins', input_dict['spins'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/basic_managed_dict_test.py',
                                               '--num_workers', str(self.args.max_workers),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())       
        self.__process_json(output, self.directory_name+'/basic_managed_dict_test_output.csv')

    def run_access_managed_dict_test(self):
        input_dict = self.__process_csv_input('automation_input/access_managed_dict_test_input.csv')
        output = []
        for i in range(len(input_dict['num_workers'])):
            if int(input_dict['num_workers'][i]) <= self.args.max_workers:
                output.append(subprocess.check_output([sys.executable, 'tests/access_managed_dict_test.py',
                                                       '--num_workers', input_dict['num_workers'][i],
                                                       '--num_iterations', input_dict['num_iterations'][i],
                                                       '--message_size', input_dict['message_size'][i],
                                                       '--spins', input_dict['spins'][i],
                                                       '--json'
                                                     ]).decode())
        output.append(subprocess.check_output([sys.executable, 'tests/access_managed_dict_test.py',
                                               '--num_workers', str(int(self.args.max_workers/2)),
                                               '--num_iterations', '5',
                                               '--json'
                                             ]).decode())
        self.__process_json(output, self.directory_name+'/access_managed_dict_test_output.csv')

    def __process_csv_input(self, file_name):
        with open(file_name, newline='', encoding='utf-8-sig') as csvfile:
            reader = csv.DictReader(csvfile)
            input_dict = {}
            first = True
            for row in reader:
                if first:
                    for key in row:
                        if (key == 'num_iterations') and (self.args.num_iterations is not None):
                            input_dict[key] = [str(self.args.num_iterations)]
                        elif (key == 'spins') and (self.args.spins is not None):
                            input_dict[key] = [str(self.args.spins)]
                        else:
                            input_dict[key] = [row[key]]
                    first = False
                else:
                    for key in row:
                        if (key == 'num_iterations') and (self.args.num_iterations is not None):
                            input_dict[key].append(str(self.args.num_iterations))
                        elif (key == 'spins') and (self.args.spins is not None):
                            input_dict[key].append(str(self.args.spins))
                        else:
                            input_dict[key].append(row[key])
            return input_dict


    def __process_json(self, json_list, file_name):
        first = json.loads(json_list[0])
        headers = []
        for key in first:
            headers.append(key)

        with open(file_name, 'w', newline='') as csvfile:
            writer = csv.DictWriter(csvfile, fieldnames=headers)
            writer.writeheader()
            for d in json_list:
                writer.writerow(json.loads(d))


if __name__ == '__main__':
    test = TestAutomation(1024)




test/mp_bench/scalable_mp_benchmarks.py
DragonHPC/dragon
#!/usr/bin/env python3

"""
Basic performance test of multiprocessing managed dictionary
"""
import argparse
import multiprocessing
import time
import test_util

class AccessManagedDictTest:

    def __init__(self):
        self.args = None
        self.results = None
        self.test = 0

    def __setup(self):
        parser = test_util.add_default_args('Does a basic speed test of managed dicts')
        parser.add_argument('--spins', type=int, default=4,
                             help='number of times for each worker to touch the dictionary')
        parser = test_util.add_connections_args(parser)
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_managed_dict_test, self.args)
        self.results = test_util.Results(self.args, times, "access_managed_dict_test")
        if self.args.json:
            return self.dump() 
        return self.results

    def dump(self):
        return self.results.dump()

    def worker_entry(self, x, key_to_set, spins, message):
        for idx in range(spins):
            x[key_to_set] = message + bytes(idx)

    def touch_dict(self, x, key_to_check, spins):
        for idx in range(spins):
            test = x[key_to_check]

    def do_managed_dict_test(self):
        """Spawns a bunch of processes that set a key in a dictionary and return.

        They all set a different key.

        :return: Time it took to do this in ns
        """

        num_workers = self.args.num_workers
        spins = self.args.spins
        assert self.args.message_size > 0
        message = bytes(self.args.message_size)
        assert num_workers > 0

        process_handles = []
        touch_process_handles = []


        with multiprocessing.Manager() as mgr:

            d = mgr.dict()
            for idx in range(num_workers):
                handle = multiprocessing.Process(target=self.worker_entry, args=(d, idx, spins, message))
                process_handles.append(handle)

            for handle in process_handles:
                handle.start()

            for handle in process_handles:
                handle.join()
            
            start_time = time.time_ns()
            
            for idy in range(num_workers):
                handle = multiprocessing.Process(target=self.touch_dict, args=(d, idy, spins,))
                touch_process_handles.append(handle)

            for handle in touch_process_handles:
                handle.start()

            for handle in touch_process_handles:
                handle.join()
                


        return time.time_ns() - start_time



if __name__ == '__main__':
    test = AccessManagedDictTest()
    print(test.run())

test/mp_bench/tests/access_managed_dict_test.py
DragonHPC/dragon
#!/usr/bin/env python3

"""
Basic performance test of multiprocessing managed dictionary
"""
import argparse
import multiprocessing
import time
import test_util

class BasicManagedDictTest:

    def __init__(self):
        self.args = None
        self.results = None
        self.test = 0

    def __setup(self):
        parser = test_util.add_default_args('Does a basic speed test of managed dicts')
        parser.add_argument('--spins', type=int, default=4,
                             help='number of times for each worker to touch the dictionary')
        parser = test_util.add_connections_args(parser)
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_managed_dict_test, self.args)
        self.results = test_util.Results(self.args, times, "basic_managed_dict_test")
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump()

    def worker_entry(self, x, key_to_set, spins, message):
        for idx in range(spins):
            x[key_to_set] = message + bytes(idx)

    def do_managed_dict_test(self):
        """Spawns a bunch of processes that set a key in a dictionary and return.

        They all set a different key.

        :return: Time it took to do this in ns
        """

        num_workers = self.args.num_workers
        spins = self.args.spins
        assert self.args.message_size > 0
        message = bytes(self.args.message_size)
        assert num_workers > 0

        process_handles = []

        start_time = time.time_ns()

        with multiprocessing.Manager() as mgr:

            d = mgr.dict()
            for idx in range(num_workers):
                handle = multiprocessing.Process(target=self.worker_entry, args=(d, idx, spins, message))
                process_handles.append(handle)

            for handle in process_handles:
                handle.start()

            for handle in process_handles:
                handle.join()


        return time.time_ns() - start_time



if __name__ == '__main__':
    test = BasicManagedDictTest()
    print(test.run())

test/mp_bench/tests/basic_managed_dict_test.py
DragonHPC/dragon
#!/usr/bin/env python3

"""
Basic performance test of multiprocessing.Pool
"""

import multiprocessing
import time
import test_util

class BasicPoolTest:

    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Does a basic speed test of spawning processes in a Pool')
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_spawn_gather, self.args)
        self.results = test_util.Results(self.args, times, 'basic_pool_test')
        if self.args.json:
            return self.dump()
        return self.results       

    def dump(self):
        return self.results.dump()


    def dummy_func(self, x):
        return 42

    def do_spawn_gather(self):
        """Spawns a bunch of processes in a pool that do nothing and return.

        :return: Time it took to do this in ns
        """

        num_workers = self.args.num_workers

        assert num_workers > 0

        start_time = time.time_ns()
        completions = []
        with multiprocessing.Pool(num_workers) as the_pool:
            for idx in range(num_workers):
                completions.append( the_pool.apply_async(self.dummy_func, (idx,)))

            for k in completions:
                k.get()

        return time.time_ns() - start_time



if __name__ == '__main__':
    test = BasicPoolTest()
    print(test.run())


test/mp_bench/tests/basic_pool_test.py
DragonHPC/dragon
#!/usr/bin/env python3

"""
Basic performance test of multiprocessing.Process
"""

import multiprocessing
import time
import test_util

class BasicProcessTest:
    
    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Does a basic speed test of spawning processes')
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_spawn_gather, self.args)
        self.results = test_util.Results(self.args, times, "basic_process_test")
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump()


    def dummy_func(self, x):
        return x

    def do_spawn_gather(self):
        """Spawns a bunch of processes that do nothing and return.

        :return: Time it took to do this in ns
        """
        num_workers = self.args.num_workers

        assert num_workers > 0

        process_handles = []

        start_time = time.time_ns()

        for _ in range(num_workers):
            handle = multiprocessing.Process()
            process_handles.append(handle)

        for handle in process_handles:
            handle.start()

        for handle in process_handles:
            handle.join()

        return time.time_ns() - start_time


if __name__ == '__main__':
    #f = open("basic_pool_test.py", "r")
    test = BasicProcessTest()
    results = test.run()
    print(results)
    #f.close()

test/mp_bench/tests/basic_process_test.py
DragonHPC/dragon
#!/usr/bin/env python3
"""
Basic test of multiprocessing message passing through multiprocessing.Pipe.
"""

import argparse
import multiprocessing
import time
import test_util

class ChainTokenPipeTest:

    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Basic token passing test, passing around a chain using Pipe')
        parser = test_util.add_token_args(parser)
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_pipe_token_ring, self.args)
        self.results = test_util.Results(self.args, times, "chain_token_pipe_test")
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump()

    def pass_it(self, is_first, spins, in_conn, out_conn, message, num_tokens):
        """Worker function to pass a cookie around.

        The is_first one generates the cookie and absorbs it.
        All the other wait and pass the cookie.
        """

        if is_first:
            for _ in range(num_tokens):
                out_conn.send(message)
            for _ in range((spins*num_tokens) - 1):
                item = in_conn.recv()
                out_conn.send(item)

            last = in_conn.recv()
            assert (message == last)
        else:
            for _ in range(spins*num_tokens):
                item = in_conn.recv()
                out_conn.send(item)


    def do_pipe_token_ring(self):

        num_workers = self.args.num_workers
        spins = self.args.spins
        num_tokens = self.args.num_tokens
        assert num_workers > 0
        assert self.args.message_size >= 0
        message = bytes(self.args.message_size)


        start_time = time.time_ns()

        pipeconns = [multiprocessing.Pipe(duplex=False) for _ in range(num_workers)]
        processes = []
        for idx in range(num_workers):
            processes.append(multiprocessing.Process(target=self.pass_it,
                                                     args=((idx == 0), spins,
                                                           pipeconns[(idx-1) % num_workers][0],
                                                           pipeconns[idx][1], 
                                                           message, num_tokens)
                                                     ))
        for handle in processes:
            handle.start()

        for handle in processes:
            handle.join()

        return time.time_ns() - start_time


if __name__ == '__main__':
    test = ChainTokenPipeTest()
    print(test.run())

test/mp_bench/tests/chain_token_pipe_test.py
DragonHPC/dragon
#!/usr/bin/env python3
"""
Basic test of multiprocessing message passing through queues.
"""

import argparse
import multiprocessing
import time
import test_util

class ChainTokenTest:

    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Basic token passing test, passing around a chain using Queue')
        parser = test_util.add_token_args(parser)
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_queue_token_ring, self.args)
        self.results = test_util.Results(self.args, times, 'chain_token_test')
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump()


    def pass_it(self, is_first, spins, in_queue, out_queue, message, num_tokens):
        """Worker function to pass a cookie around.

        The is_first one generates the cookie and absorbs it.
        All the other wait and pass the cookie.
        """

        if is_first:
            for _ in range(num_tokens):
                out_queue.put(message)
            for _ in range((spins*num_tokens) - 1):
                item = in_queue.get()
                out_queue.put(item)

            last = in_queue.get()
            #assert (message == last)
        else:
            for _ in range(spins*num_tokens):
                item = in_queue.get()
                out_queue.put(item)


    def do_queue_token_ring(self):

        num_workers = self.args.num_workers
        spins = self.args.spins
        message = bytes(self.args.message_size)
        num_tokens = self.args.num_tokens
        assert num_workers > 0
        assert num_tokens > 0

        start_time = time.time_ns()

        queues = [multiprocessing.Queue() for _ in range(num_workers)]
        processes = []
        for idx in range(num_workers):
            processes.append(multiprocessing.Process(target=self.pass_it,
                                                     args=((idx == 0), spins,
                                                           queues[idx], 
                                                           queues[(idx + 1) % num_workers],
                                                           message, num_tokens)
                                                     ))
        for handle in processes:
            handle.start()

        for handle in processes:
            handle.join()

        return time.time_ns() - start_time


if __name__ == '__main__':
    test = ChainTokenTest()
    print(test.run())

test/mp_bench/tests/chain_token_test.py
DragonHPC/dragon
#!/usr/bin/env python3
"""
Basic test of multiprocessing message passing through simple queues.
"""

import argparse
import multiprocessing
import time
import test_util

class ChainTokenTestSimple:

    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Basic token passing test, passing around a chain using SimpleQueue')
        parser = test_util.add_token_args(parser)
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_queue_token_ring, self.args)
        self.results = test_util.Results(self.args, times, 'chain_token_test_simple')
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump()


    def pass_it(self, is_first, spins, in_queue, out_queue, message, num_tokens):
        """Worker function to pass a cookie around.

        The is_first one generates the cookie and absorbs it.
        All the other wait and pass the cookie.
        """

        if is_first:
            for _ in range(num_tokens):
                out_queue.put(message)
            for _ in range((spins*num_tokens) - 1):
                item = in_queue.get()
                out_queue.put(item)

            last = in_queue.get()
            assert (message == last)
        else:
            for _ in range(spins*num_tokens):
                item = in_queue.get()
                out_queue.put(item)


    def do_queue_token_ring(self):

        num_workers = self.args.num_workers
        spins = self.args.spins
        num_tokens = self.args.num_tokens
        assert num_tokens > 0
        assert self.args.message_size > 0
        message = bytes(self.args.message_size)
        assert num_workers > 0

        start_time = time.time_ns()

        queues = [multiprocessing.SimpleQueue() for _ in range(num_workers)]
        processes = []
        for idx in range(num_workers):
            processes.append(multiprocessing.Process(target=self.pass_it,
                                                     args=((idx == 0), spins,
                                                           queues[idx], 
                                                           queues[(idx + 1) % num_workers],
                                                           message, num_tokens)
                                                     ))

        for handle in processes:
            handle.start()

        for handle in processes:
            handle.join()

        return time.time_ns() - start_time


if __name__ == '__main__':
    test = ChainTokenTestSimple()
    print(test.run())

test/mp_bench/tests/chain_token_test_simple.py
DragonHPC/dragon
#!/usr/bin/env python3

"""
Basic performance test of multiprocessing.Pool
"""

import multiprocessing
import time
import test_util

class PoolInvocationTest:

    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Does a basic speed test of dispatching a task in an existingPool')
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_spawn_gather, self.args)
        self.results = test_util.Results(self.args, times, 'pool_invocation_test')
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump() 


    def dummy_func(self, x):
        return 42

    def do_spawn_gather(self):
        """Spawns a bunch of processes in a pool that do nothing and return.

        :return: Time it took to do this in ns
        """

        num_workers = self.args.num_workers

        assert num_workers > 0

        completions = []
        completions2 = []
        with multiprocessing.Pool(num_workers) as the_pool:
            for idx in range(num_workers):
                completions.append( the_pool.apply_async(self.dummy_func, (idx,)))

            for k in completions:
                k.get()

            start_time = time.time_ns()
            for idy in range(num_workers):
                completions2.append(the_pool.apply_async(self.dummy_func, (idy,)))

            for j in completions2:
                j.get()

        return time.time_ns() - start_time


if __name__ == '__main__':
    test = PoolInvocationTest()
    print(test.run())



test/mp_bench/tests/pool_invocation_test.py
DragonHPC/dragon
#!/usr/bin/env python3
"""
Basic test of multiprocessing message passing through multiprocessing.Pipe.
"""

import argparse
import multiprocessing
import time
import test_util

class SendReceiveTest:

    def __init__(self):
        self.args = None
        self.results = None

    def __setup(self):
        parser = test_util.add_default_args('Send and receive a message of varying size using a Pipe')
        parser = test_util.add_connections_args(parser)
        parser.add_argument('--response_size', type=int, default=4, help='size of message pass back')
        self.args = test_util.process_args(parser)
        multiprocessing.set_start_method('spawn')

    def run(self):
        self.__setup()
        times = test_util.iterations(self.do_send_receive, self.args)
        self.results = test_util.Results(self.args, times, "send_receive_test")
        if self.args.json:
            return self.dump()
        return self.results

    def dump(self):
        return self.results.dump()

    def dummy_func(self, out_conn, response):
        out_conn.recv()
        out_conn.send(response)

    def do_send_receive(self):

        num_workers = self.args.num_workers
        #assert self.args.message_size > 0
        message = bytes(self.args.message_size)
        response = bytes(self.args.response_size)
        assert num_workers > 0


        pipeconns = [multiprocessing.Pipe() for _ in range(num_workers)]

        processes = []
        for idx in range(num_workers):
            processes.append(multiprocessing.Process(target=self.dummy_func,
                                                     args=(pipeconns[idx][1], 
                                                     response,)
                                                     ))
        
        start_time = time.time_ns()
        for handle in processes:
            handle.start()
        
        for i in range(len(pipeconns)):
            pipeconns[i][0].send(message)
        for i in range(len(pipeconns)):
            pipeconns[i][0].recv()


        for handle in processes:
            handle.join()

        return time.time_ns() - start_time


if __name__ == '__main__':
    test = SendReceiveTest()
    print(test.run())

test/mp_bench/tests/send_receive_test.py
DragonHPC/dragon
import user_executable
import random
import argparse
from user_parameter_study import ParameterStudy as ParameterStudy

def pi_work_items(ps, num_points):
    for i in range(num_points):
        x = (2*random.random()) - 1
        y = (2*random.random()) - 1
        ps.add_work_item([x,y])

def in_circle(x,y):
    if ((x**2)+(y**2)) < 1:
        return True
    else:
        return False

def process_args():
    parser = argparse.ArgumentParser(description='Get Monte Carlo integration options.')
    parser.add_argument('--num_points', type=int, default=10000,
                            help='number of points to generate')
    parser.add_argument('--num_workers', type=int, default=10,
                            help='number of workers to spawn at a time')
    parser.add_argument('--c_exec', action='store_true',
                            help='use C executable instead of Python function')
    parser.add_argument('--python_exec', action='store_true',
                            help='use Python executable with subprocess')
    parser.add_argument('--func_out_file', action='store_true',
                            help='use Python function from a separate file')
    parser.add_argument('--func_in_file', action='store_true',
                            help='user Python function within file. Default.')
    args = parser.parse_args()
    check_sum = args.c_exec + args.python_exec + args.func_out_file + args.func_in_file
    if check_sum > 1:
        raise Exception("Please choose only one type of executable/function")
    return args



if __name__ == '__main__':
    args = process_args()
    if args.c_exec:
        circ = ParameterStudy(num_dimensions=2, executable='./a.out' ,
                              num_workers=args.num_workers)
    elif args.python_exec:
        circ = ParameterStudy(num_dimensions=2, executable="python user_executable.py",
                              num_workers=args.num_workers)
    elif args.func_out_file:
        circ = ParameterStudy(num_dimensions=2, func=user_executable.in_circle, 
                              num_workers=args.num_workers)
    else:
        circ = ParameterStudy(num_dimensions=2, func=in_circle, 
                              num_workers=args.num_workers)
    num_points = args.num_points
    pi_work_items(circ, num_points)
    circ.go()
    inside = 0
    for item in circ.results:
        if item == True:
            inside += 1
        elif item == False:
            pass
        elif item.strip() == "True":
            inside += 1
    print("Estimated value of pi: ",(inside/num_points)*4,"\n")
    print(circ.performance_info(output="str"))


test/mp_bench/user_parameter_study/monte_carlo_pi_param_study.py
DragonHPC/dragon
#!/usr/bin/env python
import sys

def in_circle(x,y):
    if ((x**2)+(y**2)) < 1:
        return True
    else:
        return False

if __name__ == '__main__':
    print(in_circle(float(sys.argv[1]), float(sys.argv[2])))

test/mp_bench/user_parameter_study/user_executable.py
DragonHPC/dragon
import multiprocessing as mp
import time
import subprocess
import pickle
import datetime
import json

class ParameterStudy:
    def __init__(self, num_dimensions=1, executable=None, func=None, num_workers=4):
        if executable==None and func==None:
            raise ParameterStudy.ParameterStudyException("No function or executable passed")
        if executable != None and func != None:
            raise ParameterStudy.ParameterStudyException("Passed both a function and an executable")
        self.__num_dims = num_dimensions
        self.__executable = executable
        self.__func = func
        self.__num_workers = num_workers
        self.__parameter_space = []
        self.results = None
        self.total_time = None
        self.compute_time = None
        self.data_dict = {}

    def add_work_item(self, parameters):
        if (len(parameters) != self.__num_dims):
            fail = f'{parameters} not of length {self.__num_dims}'
            raise ParameterStudy.ParameterStudyException(fail)
        if self.__executable != None:
            parameters = list(map(str, parameters))
        self.__parameter_space.append(parameters)
   
    def shell_func(self, parameters_list):
        if self.__executable == None:
            result = self.__func(*parameters_list)
        else:
            output = subprocess.Popen(self.__executable.split()+parameters_list,stdout=subprocess.PIPE)
            result = output.communicate()[0].decode()
        return result


    #runs test returns results
    def go(self):
        mp.set_start_method('spawn')
        start_time_setup = time.time_ns()
        with mp.Pool(self.__num_workers) as the_pool:
            start_time_no_setup = time.time_ns()
            map_results = the_pool.map(self.shell_func, self.__parameter_space)
            end_time_no_setup = time.time_ns()
        end_time_setup = time.time_ns()

        self.results = list(map_results)
        self.total_time = end_time_setup - start_time_no_setup
        self.compute_time = end_time_no_setup - start_time_no_setup
    
    def performance_info(self, output='dict'):
        self.__process_data()
        if output == 'dict':
            return self.data_dict
        elif output == 'json':
            return json.dumps(self.data_dict, indent=4, default=str)
        elif output == 'str':
            return self.__string_result()
        else:
            raise ParameterStudy.ParameterStudyException('Invalid output type ("dict", "str", or "json")')

    def __dict_setup(self):
        self.data_dict['hostname'] = subprocess.check_output('hostname').decode().strip()
        self.data_dict['date'] = datetime.datetime.now()
        self.data_dict['function'] = self.__func
        self.data_dict['executable'] = self.__executable
        self.data_dict['num_workers'] = self.__num_workers
        self.data_dict['work_items'] = len(self.__parameter_space)
        self.data_dict['avg_response_size'] = len(pickle.dumps(self.results)) / self.data_dict['work_items']
        self.data_dict['avg_message_size'] = len(pickle.dumps(self.__parameter_space)) / self.data_dict['work_items']
        self.data_dict['latency'] = self.total_time / 1000
        self.data_dict['compute_time'] = self.compute_time / 1000
        self.data_dict['setup_time'] = self.data_dict['latency'] - self.data_dict['compute_time']

    def __process_data(self):
        self.__dict_setup()
        transmitted_bytes = len(pickle.dumps(self.results)) + len(pickle.dumps(self.__parameter_space))
        avg_compute_time_per_item = self.compute_time / self.data_dict['work_items']
        avg_setup_time_pw = self.data_dict['setup_time'] / self.__num_workers
        avg_bandwidth = transmitted_bytes / self.compute_time

        self.data_dict['avg_setup_time_per_worker'] = avg_setup_time_pw / 1000
        self.data_dict['avg_compute_time_per_work_item'] = avg_compute_time_per_item / 1000
        self.data_dict['avg_bandwidth'] = avg_bandwidth * (1000000000 / (2**20))

    def __string_result(self):
        return (
                f'Parameter Study Performance\n'
                f'Hostname: {self.data_dict["hostname"]}\n'
                f'Date: {self.data_dict["date"]}\n'
                f'Function: {self.data_dict["function"]}\n'
                f'Executable: {self.data_dict["executable"]}\n'
                f'Number of workers: {self.data_dict["num_workers"]}\n'
                f'Number of work items: {self.data_dict["work_items"]}\n'
                f'Average message size (bytes): {self.data_dict["avg_message_size"]}\n'
                f'Average response size (bytes): {self.data_dict["avg_response_size"]}\n'
                f'Latency (\xb5s): {self.data_dict["latency"]}\n'
                f'Compute time (\xb5s): {self.data_dict["compute_time"]}\n'
                f'Setup time (\xb5s): {self.data_dict["setup_time"]}\n'
                f'Average setup time per worker (\xb5s): {self.data_dict["avg_setup_time_per_worker"]}\n'
                f'Average compute time per work item (\xb5s): {self.data_dict["avg_compute_time_per_work_item"]}\n'
                f'Average bandwidth (mB/s): {self.data_dict["avg_bandwidth"]}'
               )

    class ParameterStudyException(Exception):
        pass

test/mp_bench/user_parameter_study/user_parameter_study.py
DragonHPC/dragon
import argparse
import multiprocessing
import subprocess
import json
import datetime

def get_hostname():
    return subprocess.check_output('hostname').decode()

def add_default_args(desc):
    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument('--num_workers', type=int, default=4, 
                        help='number of workers to spawn at a time')
    parser.add_argument('--num_iterations', type=int, default=4, 
                        help='number of iterations to do')
    parser.add_argument('--burn_iterations', type=int, default=1, 
                        help='number of iterations to burn first time')
    parser.add_argument('--json', action='store_true',
                        help='output json')
    return parser

def add_connections_args(parser):
    parser.add_argument('--message_size', type=int, default=4, 
                        help='size of message to pass in')
    return parser

def add_token_args(parser):
    parser = add_connections_args(parser)
    parser.add_argument('--spins', type=int, default=100,
                        help='number of times to send the token around a chain')
    parser.add_argument('--num_tokens', type=int, default=1, 
                        help='number of tokens to pass around the ring')
    return parser

def process_args(parser):
    args = parser.parse_args()
    assert args.num_workers > 0
    assert args.num_iterations > 0
    #if not args.json:
    #    print('Testing with {} workers, {} iterations'.format(args.num_workers, args.num_iterations))
    return args

def iterations(fcn, args):
    for _ in range(args.burn_iterations):
        _ = fcn()
    
    iteration_times = []

    for _ in range(args.num_iterations):
        this_time_ns = fcn()
        iteration_times.append(this_time_ns)

    return iteration_times

class Results:
    def __init__(self, args, times, test_type):
        #self.args = args
        self.times = times
        self.num_iterations = args.num_iterations
        self.num_workers = args.num_workers
        self.times_pw = [] 
        self.avg_time = 0
        self.avg_time_pw = 0
        self.hostname = get_hostname().strip()
        self.date = datetime.datetime.now()
        self.test_type = test_type
        try:
            self.spins = args.spins
        except:
            self.spins = None
        try: 
            self.message_size = args.message_size
            self.avg_bandwidth = 0
            self.bandwidth = []
        except:
            self.message_size = None
            self.bandwidth = None
            self.avg_bandwidth = None
        try:
            self.response_size = args.response_size
        except:
            self.response_size = None
        try:
            self.num_tokens = args.num_tokens
        except:
            self.num_tokens = None
        self.times_ps = []
        self.avg_time_ps = 0
        self.__total = 0
        self.__total_pw = 0
        self.__total_ps = 0
        self.__total_bw = 0
        self.__process_data()

    def __process_data(self):
        for i in range(len(self.times)):
            self.__total = self.__total + self.times[i]
            self.times_pw.append(self.times[i] / self.num_workers)
            self.__total_pw = self.__total_pw + self.times_pw[i]
            if self.spins is not None:
                self.__process_spin_data(i)
            if self.response_size is not None:
                self.__process_size_data(i)
            self.times[i] = self.times[i] / 1000
            self.times_pw[i] = self.times_pw[i] / 1000

        self.avg_time = (self.__total / self.num_iterations) / 1000
        self.avg_time_pw = (self.__total_pw / self.num_iterations) / 1000
        self.avg_time_ps = (self.__total_ps / self.num_iterations) / 1000
        self.avg_bandwidth = (self.__total_bw / self.num_iterations) * (1000000000 / (2**20))

    def __process_spin_data(self, i):
        self.times_ps.append(self.times[i] / self.spins)
        self.__total_ps = self.__total_ps + self.times_ps[i]
        self.times_ps[i] = self.times_ps[i] / 1000
        self.bandwidth.append((self.message_size * self.num_workers * self.spins) / self.times[i])
        if self.num_tokens is not None:
            self.bandwidth[i] = self.bandwidth[i] * self.num_tokens
        self.__total_bw = self.__total_bw + self.bandwidth[i]
        self.bandwidth[i] = self.bandwidth[i] * (1000 / (2**20))

    def __process_size_data(self, i):
        self.bandwidth.append(((self.message_size + self.response_size) * self.num_workers) / self.times[i])
        self.__total_bw = self.__total_bw + self.bandwidth[i]
        self.bandwidth[i] = self.bandwidth[i] * (1000 / (2**20))

    def not__str__(self):
        if self.response_size is not None: 
            return (
                f'Test: {self.test_type}\n'
                f'Hostname: {self.hostname}\n'
                f'Date: {self.date}\n'
                f'Message size (bytes): {self.message_size}\n'
                f'Response size (bytes): {self.response_size}\n'
                f'Number of workers: {self.num_workers}\n'
                f'Average total time (\xB5s): {self.avg_time}\n'
                f'Average time per worker (\xB5s): {self.avg_time_pw}\n'
                f'Average bandwidth (MiB/s): {self.avg_bandwidth}\n'
            )

        elif self.num_tokens is not None:
            return (
                f'Test: {self.test_type}\n'
                f'Hostname: {self.hostname}\n'
                f'Date: {self.date}\n'
                f'Message size (bytes): {self.message_size}\n'
                f'Number of workers: {self.num_workers}\n'
                f'Number of spins: {self.spins}\n'
                f'Number of tokens: {self.num_tokens}\n'
                f'Average total time (\xb5s): {self.avg_time}\n'
                f'Average time per worker (\xb5s): {self.avg_time_pw}\n'
                f'Average time per spin (\xb5s): {self.avg_time_ps}\n'
                f'Average bandwidth (MiB/s): {self.avg_bandwidth}\n'
            )

        elif self.spins is not None:
            return (
                f'Test: {self.test_type}\n'
                f'Hostname: {self.hostname}\n'
                f'Date: {self.date}\n'
                f'Message size (bytes): {self.message_size}\n'
                f'Number of workers: {self.num_workers}\n'
                f'Number of spins: {self.spins}\n'
                f'Average total time (\xb5s): {self.avg_time}\n'
                f'Average time per worker (\xb5s): {self.avg_time_pw}\n'
                f'Average time per spin (\xb5s): {self.avg_time_ps}\n'
                f'Average bandwidth (MiB/s): {self.avg_bandwidth}\n'
            )             
        
        else:
            return (
                f'Test: {self.test_type}\n'
                f'Hostname: {self.hostname}\n'
                f'Date: {self.date}\n'
                f'Number of workers: {self.num_workers}\n'
                f'Average total time (\xB5s): {self.avg_time}\n'
                f'Average time per worker (\xB5s): {self.avg_time_pw}\n'
            )
    def __str__(self):
        return f'{self.avg_time}'

    def dump(self):
        data_dict = {}
        data_dict["test_type"] = self.test_type
        data_dict["hostname"] = self.hostname
        data_dict["date"] = self.date
        data_dict["num_workers"] = self.num_workers
        #data_dict["times"] = self.times
        #data_dict["times_pw"] = self.times_pw
        data_dict["avg_time"] = self.avg_time
        data_dict["avg_time_pw"] = self.avg_time_pw
        if self.spins is not None:
            data_dict["spins"] = self.spins
            #data_dict["times_ps"] = self.times_ps
            data_dict["avg_time_ps"] = self.avg_time_ps
            data_dict["message_size"] = self.message_size
            #data_dict["bandwidth"] = self.bandwidth
            data_dict["avg_bandwidth"] = self.avg_bandwidth
        if self.response_size is not None:
            data_dict["response_size"] = self.response_size
            data_dict["message_size"] = self.message_size
            data_dict["avg_bandwidth"] = self.avg_bandwidth
        if self.num_tokens is not None:
            data_dict["num_tokens"] = self.num_tokens
        return json.dumps(data_dict, indent=4, default=str)

test/mp_bench/util/test_util.py
DragonHPC/dragon
""" This file contains tests of the mpbridge interfaces, i.e. classes, methods, attributes and how they
are hooked into multiprocessing context.

This is driven by the requirement that Python multiprocessing should **just work** with Dragon.
"""

import unittest
import copy

import dragon
import multiprocessing as mp

# these parts of the "public" API are apparently not reachable through the context or top level import
from multiprocessing.managers import BaseManager, SyncManager, BaseProxy, Token, SharedMemoryManager
from multiprocessing.shared_memory import SharedMemory, ShareableList
from multiprocessing.util import get_logger, log_to_stderr
from multiprocessing.heap import Arena, BufferWrapper


def setUpModule():
    mp.set_start_method("dragon", force=True)


class TestMultiprocessingAPI(unittest.TestCase):
    """An attempt at testing that we are monkeypatching correctly.
    In particular, that that the multiprocessing API is correctly
    connected to the Dragon Native API.
    """

    def setUp(self):
        self.assertEqual(mp.get_start_method(), "dragon")

    def test_monkeypatching(self):

        ctx = mp.get_context()
        self.assertTrue(ctx._name == "dragon")

        # enable all Dragon API
        ctx_backup = copy.copy(ctx)  # save orig for copy back later
        ctx.USE_MPFUNCTIONS = ctx.USE_MPQUEUE = ctx.USE_MPJQUEUE = ctx.USE_MPSQUEUE = False
        ctx.USE_MPPOOL = ctx.USE_MPPIPE = ctx.USE_MPMANAGER = ctx.USE_MPLOCK = False
        ctx.USE_MPRLOCK = ctx.USE_MPCONDITION = ctx.USE_MPSEMAPHORE = ctx.USE_MPBOUNDSEMAPHORE = False
        ctx.USE_MPEVENT = ctx.USE_MPBARRIER = ctx.USE_MPVALUE = ctx.USE_MPRAWVALUE = ctx.USE_MPARRAY = False
        ctx.USE_MPRAWARRAY = ctx.USE_MPSHAREDMEMORY = ctx.USE_MPSHAREABLELIST = ctx.USE_MPLISTENER = False
        ctx.USE_MPCLIENT = False

        # Functions
        self.assertRaises(NotImplementedError, mp.freeze_support)  # win32 only
        self.assertTrue(callable(getattr(mp, "cpu_count")))
        self.assertTrue(callable(getattr(mp, "log_to_stderr")))
        self.assertTrue(callable(getattr(mp, "get_logger")))
        self.assertTrue(callable(log_to_stderr))
        self.assertTrue(callable(get_logger))

        self.assertTrue(mp.connection.wait == mp.get_context().wait)
        self.assertTrue(mp.pool.wait == mp.get_context().wait)

        self.assertTrue(mp.reduction.DupFd == dragon.mpbridge.reduction.dragon_DupFd)
        self.assertTrue(mp.reduction.recvfds == dragon.mpbridge.reduction.dragon_recvfds)
        self.assertTrue(mp.reduction.sendfds == dragon.mpbridge.reduction.dragon_sendfds)
        self.assertTrue(mp.get_context().reducer.DupFd == dragon.mpbridge.reduction.dragon_DupFd)
        self.assertTrue(mp.get_context().reducer.recvfds == dragon.mpbridge.reduction.dragon_recvfds)
        self.assertTrue(mp.get_context().reducer.sendfds == dragon.mpbridge.reduction.dragon_sendfds)

        # Classes
        self.assertIsInstance(mp.Queue(), dragon.mpbridge.queues.DragonQueue)
        self.assertIsInstance(mp.SimpleQueue(), dragon.mpbridge.queues.DragonSimpleQueue)
        self.assertIsInstance(mp.JoinableQueue(), dragon.mpbridge.queues.DragonJoinableQueue)

        self.assertIsInstance(mp.Semaphore(), dragon.mpbridge.synchronize.DragonSemaphore)
        self.assertIsInstance(mp.BoundedSemaphore(), dragon.mpbridge.synchronize.DragonBoundedSemaphore)
        self.assertIsInstance(mp.Lock(), dragon.mpbridge.synchronize.DragonLock)
        self.assertIsInstance(mp.RLock(), dragon.mpbridge.synchronize.DragonRLock)
        self.assertIsInstance(mp.Event(), dragon.mpbridge.synchronize.DragonEvent)
        self.assertIsInstance(mp.Value(typecode_or_type="i"), dragon.mpbridge.sharedctypes.DragonValue)
        self.assertIsInstance(mp.RawValue(typecode_or_type="i"), dragon.mpbridge.sharedctypes.DragonRawValue)
        self.assertIsInstance(mp.Barrier(5), dragon.mpbridge.synchronize.DragonBarrier)
        self.assertIsInstance(mp.Condition(), dragon.mpbridge.synchronize.DragonCondition)
        self.assertIsInstance(mp.Array('i', []), dragon.mpbridge.sharedctypes.DragonArray)
        self.assertIsInstance(mp.RawArray('i', []), dragon.mpbridge.sharedctypes.DragonRawArray)

        self.assertIsInstance(Arena(1024), dragon.mpbridge.heap.DragonArena)
        self.assertIsInstance(BufferWrapper(1024)._heap._arenas[0], dragon.mpbridge.heap.DragonArena)

        with mp.Pool(1) as dut:  # it is not correct to rely on the garbage collector to destroy the pool
            self.assertIsInstance(dut, dragon.mpbridge.pool.DragonPool)

        self.assertTrue(
            callable(getattr(mp, "Pipe"))
        )  # this should be dragon.mpbridge.connection.DragonConnection
        r, w = mp.Pipe()
        self.assertIsInstance(r, dragon.infrastructure.connection.Connection)  # TODO: PE-41642
        self.assertIsInstance(w, dragon.infrastructure.connection.Connection)

        # Missing features
        self.assertRaises(NotImplementedError, mp.Manager)
        self.assertRaises(NotImplementedError, BaseManager)
        self.assertRaises(NotImplementedError, SyncManager)
        self.assertRaises(NotImplementedError, BaseProxy)
        self.assertRaises(NotImplementedError, Token)
        self.assertRaises(NotImplementedError, SharedMemoryManager)

        self.assertRaises(NotImplementedError, SharedMemory)
        self.assertRaises(NotImplementedError, ShareableList)

        self.assertRaises(NotImplementedError, mp.connection.Listener)
        self.assertRaises(NotImplementedError, mp.connection.Client, "")

        # copy back original values
        ctx.USE_MPFUNCTIONS = ctx_backup.USE_MPFUNCTIONS
        ctx.USE_MPJQUEUE = ctx_backup.USE_MPJQUEUE
        ctx.USE_MPMANAGER = ctx_backup.USE_MPMANAGER
        ctx.USE_MPLOCK = ctx_backup.USE_MPLOCK
        ctx.USE_MPRLOCK = ctx_backup.USE_MPRLOCK
        ctx.USE_MPCONDITION = ctx_backup.USE_MPCONDITION
        ctx.USE_MPSEMAPHORE = ctx_backup.USE_MPSEMAPHORE
        ctx.USE_MPBOUNDSEMAPHORE = ctx_backup.USE_MPBOUNDSEMAPHORE
        ctx.USE_MPEVENT = ctx_backup.USE_MPEVENT
        ctx.USE_MPBARRIER = ctx_backup.USE_MPBARRIER
        ctx.USE_MPVALUE = ctx_backup.USE_MPVALUE
        ctx.USE_MPRAWVALUE = ctx_backup.USE_MPRAWVALUE
        ctx.USE_MPARRAY = ctx_backup.USE_MPARRAY
        ctx.USE_MPRAWARRAY = ctx_backup.USE_MPRAWARRAY
        ctx.USE_MPSHAREDMEMORY = ctx_backup.USE_MPSHAREDMEMORY
        ctx.USE_MPSHAREABLELIST = ctx_backup.USE_MPSHAREABLELIST
        ctx.USE_MPLISTENER = ctx_backup.USE_MPLISTENER
        ctx.USE_MPCLIENT = ctx_backup.USE_MPLISTENER

    def test_queue_public_interfaces(self):
        """Test that our metaclasses create the correct public interfaces, i.e.
        clean the dragon.native objects correctly of any non-standard attributes
        """

        import dragon.mpbridge.queues

        sq = mp.SimpleQueue()
        sq_attr = ["close", "empty", "get", "put"]
        self.assertListEqual([x for x in dir(sq) if not x.startswith("_")], sq_attr)

        q = mp.Queue()
        q_attr = [
            "cancel_join_thread",
            "close",
            "empty",
            "full",
            "get",
            "get_nowait",
            "join_thread",
            "put",
            "put_nowait",
            "qsize",
        ]
        self.assertListEqual([x for x in dir(q) if not x.startswith("_")], q_attr)

        jq = mp.JoinableQueue()
        jq_attr = [
            "cancel_join_thread",
            "close",
            "empty",
            "full",
            "get",
            "get_nowait",
            "join",
            "join_thread",
            "put",
            "put_nowait",
            "qsize",
            "task_done",
        ]

        self.assertListEqual([x for x in dir(jq) if not x.startswith("_")], jq_attr)


class TestMultiprocessingInternalPatching(unittest.TestCase):
    """Test that internal instances of Multiprocessing classes correctly reference
    Dragon components, not Multiprocessing components.
    """

    def test_condition(self):
        cond = mp.Condition()
        self.assertIsInstance(cond._lock, dragon.mpbridge.synchronize.DragonRLock)
        self.assertIsInstance(cond._sleeping_count, dragon.mpbridge.synchronize.DragonSemaphore)
        self.assertIsInstance(cond._woken_count, dragon.mpbridge.synchronize.DragonSemaphore)
        self.assertIsInstance(cond._wait_semaphore, dragon.mpbridge.synchronize.DragonSemaphore)


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_api.py
DragonHPC/dragon
#!/usr/bin/env python3

import dragon
import unittest
from dragon.native.lock import Lock
from dragon.mpbridge.sharedctypes import Array
import ctypes
import multiprocessing as mp


class TestArray(unittest.TestCase):
    def test_requirement_1_1(self,):
        """tests that the attributes get_obj and get_lock are not available"""
        self.assertFalse(hasattr(Array("c",range(10),lock=False,),"get_lock",))
        self.assertFalse(hasattr(Array("c",range(10),lock=False,),"get_obj",))

    def test_requirement_1_2(self,):
        """tests that the rlock is set properly"""
        array = Array("c", range(10), lock=True,)
        lock = array.get_lock()
        self.assertIsInstance(lock, type(Lock(recursive=True)),)
        #tests recursive attributes
        self.assertTrue(hasattr(lock,"_recursive",))
        self.assertTrue(hasattr(lock,"is_recursive",))


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_array.py
DragonHPC/dragon
import unittest
import time
import threading

import dragon
from dragon.mpbridge.synchronize import Barrier, BrokenBarrierError


class TestBarrier(unittest.TestCase):
    def test_requirement_1_3(self):
        """Requirement 1.3: If the method is called with timeout <0, it has to assume timeout=0."""
        start = time.monotonic()
        Barrier(parties=2, timeout=-0.1)
        stop = time.monotonic()
        timeout = int(int(stop) - int(start))
        self.assertAlmostEqual(0, timeout, "The negative timeout for a barrier is 0.")

    def test_requirement_2_1(self):
        """If the wait is called with timeout <0, it has to assume timeout=0."""
        start = float(time.monotonic())
        barrier = Barrier(parties=2)
        try:
            barrier.wait(timeout=-1)
        except threading.BrokenBarrierError:
            pass
        stop = float(time.monotonic())
        timeout = round(float(float(stop) - float(start)), 1)
        self.assertAlmostEqual(0, timeout, "The negative timeout for a barrier is 0.")


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_barrier.py
DragonHPC/dragon
from time import sleep
import unittest
import pickle
import os

import dragon
import multiprocessing as mp


class TestCondition(unittest.TestCase):
    def setUp(self):
        self.assertEqual(mp.get_start_method(), "dragon")

    @staticmethod
    def cond_wait(cond: mp.Condition, queue: mp.Queue, item: object) -> None:
        queue.put(item)

        assert cond.acquire(timeout=None)
        cond.wait(timeout=None)
        cond.release()

        queue.put(item)

    # Test cases
    def test_basic_condition_test(self):
        """
        Basic condition test that checks that acquire, wait, and release work.
        """
        condition = mp.Condition()
        result = condition.acquire()
        self.assertTrue(result)
        condition.release()

    def test_process_multiple_notify(self):
        """
        Condition test that passes a condition to multiple processes. Notify is called
        by the parent and the first process exits.
        The other processes continue to run until is notify_all called on them.
        """

        cond = mp.Condition()
        queue = mp.Queue()
        processes = []

        nproc = 4

        # start a bunch of processes
        for i in range(nproc):
            proc = mp.Process(target=self.cond_wait, args=(cond, queue, i))
            proc.start()
            processes.append(proc)
            j = queue.get(timeout=10)  # make sure process is actually running

        # notify one of them
        cond.acquire()
        cond.notify()
        cond.release()

        # check that one of them exited
        i = queue.get(timeout=10)
        processes[i].join(timeout=None)
        self.assertTrue(processes[i].exitcode == 0)

        # notify them all and check they all joined
        cond.acquire()
        cond.notify_all()
        cond.release()

        for _ in range(nproc - 1):
            i = queue.get(timeout=10)
            processes[i].join(timeout=None)
            self.assertTrue(processes[i].exitcode == 0)

    def test_lock_hack_for_condition(self):
        """Check that we have correctly faked a _semlock object in _lock"""
        cond = mp.Condition()

        self.assertTrue(cond.acquire == cond._lock.acquire)
        self.assertTrue(cond.release == cond._lock.release)

        success = cond.acquire()

        self.assertTrue(cond._lock._semlock._is_mine())
        self.assertTrue(1 == cond._lock._semlock._count())

        bytes = pickle.dumps(cond)
        cond2 = pickle.loads(bytes)
        self.assertTrue(cond2._lock._semlock._is_mine())


if __name__ == "__main__":

    mp.set_start_method("dragon", force=True)
    unittest.main()

test/mpbridge/test_condition.py
DragonHPC/dragon
import unittest

import dragon
import multiprocessing
import multiprocessing.heap


class TestMultiprocessingAPI(unittest.TestCase):
    def test_arena(self):
        size = 1024
        fd = 12
        a = multiprocessing.heap.Arena(size, fd=fd)

        self.assertTrue(a.size == size)
        self.assertTrue(a.fd == fd)
        self.assertIsInstance(a.buffer, memoryview)

    def test_with_BufferWrapper(self):
        bw = multiprocessing.heap.BufferWrapper(8)
        memory = bw.create_memoryview()
        self.assertIsInstance(memory, memoryview)


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_heap.py
DragonHPC/dragon
#!/usr/bin/env python3
import unittest

import dragon
import multiprocessing as mp


def setUpModule():
    mp.set_start_method("dragon", force=True)


class TestDragonLocks(unittest.TestCase):
    def setUp(self):
        self.assertEqual(mp.get_start_method(), "dragon")

    def test_semlock_hack(self):

        rlock = mp.RLock()

        success = rlock.acquire()
        self.assertTrue(success)

        self.assertTrue(rlock._semlock._is_mine())
        self.assertTrue(1 == rlock._semlock._count())

        for i in range(13):
            self.assertTrue(rlock.acquire())
            self.assertTrue(i + 2 == rlock._semlock._count())


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_lock.py
DragonHPC/dragon
import dragon
from dragon import mpbridge
import multiprocessing as mp
import os
import unittest


mpbridge.inspections = True  # Inspect and report origin of any direct invocations of os.pipe
# logger = mp.log_to_stderr()
# logger.setLevel(mp.SUBWARNING)
# logger.setLevel(mp.SUBDEBUG)  # Verbosity to aid detailed debugging.


def sleepy_getpid(delay):
    """Convenient function to use as target when creating new processes."""
    import time, os

    time.sleep(delay)
    mp.util.info(f"sleepy_getpid done on {os.getpid()}")
    return os.getpid()


class TestStartMethod(unittest.TestCase):
    def test_AAA_available_start_method(self):
        start_methods = mp.get_all_start_methods()
        self.assertTrue(mpbridge.DragonPopen.method in start_methods)
        self.assertTrue("spawn" in start_methods)
        self.assertTrue("fork" in start_methods)
        self.assertTrue("forkserver" in start_methods)

    def test_AAB_set_dragon_start_method(self):
        mp.set_start_method("spawn", force=True)
        self.assertFalse(mp.get_start_method(allow_none=True) == "dragon")
        mp.set_start_method(mpbridge.DragonPopen.method, force=True)
        self.assertEqual(mp.get_start_method(allow_none=True), "dragon")

    def test_dragon_is_available_via_default_context(self):
        from dragon.mpbridge.monkeypatching import AugmentedDefaultContext

        self.assertTrue(isinstance(mp.context._default_context, AugmentedDefaultContext))

    # check we're no longer using file descriptors
    @unittest.skipIf("DRAGON_MY_PUID" not in os.environ, "Not launched via dragon")
    @unittest.skip("filed as part of CIRRUS-1473")
    def test_DragonProcess_pool_map(self):
        mp.set_start_method(mpbridge.DragonPopen.method, force=True)
        with mp.Pool(2) as pool:
            results = pool.map(sleepy_getpid, [0.1, 0.3, 0.2])
            self.assertEqual(len(mp.process._children), 2)
            for live_process in mp.process._children:
                mp.util.info(f"{live_process}:\n\t{live_process._popen._fds}")
                self.assertEqual(len(live_process._popen._fds), 0)
        del pool
        self.assertEqual(len(results), 3)
        self.assertEqual(len(set(results)), 2)
        mp.set_start_method(mpbridge.DragonPopen.method, force=True)

    @unittest.skip(f"Currently fails, mp.Process contains _start_method, PE-37293")
    def test_basic_process_start_join(self):
        mp.set_start_method(mpbridge.DragonPopen.method, force=True)
        p1 = mpbridge.DragonProcess(target=sleepy_getpid, args=(0.1,))
        p2 = mp.Process(target=sleepy_getpid, args=(0.1,))
        p1.start()
        p2.start()
        p1.join()
        self.assertEqual(p1.exitcode, 0)
        p2.join()
        self.assertEqual(p2.exitcode, 0)
        self.assertEqual(p1._start_method, mpbridge.DragonPopen.method)
        self.assertEqual(p2._start_method, None)
        self.assertTrue(isinstance(p1, mpbridge.DragonProcess))
        self.assertTrue(isinstance(p2, mp.context.Process))

    @classmethod
    def _check_context(cls, conn):
        conn.send(mp.get_start_method())

    def check_context(self, ctx):
        r, w = ctx.Pipe(duplex=False)
        p = ctx.Process(target=self._check_context, args=(w,))
        p.start()
        w.close()
        child_method = r.recv()
        r.close()
        p.join()
        self.assertEqual(child_method, ctx.get_start_method())

    def test_set_get(self):
        # To speed up tests when using the forkserver, we can preload these:
        PRELOAD = ["__main__", "test.test_multiprocessing_forkserver"]
        mp.set_forkserver_preload(PRELOAD)
        count = 0
        # mp.set_start_method('fork', force=True)
        old_method = mp.get_start_method()
        try:
            # for method in ('fork', 'spawn', 'forkserver'):  # Modified to avoid issue when testing on WSL.
            for method in ("fork", "spawn"):
                try:
                    mp.set_start_method(method, force=True)
                except ValueError:
                    raise
                self.assertEqual(mp.get_start_method(), method)
                ctx = mp.get_context()
                self.assertEqual(ctx.get_start_method(), method)
                self.assertTrue(type(ctx).__name__.lower().startswith(method))
                self.assertTrue(ctx.Process.__name__.lower().startswith(method))
                self.check_context(mp)
                count += 1
        finally:
            mp.set_start_method(old_method, force=True)
        self.assertGreaterEqual(count, 1)

    def test_get_all(self):
        import sys

        methods = mp.get_all_start_methods()
        if sys.platform == "win32":
            self.assertEqual(methods, ["spawn"])
        else:
            self.assertTrue("dragon" in methods)
            methods.remove("dragon")
            self.assertTrue(
                methods == ["fork", "spawn"]
                or methods == ["spawn", "fork"]
                or methods == ["fork", "spawn", "forkserver"]
                or methods == ["spawn", "fork", "forkserver"]
            )


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_mpbridge_basic.py
DragonHPC/dragon
""" Test for our context wait.
When the wait is lowered into dragon.native, these tests should move along.
"""

import unittest
from test.support import socket_helper

import socket
import time
import random
import os
import random
from sys import platform

import dragon

import multiprocessing as mp
import threading

import dragon.mpbridge.context


def setUpModule():
    mp.set_start_method("dragon", force=True)


class TestDragonContextWait(unittest.TestCase):
    """Software tests for Dragons wait() working on PUIDs, Connections, CUIDs,
    DragonQueues, DragonQueue.Fakers and sockets.
    """

    def setUp(self):
        self.assertEqual(mp.get_start_method(), "dragon")

    @classmethod
    def _waitev_and_ready(self, ev, send_handle, msg):

        ev.wait(timeout=None)

        # when many processes are waiting at the same time, randomize the return here,
        # so that the parent can get one or several at the same time.
        time.sleep(random.uniform(0, 0.1))

        if send_handle is None:  # sentinel
            pass  # just exit
        elif isinstance(send_handle, dragon.infrastructure.connection.Connection):
            send_handle.send(msg)
        elif isinstance(send_handle, dragon.infrastructure.connection.CUID):
            send_handle.send(msg)
        elif isinstance(send_handle, dragon.mpbridge.queues.DragonQueue):
            send_handle.put(msg)
        elif isinstance(send_handle, dragon.mpbridge.queues.BaseImplQueue):
            send_handle.put(msg)
        elif isinstance(send_handle, dragon.mpbridge.queues.BaseImplJoinableQueue):
            send_handle.put(msg)
        elif isinstance(send_handle, dragon.mpbridge.queues.BaseImplSimpleQueue):
            send_handle.put(msg)
        elif isinstance(send_handle, dragon.mpbridge.queues.FakeConnectionHandle):
            send_handle.send(msg)
        elif isinstance(send_handle, tuple):  # socket address
            s = socket.socket()
            s.connect(send_handle)
            s.sendall(("%s" % msg).encode("ascii"))
            s.close()
        else:
            raise (TypeError(f"Don't know how to set ready for a handle of type: {type(send_handle)}"))

    def test_sentinel(self):

        num_workers = 4

        events = [dragon.native.event.Event() for _ in range(num_workers)]
        procs = []

        for i in range(num_workers):
            p = dragon.mpbridge.DragonProcess(target=self._waitev_and_ready, args=(events[i], None, ""))
            procs.append(p)

        for p in procs:
            p.start()

        sentinels = [p.sentinel for p in procs]

        # they should all be sleeping
        ready_sentinels = dragon.mpbridge.DragonContext.wait(sentinels, timeout=0)
        self.assertEqual(ready_sentinels, [])

        # they should wake up and exit roughly every 0.1 seconds
        while sentinels:

            events[0].set()

            ready_sentinels = dragon.mpbridge.DragonContext.wait(sentinels, timeout=None)

            self.assertNotEqual(ready_sentinels, [])
            self.assertEqual(len(ready_sentinels), 1)

            s = ready_sentinels[0]
            self.assertIsInstance(s, dragon.mpbridge.process.PUID)
            self.assertTrue(s == sentinels[0])

            sentinels.pop(0)
            events.pop(0)

        # make sure we shut down all service threads of WaitService class
        # before moving to the next test
        dragon.mpbridge.context._wait_service.__del__()
        # Reset global _wait_service so that init() is called again
        # and the garbage_collect thread initiated
        dragon.mpbridge.context._wait_service = None

    def test_connection(self):

        num_workers = 4

        impbridge_ctx = mp.get_context()  # Pipe() requires a context object

        pipes = [impbridge_ctx.Pipe(duplex=False) for _ in range(num_workers)]
        readers = [pipes[i][0] for i in range(num_workers)]
        writers = [pipes[i][1] for i in range(num_workers)]

        messages = [f"Test-Message-{i}" for i in range(num_workers)]
        events = [dragon.native.event.Event() for _ in range(num_workers)]

        procs = []

        for i in range(num_workers):
            p = dragon.mpbridge.DragonProcess(
                target=self._waitev_and_ready, args=(events[i], writers[i], messages[i])
            )
            procs.append(p)

        # there should be no message yet
        self.assertFalse(dragon.mpbridge.DragonContext.wait(readers, timeout=0))

        for p in procs:
            p.start()

        #  a new message should be ready every 0.1 sec
        while readers:
            events[0].set()
            ready_readers = dragon.mpbridge.DragonContext.wait(readers, timeout=10)

            self.assertNotEqual(ready_readers, [])

            self.assertEqual(len(ready_readers), 1)
            r = ready_readers[0]
            self.assertEqual(r, readers[0])
            msg = r.recv()  # this shouldn't block now
            self.assertEqual(msg, messages[0])
            messages.pop(0)
            readers.pop(0)
            events.pop(0)

        # make sure we shut down all service threads of WaitService class
        # before moving to the next test
        dragon.mpbridge.context._wait_service.__del__()
        # Reset global _wait_service so that init() is called again
        # and the garbage_collect thread initiated
        dragon.mpbridge.context._wait_service = None

    def test_queue(self):

        num_workers = 4

        impbridge_ctx = mp.get_context()  # Queue() requires a context object

        queues = [impbridge_ctx.Queue() for _ in range(num_workers)]
        items = [f"Test-Item-{i}" for i in range(num_workers)]
        events = [dragon.native.event.Event() for _ in range(num_workers)]

        procs = []

        for i in range(num_workers):
            p = dragon.mpbridge.DragonProcess(
                target=self._waitev_and_ready, args=(events[i], queues[i], items[i])
            )
            procs.append(p)

        # there should be no message yet
        self.assertFalse(dragon.mpbridge.DragonContext.wait(queues, timeout=0.0))

        for p in procs:
            p.start()

        # a queue should be ready every 0.1 sec
        while queues:
            events[0].set()
            ready_queues = dragon.mpbridge.DragonContext.wait(queues, timeout=10)

            self.assertNotEqual(ready_queues, [])  # cannot hang, delay is bound < 0.1 sec

            self.assertEqual(len(ready_queues), 1)
            q = ready_queues[0]
            self.assertEqual(q, queues[0])
            item = q.get()
            self.assertEqual(item, items[0])
            items.pop(0)
            queues.pop(0)
            events.pop(0)

        # make sure we shut down all service threads of WaitService class
        # before moving to the next test
        dragon.mpbridge.context._wait_service.__del__()
        # Reset global _wait_service so that init() is called again
        # and the garbage_collect thread initiated
        dragon.mpbridge.context._wait_service = None

    @unittest.skip(
        "This is broken, because dragon.mpbridge.queue.Faker() cannot be hashed. So Ben's solution actually "
        + "doesn't work either. The relevant MP unit tests seem to be:  WithProcessesTestPool.test_context,  "
        + "WithProcessesTestPool.test_empty_iterable"
    )
    def test_faker(self):

        num_workers = 4

        impbridge_ctx = mp.get_context()  # Queue() requires a context object

        queues = [impbridge_ctx.Queue() for _ in range(num_workers)]
        events = [dragon.native.event.Event() for _ in range(num_workers)]

        fake_writers = [q._reader for q in queues]
        fake_readers = [q._writer for q in queues]

        messages = [f"Test-Message-{i}" for i in range(num_workers)]

        procs = []

        for i in range(num_workers):
            p = dragon.mpbridge.DragonProcess(
                target=self._waitev_and_ready, args=(events[i], fake_writers[i], messages[i])
            )
            procs.append(p)

        for p in procs:
            p.start()

        # there should be no message yet
        self.assertFalse(dragon.mpbridge.DragonContext.wait(fake_readers, timeout=0.0))

        #  a new pipe should be ready every 0.1 sec
        while fake_readers:
            events[0].set()
            ready_readers = dragon.mpbridge.DragonContext.wait(fake_readers, timeout=10)

            self.assertNotEqual(ready_readers, [])  # cannot hang, delay is bound < 0.1 sec

            self.assertEqual(len(ready_readers), 1)
            r = ready_readers[0]
            self.assertEqual(r, fake_readers[0])
            msg = r.recv()  # this shouldn't block now
            self.assertEqual(msg, messages[0])
            messages.pop(0)
            fake_readers.pop(0)
            events.pop(0)

        # make sure we shut down all service threads of WaitService class
        # before moving to the next test
        dragon.mpbridge.context._wait_service.__del__()
        # Reset global _wait_service so that init() is called again
        # and the garbage_collect thread initiated
        dragon.mpbridge.context._wait_service = None

    # Make sure we can wait on a large list of different Dragon and multiprocessing objects with random return.
    # This also tests the capability of wait() to return several handles at once, as some delays will be
    # close to one another. We are using sockets to emulate a multiprocessing object that is certainly not Dragon.
    # TODO: Test this with a much large number of objects on a slow system. We should be able to handle 100s of
    # objects getting ready at very odd times.
    # JD: Add the sockets back in. This stopped working in the devcontainer and it is unclear to me why.
    def test_mixed(self):

        impbridge_ctx = mp.get_context()
        self.assertIsInstance(impbridge_ctx, dragon.mpbridge.DragonContext)

        num_pipes = num_queues = num_sentinels = num_fakers = num_channels = num_mp_objects = 4

        num_workers = (
            num_sentinels + num_pipes + num_queues
        )  # + num_mp_objects  # + num_channels + num_fakers

        items = [f"Test-Item-{i}" for i in range(num_workers)]  # wont need all of them

        # create handles for the processes to write to and if possible to read from
        pseudo_sentinels = [None for _ in range(num_sentinels)]  # don't have the real ones yet

        pipes = [impbridge_ctx.Pipe(duplex=False) for _ in range(num_pipes)]
        writers = [pipes[i][1] for i in range(num_pipes)]
        readers = [pipes[i][0] for i in range(num_pipes)]

        queues = [impbridge_ctx.Queue() for _ in range(num_queues)]

        faker_queues = [impbridge_ctx.Queue() for _ in range(num_fakers)]
        faker_writers = [q._reader for q in faker_queues]
        faker_readers = [q._writer for q in faker_queues]

        chan_pipes = [impbridge_ctx.Pipe(duplex=False) for _ in range(num_channels)]
        chan_writers = [chan_pipes[i][1] for i in range(num_channels)]
        chan_readers = [chan_pipes[i][0] for i in range(num_channels)]

        # socket_server = socket.create_server((socket_helper.HOST, 0))
        # addr = socket_server.getsockname()
        # mp_writers = [addr for _ in range(num_mp_objects)]
        # mp_readers = []

        proc_handles = pseudo_sentinels + writers + queues  # + mp_writers  # + chan_writers + faker_writers
        events = [
            dragon.native.event.Event() for _ in range(num_workers // 4)
        ]  # that'll make it pretty messy ...

        # start the processes
        procs = []
        random.seed(a=42)

        for i in range(num_workers):
            p = dragon.mpbridge.DragonProcess(
                target=self._waitev_and_ready, args=(events[i % len(events)], proc_handles[i], items[i])
            )
            procs.append(p)

        for p in procs:
            p.start()

        # create rest of the read handles to wait on
        sentinels = [p.sentinel for p in procs[:num_sentinels]]
        # cuids = [r.fileno() for r in chan_readers]
        # for _ in range(num_mp_objects):
        #     r, __ = socket_server.accept()
        #     mp_readers.append(r)
        # socket_server.close()

        # we should be able to wait on all of those simultaneously
        # this works for connections and queues, so I am assuming the test itself is OK.
        wait_handles = sentinels + readers + queues  # + mp_readers  # + cuids + faker_readers

        while wait_handles:

            if events:  # notify groups of processes
                ev = events.pop(0)
                ev.set()

            ready_handles = dragon.mpbridge.DragonContext.wait(wait_handles, timeout=None)

            self.assertNotEqual(ready_handles, [])

            for h in ready_handles:

                wait_handles.remove(h)

                if isinstance(h, dragon.mpbridge.process.PUID):
                    self.assertIn(h, sentinels)
                    sentinels.remove(h)
                elif isinstance(h, dragon.infrastructure.connection.Connection):
                    self.assertIn(h, readers)
                    readers.remove(h)
                    item = h.recv()
                    self.assertIn(item, items)
                    items.remove(item)
                elif isinstance(h, dragon.infrastructure.connection.CUID):
                    self.assertIn(h, cuids)
                    cuids.remove(h)
                    item = h.conn.recv()
                    self.assertIn(item, items)
                    items.remove(item)
                elif isinstance(h, dragon.mpbridge.queues.DragonQueue):
                    self.assertIn(h, queues)
                    queues.remove(h)
                    item = h.get()
                    self.assertIn(item, items)
                    items.remove(item)
                elif isinstance(h, dragon.mpbridge.queues.BaseImplQueue):
                    self.assertIn(h, queues)
                    queues.remove(h)
                    item = h.get()
                    self.assertIn(item, items)
                    items.remove(item)
                elif isinstance(h, dragon.mpbridge.queues.BaseImplJoinableQueue):
                    self.assertIn(h, queues)
                    queues.remove(h)
                    item = h.get()
                    self.assertIn(item, items)
                    items.remove(item)
                elif isinstance(h, dragon.mpbridge.queues.BaseImplSimpleQueue):
                    self.assertIn(h, queues)
                    queues.remove(h)
                    item = h.get()
                    self.assertIn(item, items)
                    items.remove(item)
                elif isinstance(h, dragon.mpbridge.queues.FakeConnectionHandle):
                    self.assertIn(h, faker_readers)
                    item = h.recv()
                    self.assertIn(item, items)
                    items.remove(item)
                elif isinstance(h, socket.socket):
                    self.assertIn(h, mp_readers)
                    mp_readers.remove(h)
                    item = h.recv(128).decode("utf-8")
                    self.assertIn(item, items)
                    items.remove(item)
                    h.close()
                else:
                    raise (TypeError(f"Wait returned a type that was not in the wait list: {type(h)}"))

        # make sure we shut down all service threads of WaitService class
        # before moving to the next test
        dragon.mpbridge.context._wait_service.__del__()
        # Reset global _wait_service so that init() is called again
        # and the garbage_collect thread initiated
        dragon.mpbridge.context._wait_service = None

    @staticmethod
    def _get_linux_num_threads(pid):
        with open(f"/proc/{pid}/status", "r") as fp:
            txt = fp.read()
        pstatus = dict(x.split(":\t") for x in txt.splitlines())

        return int(pstatus["Threads"])

    @classmethod
    def _check_num_threads(self, max_num_threads) -> bool:

        max_wait_time = 5  # sec
        pid = os.getpid()
        num_threads = self._get_linux_num_threads(pid)
        start = time.monotonic()
        while (time.monotonic() - start) < max_wait_time:
            if num_threads == max_num_threads:
                break
            time.sleep(0.1)
            num_threads = self._get_linux_num_threads(pid)

        return num_threads == max_num_threads

    # Ensure that the number of threads used by wait is constant, if we issue multiple waits on
    # the same connection object. Check that no service threads are spawned if we wait only on a
    # single object.
    # NOTE: This will break, when we implement a multi-poll on channels - that's the intention.
    # @unittest.skipIf(platform == "win32" or platform == "darwin", "Works on Linux only")
    def test_thread_hoarding(self):

        # Create something to wait on
        num_queues = 6

        impbridge_ctx = mp.get_context()

        queues = [impbridge_ctx.Queue() for _ in range(num_queues)]
        items = [f"Test-Item-{i}" for i in range(num_queues)]
        events = [dragon.native.event.Event() for _ in range(num_queues)]

        # start evaluating thread behaviour of the wait implementation
        wait_threads = list()
        spawn_counter = list()

        # initial state, should really be 1
        pid = os.getpid()
        spawn_counter.append(self._get_linux_num_threads(pid))

        # start waiting on two queues
        t = threading.Thread(target=dragon.mpbridge.DragonContext.wait, args=(queues[0:2],))
        t.start()
        wait_threads.append(t)
        self.assertTrue(
            self._check_num_threads(spawn_counter[0] + 4)
        )  # + 2 service threads (conns) + 1 wait thread + 1 garbage_collect thread
        spawn_counter.append(self._get_linux_num_threads(pid))

        # have another thread wait on the same 2 queues
        t = threading.Thread(target=dragon.mpbridge.DragonContext.wait, args=(queues[0:2],))
        t.start()
        wait_threads.append(t)
        self.assertTrue(self._check_num_threads(spawn_counter[1] + 1))  # + 0 service threads + 1 wait thread
        spawn_counter.append(self._get_linux_num_threads(pid))

        # have a third thread wait on another 3 queues
        t = threading.Thread(target=dragon.mpbridge.DragonContext.wait, args=(queues[2:5],))
        t.start()
        wait_threads.append(t)
        self.assertTrue(
            self._check_num_threads(spawn_counter[2] + 4)
        )  # + 3 service threads (conns) + 1 wait thread
        spawn_counter.append(self._get_linux_num_threads(pid))

        # have a fourth thread wait on a single queue
        t = threading.Thread(target=dragon.mpbridge.DragonContext.wait, args=([queues[5]],))
        t.start()
        wait_threads.append(t)
        self.assertTrue(
            self._check_num_threads(spawn_counter[3] + 2)
        )  # + 1 service threads (conns) + 1 wait thread
        spawn_counter.append(self._get_linux_num_threads(pid))

        # have the third wait thread and 4 service thread exit
        queues[4].put(items[4])
        wait_threads[2].join()
        self.assertTrue(
            self._check_num_threads(spawn_counter[4] - 4)
        )  # - 3 service threads (conns) - 1 wait thread
        spawn_counter.append(self._get_linux_num_threads(pid))

        # have the first 2 service threads + 2x1 wait threads exit
        queues[0].put(items[0])
        wait_threads[0].join()
        wait_threads[1].join()  # this one should wake up and exit too.
        self.assertTrue(
            self._check_num_threads(spawn_counter[5] - 4)
        )  # - 2 service threads (conns) - 2 wait threads
        spawn_counter.append(self._get_linux_num_threads(pid))

        # have the single wait thread exit last
        queues[5].put(items[5])
        wait_threads[3].join()
        self.assertTrue(
            self._check_num_threads(spawn_counter[6] - 2)
        )  # - 1 service threads (conns) - 1 wait thread

        # make sure we shut down all service threads of WaitService class
        # before moving to the next test
        dragon.mpbridge.context._wait_service.__del__()
        # Reset global _wait_service so that init() is called again
        # and the garbage_collect thread initiated
        dragon.mpbridge.context._wait_service = None

        # wait for the garbage_colloct thread to exit
        time.sleep(1)
        spawn_counter.append(self._get_linux_num_threads(pid))

        self.assertEqual(spawn_counter[7], spawn_counter[0])  # as we started

    # Test that we get the timeouts right.
    def test_timeout(self):

        # Create something to wait on
        num_queues = 6

        impbridge_ctx = mp.get_context()

        queues = [impbridge_ctx.Queue() for _ in range(num_queues)]
        items = [f"Test-Item-{i}" for i in range(num_queues)]

        # check that the timeout is reasonably correct, even when small
        timeout = 0.01
        t = threading.Thread(
            target=dragon.mpbridge.DragonContext.wait,
            args=(queues,),
            kwargs={
                "timeout": timeout,
            },
        )
        t.start()
        start = time.monotonic()
        t.join()
        end = time.monotonic()
        self.assertTrue(end - start < 1.5 * timeout)

        # check that waiting on the same objects a second time with a shorter timout works
        t0 = threading.Thread(
            target=dragon.mpbridge.DragonContext.wait,
            args=(queues,),
            kwargs={
                "timeout": 0.1,
            },
        )
        t1 = threading.Thread(
            target=dragon.mpbridge.DragonContext.wait,
            args=(queues,),
            kwargs={
                "timeout": 0.05,
            },
        )
        t0.start()
        t1.start()
        self.assertTrue(t1.is_alive())
        self.assertTrue(t0.is_alive())
        t1.join()
        self.assertTrue(t0.is_alive())
        self.assertFalse(t1.is_alive())
        t0.join()

        # check that waiting on the same objects a second time with a longer timeout works
        t0 = threading.Thread(
            target=dragon.mpbridge.DragonContext.wait,
            args=(queues,),
            kwargs={
                "timeout": 0.05,
            },
        )
        t1 = threading.Thread(
            target=dragon.mpbridge.DragonContext.wait,
            args=(queues,),
            kwargs={
                "timeout": 0.1,
            },
        )
        t0.start()
        t1.start()
        self.assertTrue(t1.is_alive())
        self.assertTrue(t0.is_alive())
        t0.join()
        self.assertTrue(t1.is_alive())
        self.assertFalse(t0.is_alive())
        t1.join()

        # the following is related to PE-44688
        # calling close() on the queues is useful when combined with
        # #define CHECK_POINTER on bcast.c
        # for q in queues:
        #     q.close()

        # make sure we shut down all service threads of WaitService class
        # before moving to the next test
        dragon.mpbridge.context._wait_service.__del__()
        # Reset global _wait_service so that init() is called again
        # and the garbage_collect thread initiated
        dragon.mpbridge.context._wait_service = None

if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_mpbridge_context_wait.py
DragonHPC/dragon
"""
Basic test of multiprocessing message passing through multiprocessing.Pipe.

Unit test adapted from original 'fun_srt.py' code.
"""

import dragon.mpbridge
import dragon.infrastructure.connection
import dragon.infrastructure.messages as dmsg

import dragon.dlogging.util as dlog
from dragon.utils import B64

import multiprocessing
import time
import logging
import mp_bench.util.test_util as test_util

from collections import namedtuple
import os
import unittest
from unittest.mock import patch

POOL_SIZE = 4194304
DEFAULT_UID = 123456
POOL_NAME = f"log_test_{os.getpid()}"


def setUpModule():
    multiprocessing.set_start_method("dragon", force=True)


# could this be a deadly embrace sort of thing
# given that open() happens and something
# is being asked for out of order in full duplex world?
# I think that could have to do with this...


@patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "INFO"})
def external_single_echo(conn, resp_size, serialized_dlogging, idx):
    dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST, logger_sdesc=B64(serialized_dlogging))
    logger = logging.getLogger(dlog.DragonLoggingServices.TEST)

    logger.info(f"{idx}: external_single_echo")
    conn.recv()
    conn.send(bytes(resp_size))


class SendReceiveTest(unittest.TestCase):
    def setUp(self):
        self.assertEqual(multiprocessing.get_start_method(), "dragon")

        # parser = test_util.add_default_args('Send and receive a message of varying size using a Pipe')
        # parser = test_util.add_connections_args(parser)
        # parser.add_argument('--response_size', type=int, default=4, help='size of message pass back')
        # parser.add_argument('--dragon', action='store_true', help='run a la single node dragon')
        # self.args = test_util.process_args(parser)
        FakeArgs = namedtuple(
            "FakeArgs",
            ["burn_iterations", "json", "message_size", "num_workers", "response_size", "num_iterations"],
        )
        self.response_size = 4
        self.args = FakeArgs(10, True, 100, 2, self.response_size, 3)
        self.results = None

        # check avail I guess.
        self.Pipe = dragon.infrastructure.connection.Pipe
        self.Process = dragon.mpbridge.DragonProcess

    @unittest.skipUnless("DRAGON_MODE" in os.environ, "Must be launched as a dragon process")
    def test_run(self):

        times = test_util.iterations(self.do_send_receive, self.args)
        self.results = test_util.Results(self.args, times, "send_receive_test")
        self.assertEqual(self.response_size, self.results.response_size)

    def do_send_receive(self):

        num_workers = self.args.num_workers

        message = bytes(self.args.message_size)

        pipeconns = [self.Pipe() for _ in range(num_workers)]

        dragon_logger = dlog.setup_dragon_logging(0)
        serialized_dlogging = dragon_logger.serialize()

        processes = []
        for idx in range(num_workers):

            processes.append(
                self.Process(
                    target=external_single_echo,
                    args=(pipeconns[idx][1], self.args.response_size, serialized_dlogging, idx),
                )
            )

        start_time = time.time_ns()

        for idx, handle in enumerate(processes):
            handle.start()

        for i in range(num_workers):
            pipeconns[i][0].send(message)

        for i in range(num_workers):
            pipeconns[i][0].recv()

        for idx, handle in enumerate(processes):
            handle.join()

        num_logs = dragon_logger.num_logs()
        self.assertEqual(
            num_logs, num_workers, f"expected {num_workers} log entries, but there's only {num_logs}"
        )
        found_worker_log_msg = [False for _ in range(num_workers)]
        for i in range(num_workers):
            msg = dmsg.parse(dragon_logger.get(i))
            worker = int(msg.msg.split(":")[0])
            self.assertFalse(
                found_worker_log_msg[worker - 1], f"duplicate log message received for node {worker}"
            )
            found_worker_log_msg[worker - 1] = True
        self.assertTrue(
            all(found_worker_log_msg), "One or more log messages not found in DragonLogging chanel"
        )
        dragon_logger.destroy()

        return time.time_ns() - start_time


if __name__ == "__main__":
    dlog.setup_logging(basename="head", the_level=logging.DEBUG)

    unittest.main()

test/mpbridge/test_pipe_with_process.py
DragonHPC/dragon
import unittest
import pickle
import time

import dragon
import multiprocessing as mp


def setUpModule():
    mp.set_start_method("dragon", force=True)


class TestMPBridgeProcess(unittest.TestCase):
    def setUp(self):
        self.assertEqual(mp.get_start_method(), "dragon")

    def test_active_children(self):

        self.assertTrue(len(mp.active_children()) == 0)

        processes = []

        for _ in range(0, 8):
            p = mp.Process(target=time.sleep, args=(100,))
            processes.append(p)

        pids = []

        for p in processes:
            p.start()
            active_pids = [c.pid for c in mp.active_children()]
            self.assertTrue(p.pid in active_pids)
            pids.append(p.pid)

        self.assertTrue({c.pid for c in mp.active_children()} == set(pids))

        for p in processes:
            p.kill()
            p.join()

        self.assertTrue(len(mp.active_children()) == 0)

    @staticmethod
    def _send_own_process_object(writer):

        p = mp.current_process()

        p.authkey = b""  # cannot pickle the authkey
        p._config["authkey"] = b""
        writer.send(p)

    def test_current_process(self):

        reader, writer = mp.Pipe()

        my_p = mp.Process(target=self._send_own_process_object, args=(writer,))
        my_p.start()

        my_p.authkey = b""  # make authkey equal
        my_p._config["authkey"] = b""

        their_p = reader.recv()

        # Note that the two process objects must be the same only in some
        # attributes. Other are different - this is true also for
        # standard multiprocessing

        self.assertTrue(their_p._closed == my_p._closed)
        self.assertTrue(their_p._identity == my_p._identity)
        self.assertTrue(their_p._name == my_p._name)
        self.assertTrue(their_p._parent_pid == my_p._parent_pid)
        self.assertTrue(their_p._parent_name == my_p._parent_name)

        self.assertTrue(their_p._p_p_uid == my_p._p_p_uid)

    @staticmethod
    def _send_parent_process_object(writer):

        pp = mp.parent_process()

        writer.send(pp)

    def test_parent_process(self):

        reader, writer = mp.Pipe()

        my_p = mp.current_process()

        cp = mp.Process(target=self._send_parent_process_object, args=(writer,))
        cp.start()

        their_p = reader.recv()

        self.assertTrue(their_p._closed == my_p._closed)
        self.assertTrue(their_p._identity == my_p._identity)
        self.assertTrue(their_p._name == my_p._name)
        self.assertTrue(their_p._popen == my_p._popen)
        self.assertTrue(their_p._parent_pid == my_p._parent_pid)


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_process.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import random
import dragon
import multiprocessing as mp
import sys
import threading
import numpy as np
import queue as pyqueue


def setUpModule():
    mp.set_start_method("dragon", force=True)


class TestQueue(unittest.TestCase):
    def setUp(self):
        self.assertEqual(mp.get_start_method(), "dragon")

    def test_queue(self):
        q = dragon.mpbridge.queues.DragonQueue()
        y = [1, 2, 3]
        q.put(y)
        x = q.get()
        self.assertEqual(x, y, "The Queue implementation did not pass a list correctly through it.")

        yy = ["a", "b", "c"]
        q.put(yy)
        q.put(y)
        x = q.get()
        self.assertEqual(x, yy, "The Queue implementation did not pass a list correctly through it.")
        x = q.get()
        self.assertEqual(x, y, "The Queue implementation did not pass a list correctly through it.")

        q.close()

    @classmethod
    def _writer(self, id, q, nitems, msg_list, qver):
        for i in range(id * nitems, id * nitems + nitems):
            q.put((i, msg_list[i]))
            qver.put((i, msg_list[i]))
        q.close()
        qver.close()

    @classmethod
    def _reader(self, id, q, nitems, msg_list):
        for _ in range(nitems):
            item = q.get()
            _id, payload = item
            assert payload == msg_list[_id]
        q.close()

    def test_queue_multi_readers_multi_writers(self):
        nitems = 20  # multiple of num_readers and num_writers
        num_readers = 4
        num_writers = 2
        msg_list = {}  # key is the msg identifier, value is the payload
        for i in range(nitems):
            msg_list[i] = "Hello msg with id " + str(i)

        q = dragon.mpbridge.queues.DragonQueue()
        q_ver = (
            dragon.mpbridge.queues.DragonQueue()
        )  # used to verify that all messages are gotten from the queue

        writers = []
        for i in range(num_writers):
            p = mp.Process(target=self._writer, args=(i, q, int(nitems / num_writers), msg_list, q_ver))
            p.start()
            writers.append(p)

        readers = []
        for i in range(num_readers):
            p = mp.Process(target=self._reader, args=(i, q, int(nitems / num_readers), msg_list))
            readers.append(p)
            p.start()

        for p in writers:
            p.join()

        for p in readers:
            p.join()

        while not q_ver.empty():
            _id, payload = q_ver.get()
            if _id in msg_list:
                del msg_list[_id]

        self.assertEqual(len(msg_list), 0)
        self.assertTrue(q_ver.empty())

        self.assertTrue(q.empty())

        q.close()
        q_ver.close()

    def pass_an_obj(self, obj):
        q = dragon.mpbridge.queues.DragonQueue()
        q.put(obj)
        rec_obj = q.get()

        if isinstance(obj, np.ndarray):
            np.testing.assert_array_equal(obj, rec_obj)
        else:
            self.assertEqual(obj, rec_obj)

        q.close()

    def test_objs_queue(self):
        """Test different python objects and different msg sizes"""
        self.pass_an_obj("hyena")
        self.pass_an_obj({"hyena": "hyena"})
        self.pass_an_obj(bytes(100000))
        self.pass_an_obj(bytes(200000))
        self.pass_an_obj(np.random.rand(1000000))

        for lg_size in range(25):
            self.pass_an_obj(bytearray(2**lg_size))

        for lg_size in range(15):
            self.pass_an_obj(random.randbytes(2**lg_size))

        for lg_size in range(22):
            my_msg = []
            my_size = sys.getsizeof(my_msg)
            while my_size < 2**lg_size:
                my_msg.append("a")
                my_size = sys.getsizeof(my_msg)
            self.pass_an_obj(my_msg)

    def test_queue_timeout(self):
        q = dragon.mpbridge.queues.DragonQueue()

        # try to send an object which is large enough
        # such that pickle will call write() multiple
        # times and we want to make sure that the timeout
        # is updated/decremented at each call
        obj = bytes(10000000)

        # set a small value for the timeout for this obj
        # so that we make sure it raises ChannelTimeout
        self.assertRaises(pyqueue.Full, q.put, obj, True, 0.0001)

        q.put(obj, timeout=5)

        rec_obj = q.get()
        if isinstance(obj, np.ndarray):
            np.testing.assert_array_equal(obj, rec_obj)
        else:
            self.assertEqual(obj, rec_obj)

        q.close()


if __name__ == "__main__":
    unittest.main(verbosity=2)

test/mpbridge/test_queue.py
DragonHPC/dragon
#!/usr/bin/env python3

import dragon
import unittest
from dragon.native.lock import (
    Lock,
)
from dragon.mpbridge.sharedctypes import (
    Value,
)
import ctypes


class TestValue(unittest.TestCase):
    def test_requirement_1_1(
        self,
    ):
        """tests that the attributes get_obj and get_lock are not available"""
        self.assertFalse(
            hasattr(
                Value(
                    "c",
                    lock=False,
                ),
                "get_lock",
            )
        )
        self.assertFalse(
            hasattr(
                Value(
                    "c",
                    lock=False,
                ),
                "get_obj",
            )
        )

    def test_requirement_1_2(
        self,
    ):
        """tests that the rlock is set properly"""
        val = Value(
            "c",
            lock=True,
        )
        self.assertIsInstance(
            val.get_lock(),
            type(Lock(recursive=True)),
        )


if __name__ == "__main__":
    unittest.main()

test/mpbridge/test_value.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance test for the
`dragon.native.Array` object.  The test scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.

The test is run with `dragon test_array.py -f -v`
"""
import unittest
import os
import sys
import numpy as np

import dragon
import multiprocessing as mp
import random


def process_function(event, array, queue, max_val, idx,):

    p = mp.current_process()
    mypuid = p.pid

    # each child process waits on parent process
    event.wait()

    array[idx] = 1

    #while the value is less than max_value
    while array[idx] < max_val:
      with array.get_lock():
            array[idx] += 1
            array_list = array[:]
            array_list.append(random.randint(-sys.maxsize - 1, sys.maxsize))
            array[:] = array_list
            queue.put((mypuid, array[idx],))


class TestArrayMultiNode(unittest.TestCase):
    def test_lock_fairness(self,):
        """Tests that array can be incremented by multiple processes in a
        multinode environment.

        1. The Test uses `nproc = mp.cpu_count()` processes.
        2. Each child process gets handed the same `Array("i", range(nproc)), lock=True)`, a
           Queue, an Event and the max_val.
        3. Each child process first waits in the Event.
        4. Once all processes have been started by the parent, the parent sets
           the event, so that all processes start working at the same time. This
           means the runtime sees a very sudden increase in activity, which is a
           test in itself.
        5. Each child process tries to acquire the lock and add elements to the array
           associated with the `Array`. This (step 6) is repeated until the element at the
           index of the Array is equal to max_val.
        6. When the lock was acquired, also adds a random element to the array while the element at the selected
            index of the array is being incremented to max_val.
        7. The parent process counts how often it received a array from any of
           the puids. It prints the minimum, maximum, mean and median of the number
           of times.
        8. Finally the parent joins on all processes and check that their
           exit_code is 0, i.e. that they exited cleanly without an exception. The size of the array is checked,
           and the the first nproc values in the array are checked.
        """

        procs = []
        puid_dict = {}

        # number of processes for the test
        nproc = mp.cpu_count() // 16
        max_val = 2

        # each child process is handed the same array, queue, and event
        array = mp.Array("i", range(128+nproc), lock=True,)
        queue = mp.Queue()
        event = mp.Event()

        for idx in range(nproc):
            p = mp.Process(target=process_function,args=(event, array, queue, max_val, idx,),)
            p.start()
            procs.append(p)
            puid_dict[p.pid] = 0

        # parent sets event for children processes to work at same time
        event.set()

        cnt = 0

        for idx in range(0,max_val,):
            puid, val, = (queue.get())
            # check that the array from child process is valid
            cnt += 1
            # increment each time child puid returns array
            puid_dict[puid] += 1

        numpy_puids = np.array(list(puid_dict.values()))

        # print minimum, maximum, mean, and median from puids_dict
        (puids_max,puids_min,puids_mean,puids_median,) = (np.max(numpy_puids),np.min(numpy_puids),np.mean(numpy_puids),np.median(numpy_puids),)

        print(f"Lock acquisition, {max_val/nproc} tries / proc, {nproc} processes:")
        print(f"    Maximum: {puids_max}, Minimum: {puids_min}, Mean: {puids_mean}, Median: {puids_median}")

        # check that the child process exited cleanly in parent process
        for p in procs:
            p.join()
            self.assertEqual(p.exitcode,0,)

        #checks the length of the array
        self.assertLessEqual(len(array), (128 + (max_val+1) * nproc))

        #checks that all the integers in the array fall in the range of ints possible
        for index in range(0, nproc):
            self.assertTrue((-sys.maxsize - 1) <= array[index] <= (sys.maxsize))



if (__name__== "__main__"):
    mp.set_start_method("dragon",force=True,)
    unittest.main()

test/multi-node/test_array.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
`dragon.native.Barrier` object. The tests scale with the total number of CPUs
reported by the allocation, i.e. they become tougher on larger allocations.

The tests are run with `dragon test_barrier.py -f -v`
"""

import os
import unittest
import time
import pickle
import sys

import dragon

from dragon.native.process import current
from dragon.native.queue import Queue
from dragon.native.barrier import Barrier, BrokenBarrierError
from dragon.native.machine import cpu_count

from dragon.globalservices.process import (
    join,
    create,
    multi_join,
)
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.utils import B64


def wait_and_timeout_gracefully(args):
    bytes = B64.str_to_bytes(args)
    barrier = pickle.loads(bytes)[0]
    try:
        start = time.monotonic()
        barrier.wait(timeout=2)
    except BrokenBarrierError:
        pass
    stop = time.monotonic()
    if abs(stop - start - 2) > 0.2:
        raise AssertionError


def wait_forever(args):
    try:
        bytes = B64.str_to_bytes(args)
        barrier = pickle.loads(bytes)[0]
        barrier.wait(timeout=None)
    except Exception as ex:
        print(f"There was an exception {ex}")


def wait_forever_and_timeout_gracefully(args):
    bytes = B64.str_to_bytes(args)
    barrier = pickle.loads(bytes)[0]
    try:
        barrier.wait(timeout=None)
    except BrokenBarrierError:
        pass


def reset_barrier(args):
    bytes = B64.str_to_bytes(args)
    arg_list = pickle.loads(bytes)
    barrier = arg_list[0]
    num_parties = arg_list[1]
    while barrier.n_waiting < num_parties:
        time.sleep(0.1)
    barrier.reset()


def abort_barrier(args):
    bytes = B64.str_to_bytes(args)
    arg_list = pickle.loads(bytes)
    barrier = arg_list[0]
    num_parties = arg_list[1]
    while barrier.n_waiting < num_parties:
        time.sleep(0.1)
    barrier.abort()


def wait_and_put_result(args):
    bytes = B64.str_to_bytes(args)
    barrier, queue = pickle.loads(bytes)
    queue.put(barrier.wait())


def wait_and_put_when_broken(args):
    bytes = B64.str_to_bytes(args)
    barrier, queue = pickle.loads(bytes)
    try:
        barrier.wait(timeout=None)
    except BrokenBarrierError:
        queue.put(True)


class TestBarrierMultiNode(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.ncpu = max(2, int(cpu_count() / 4))  # control the number of CPU for all tests
        # print(f'The tests were run with {cls.ncpu} parties for the Barriers.')

    def test_requirement_2_2(self):
        """
        If wait is called by less than parties processes within timeout seconds
        from the first process calling the method, all processes who have called
        it and will call it until reset() was called, have to raise a
        BrokenBarrierError immediately.  Note also Req 2.4.
        """
        barrier = Barrier(parties=2)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        arg = B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_and_timeout_gracefully", arg],
            None,
            options=options,
        )
        join(process_1.p_uid)

        self.assertRaises(BrokenBarrierError, lambda: barrier.wait(timeout=0))

    def test_requirement_2_3_a(self):
        """
        If wait is called with timeout=None, the calling process has to block
        until the method has been called by parties processes, or abort() or
        reset() has been called by any other process.
        """

        barrier = Barrier(parties=self.ncpu)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)

        arg = B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )

        puids = []
        for _ in range(self.ncpu):
            descr = create(
                cmd,
                wdir,
                [__file__, "wait_forever", arg],
                None,
                options=options,
            )
            puids.append(descr.p_uid)

        ready = multi_join(puids, join_all=True)

        for puid, ecode in ready[0]:
            self.assertTrue(ecode == 0)
            self.assertTrue(puid in puids)

    def test_requirement_2_3_b(self):
        """
        If wait is called with timeout=None, the calling process has to block
        until the method has been called by parties processes, or abort() or
        reset() has been called by any other process.
        """

        barrier = Barrier(parties=self.ncpu)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        arg = B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )

        puids = []
        for _ in range(self.ncpu - 1):
            descr = create(
                cmd,
                wdir,
                [__file__, "wait_forever_and_timeout_gracefully", arg],
                None,
                options=options,
            )
            puids.append(descr.p_uid)

        arg2 = B64.bytes_to_str(
            pickle.dumps(
                [barrier, self.ncpu - 1],
                protocol=5,
            )
        )

        descr = create(
            cmd,
            wdir,
            [__file__, "reset_barrier", arg2],
            None,
            options=options,
        )
        puids.append(descr.p_uid)

        ready = multi_join(puids, join_all=True)

        for puid, ecode in ready[0]:
            self.assertTrue(ecode == 0)
            self.assertTrue(puid in puids)

    def test_requirement_2_3_c(self):
        """
        If wait is called with timeout=None, the calling process has to block
        until the method has been called by parties processes, or abort() or
        reset() has been called by any other process.
        """

        barrier = Barrier(parties=self.ncpu)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        arg = B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )
        puids = []
        for _ in range(self.ncpu - 1):
            descr = create(
                cmd,
                wdir,
                [__file__, "wait_forever_and_timeout_gracefully", arg],
                None,
                options=options,
            )
            puids.append(descr.p_uid)

        arg2 = B64.bytes_to_str(
            pickle.dumps(
                [barrier, self.ncpu - 1],
                protocol=5,
            )
        )

        descr = create(
            cmd,
            wdir,
            [__file__, "abort_barrier", arg2],
            None,
            options=options,
        )
        puids.append(descr.p_uid)

        ready = multi_join(puids, join_all=True)

        for puid, ecode in ready[0]:
            self.assertTrue(ecode == 0)
            self.assertTrue(puid in puids)

    def test_requirement_2_4(self):
        """
        If wait() is called after abort() has been called by any process, raise
        a BrokenBarrierError until reset() is called.
        """

        barrier = Barrier(parties=self.ncpu)

        barrier.abort()
        self.assertRaises(BrokenBarrierError, lambda: barrier.wait(timeout=None))

        barrier.reset()

        # check that we are fully functional again
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        arg = B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )

        puids = []
        for _ in range(self.ncpu):
            descr = create(
                cmd,
                wdir,
                [__file__, "wait_forever", arg],
                None,
                options=options,
            )
            puids.append(descr.p_uid)

        ready = multi_join(puids, join_all=True)

        for puid, ecode in ready[0]:
            self.assertTrue(ecode == 0)
            self.assertTrue(puid in puids)

    def test_requirement_2_6(self):
        """
        If wait has been called parties times within timeout seconds,
        it has to return a unique integer on every process in range 0<=i<parties.
        """

        barrier = Barrier(action=None, parties=self.ncpu, timeout=None)

        queue = Queue(joinable=False)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        arg = B64.bytes_to_str(pickle.dumps([barrier, queue], protocol=5))

        puids = []
        for _ in range(self.ncpu):
            process_1 = create(
                cmd,
                wdir,
                [__file__, "wait_and_put_result", arg],
                None,
                options=options,
            )
            puids.append(process_1.p_uid)

        queue_vals = []
        for _ in range(self.ncpu):
            queue_vals.append(queue.get())

        multi_join(puids, join_all=True)

        queue_vals = sorted(queue_vals)
        self.assertListEqual(list(range(self.ncpu)), queue_vals)

    def test_requirement_3_1(self):
        """
        If the reset has been called while processes are waiting in wait(), they
        have to raise a BrokenBarrierError immediately.
        """
        barrier = Barrier(parties=self.ncpu)

        queue = Queue(joinable=False)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        arg = B64.bytes_to_str(pickle.dumps([barrier, queue], protocol=5))

        puids = []
        for _ in range(self.ncpu - 1):
            process_1 = create(
                cmd,
                wdir,
                [__file__, "wait_and_put_when_broken", arg],
                None,
                options=options,
            )
            puids.append(process_1.p_uid)

        arg2 = B64.bytes_to_str(
            pickle.dumps(
                [barrier, self.ncpu - 1],
                protocol=5,
            )
        )

        descr = create(
            cmd,
            wdir,
            [__file__, "reset_barrier", arg2],
            None,
            options=options,
        )

        for _ in range(self.ncpu - 1):  # hangs here ...
            self.assertTrue(queue.get())

        ready = multi_join(puids, join_all=True)

        for puid, ecode in ready[0]:
            self.assertTrue(ecode == 0)
            self.assertTrue(puid in puids)


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "wait_forever":
        wait_forever(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "wait_and_put_result":
        wait_and_put_result(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "wait_and_timeout_gracefully":
        wait_and_timeout_gracefully(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "wait_and_put_when_broken":
        wait_and_put_when_broken(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "reset_barrier":
        reset_barrier(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "abort_barrier":
        abort_barrier(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "wait_forever_and_timeout_gracefully":
        wait_forever_and_timeout_gracefully(sys.argv[2])
    else:
        unittest.main()

test/multi-node/test_barrier.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
`multiprocessing.Pipe` object.  The test scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.

The test is run with `dragon test_connection.py -f -v`
"""

import numpy as np
import unittest

import dragon
import multiprocessing as mp

import hashlib  # not using standard hash, as it is salted :-(


def ring_send_recv(id, bwd_recv, fwd_send, q, ev, sizes):
    """Send random bytes through a connection. Then recv the same."""

    data_all = np.random.bytes(max(sizes) * 1024**2)

    q.put(True)

    ev.wait()

    for size in sizes:

        idx = int(size * 1024**2)
        data = data_all[0:idx]

        m = hashlib.sha256()
        m.update(data)
        checksum = m.hexdigest()

        fwd_send.send((data, checksum))

        data, checksum = bwd_recv.recv()

        m = hashlib.sha256()
        m.update(data)
        digest = m.hexdigest()

        if digest != checksum:
            raise AssertionError("Data Corruption detected")

        q.put(len(data))

    fwd_send.close()
    bwd_recv.close()


def ring_recv_send(id, bwd_recv, fwd_send, q, ev, sizes):
    """Recv random bytes through a connection. Then send the same."""

    data_all = np.random.bytes(max(sizes) * 1024**2)

    q.put(True)

    ev.wait()

    for size in sizes:
        data, checksum = bwd_recv.recv()

        m = hashlib.sha256()
        m.update(data)
        digest = m.hexdigest()

        if digest != checksum:
            raise AssertionError("Data Corruption detected")

        q.put(len(data))

        idx = int(size * 1024**2)
        data = data_all[0:idx]

        m = hashlib.sha256()
        m.update(data)
        checksum = m.hexdigest()

        fwd_send.send((data, checksum))

    fwd_send.close()
    bwd_recv.close()


# @unittest.skip(f"Fails for larger CPU counts")
class TestConnectionMultiNode(unittest.TestCase):
    def test_ring_multi_node(self):
        """This test creates 2 processes per cpu. They send increasingly large
        messages in a ring to each other."""

        sizes = [0.1, 1, 10, 20]  # , 100, 500]  # MB !

        num_processes = max(2, mp.cpu_count() // 8)

        qs = [mp.Queue() for i in range(num_processes)]
        ev = mp.Event()

        first_recv, last_send = mp.Pipe(duplex=False)

        this_recv = first_recv

        processes = []
        for i in range(num_processes):

            if i == num_processes - 1:
                this_send = last_send
            else:
                next_recv, this_send = mp.Pipe(duplex=False)

            if i % 2 == 0:
                target = ring_send_recv
            else:
                target = ring_recv_send

            p = mp.Process(target=target, args=(i, this_recv, this_send, qs[i], ev, sizes))
            p.start()
            processes.append(p)

            this_recv = next_recv

        for q in qs:
            self.assertTrue(q.get())

        ev.set()  # start communication all at once

        for size in sizes:
            for q in qs:
                recv_size = q.get()
                self.assertTrue(recv_size == int(size * 1024**2))

        for p in processes:
            p.join()

        for q in qs:
            q.close()


if __name__ == "__main__":
    mp.set_start_method("dragon")
    unittest.main()

test/multi-node/test_connection.py
DragonHPC/dragon
#!/usr/bin/env python3

import os
import sys
import string
import random
import dragon
import unittest
import multiprocessing as mp

from dragon.data.distdictionary.dragon_dict import DragonDict
from dragon.globalservices.node import get_list
from dragon.infrastructure import parameters

def do_client_ops(ddict, key, value):
    """Each client will do the operation of adding the sample key,value pair
    into the dictionary and again delete the same pair.

    :param ddict: dragon distributed dictionary
    :type ddict: dragon dictionary object
    :param key: Information about the SET operation to the manager
    :type key: Any hashable structure
    :param value: Value to be stored inside the manager
    :type value: Any structure that can be serialized
    """
    ddict[key] = value
    del ddict[key]
    # Close the client
    ddict.close()

def generate_keys():
    keys = list()
    # Generate the list of 100 keys to add to the dictionary
    dict_size = 100
    letters = string.ascii_letters

    for _ in range(dict_size):
        # each key is 20 characters long and characters can be repeated
        key = ''.join(random.choice(letters) for i in range(20))
        keys.append(key)
    return keys

def do_set_ops(_keys, ddict, value_size):
    """Each client will generate the value of given size, and perform
    the SET operation to all the keys of the dictionary

    :param _keys: list of all keys to be added to the dictionary
    :type _keys: list
    :param ddict: dragon distributed dictionary
    :type ddict: dragon dictionary object
    :param value_size: size of the value in bytes added to each dictionary key
    :type value_size: int
    """
    num_keys = len(_keys)
    letters = string.ascii_letters
    value = ''.join(random.choice(letters) for i in range(value_size))

    # Add each key to the dictionary
    for key in _keys:
        ddict[key] = value
    # Close the client
    ddict.close()

def do_del_ops(_keys, ddict):
    """Each client will delete the given list of keys in the dictionary
    
    :param _keys: list of all keys to be deleted from the dictionary
    :type _keys: list
    :param ddict: dragon distributed dictionary
    :type ddict: dragon dictionary object
    """
    for key in _keys:
        del ddict[key]
    # Close the client
    ddict.close()

def do_get_ops(_keys, ddict):
    """Each client will retrieve the values for the given list of keys in the dictionary
    
    :param _keys: list of all keys to be fetched from the dictionary
    :type _keys: list
    :param ddict: dragon distributed dictionary
    :type ddict: dragon dictionary object
    """
    for key in _keys:
        value = ddict[key]
    # Close the client
    ddict.close()


class TestDragonDictMultiNode(unittest.TestCase):
    @classmethod
    def setUpClass(self) -> None:
        # Create a dragon dictionary on a single node with multiple manager processes
        self.managers_per_node = 1 # 1 Managers per node
        self.num_nodes = len(get_list()) # Collect the total number of nodes
        self.total_mem_size = 2*(1024*1024*1024) # 2 GB total size
        self.DD = DragonDict(self.managers_per_node, self.num_nodes, self.total_mem_size)

    @classmethod
    def tearDownClass(self) -> None:
        self.DD.stop()

    def test_dict_params(self):
        # Verify the dictionary params
        total_nodes = len(get_list())
        dict_nodes = self.DD._dist_dict.nodes
        managers_per_node = self.DD._dist_dict.managers_per_node
        dict_managers = dict_nodes * managers_per_node
        self.assertEqual(dict_nodes, total_nodes, "Number of dictionary nodes are incorrect")
        self.assertEqual(managers_per_node, 1, "Number of managers per node are incorrect")
        self.assertEqual(dict_managers, total_nodes, "Number of dictionary managers are incorrect")

    def test_set_and_get_item(self):
        key = "Who are you Dragon?"
        value = "A Fire Breathing Monster!"
        self.DD[key] = value
        # Compare if the fetched values are correct
        self.assertEqual(self.DD[key], value, "Retrieved value from the dict is incorrect")

    def test_setup_and_close_client(self):
        # Create a client process and pass the dictionary
        client_proc = mp.Process(target=do_client_ops, args=(self.DD, "Dragon", "Dictionary"))
        client_proc.start()
        client_proc.join()

    def test_set_ops(self):
        keys = generate_keys()
        # Create client procs that does dictionary operations
        num_clients = 8
        value_size = 64
        procs = []
        for i in range(num_clients):
            client_proc = mp.Process(target=do_set_ops, args=(keys, self.DD, value_size))
            client_proc.start()
            procs.append(client_proc)

        for i in range(len(procs)):
            procs[i].join()
            procs[i].kill()

    def test_dictionary_length(self):
        num_kv_pairs = 100
        # Assign key value pairs to the dictionary
        for i in range(0, num_kv_pairs):
            key = "Hello" + "." + str(i)
            value = "Dragon" + "." + str(i)
            self.DD[key] = value

        # Calculate the length of the dictionary
        dict_length = len(self.DD)
        # Verify that length of dictionary is equal to number of kv pairs
        self.assertEqual(dict_length, num_kv_pairs, "Calculated length of the dictionary is incorrect")

    def test_set_and_del_item(self):
        key = "Hello"
        value = "Dictionary"
        self.DD[key] = value
        # Delete operation should be successful
        del self.DD[key]

    def test_del_item_with_no_key(self):
        # Delete operation should be successful, without key present also
        # Log the information that the key is not present.
        letters = string.ascii_letters
        key = ''.join(random.choice(letters) for i in range(10))
        del self.DD[key]

    def test_del_ops(self):
        keys = list()
        num_kv_pairs = 100
        # Collect the set of keys for the dictionary
        for i in range(0, num_kv_pairs):
            key = "Hello" + "." + str(i)
            keys.append(key)

        # Two client processes delete keys in parallel
        # This might conduct extra deletes with keys deleted by other proc before
        num_clients = 2
        procs = []
        for i in range(num_clients):
            client_proc = mp.Process(target=do_del_ops, args=(keys, self.DD))
            client_proc.start()
            procs.append(client_proc)

        for i in range(len(procs)):
            procs[i].join()
            procs[i].kill()

    def test_keys(self):
        # Collect the keys from the dictionary
        keys = self.DD.keys()
        num_keys = len(keys)
        dict_length = len(self.DD)
        # Verify that length of dictionary is equal to number of kv pairs
        self.assertEqual(dict_length, num_keys, "Calculated length of the dictionary is incorrect")

    def test_get_ops(self):
        keys = self.DD.keys()
        # Create client procs that does dictionary operations
        num_clients = 8
        procs = []
        for i in range(num_clients):
            client_proc = mp.Process(target=do_get_ops, args=(keys, self.DD))
            client_proc.start()
            procs.append(client_proc)

        for i in range(len(procs)):
            procs[i].join()
            procs[i].kill()

    def test_existing_key(self):
        key = "Testing"
        value = "Key"
        self.DD[key] = value
        key_found = False
        # Verify if the key is present in the dictionary
        if key in self.DD:
            key_found = True
        self.assertEqual(key_found, True, "Key has not found in the dictionary")

    def test_not_existing_key(self):
        key = "Test"
        key_found = False
        # Verify if the key is present in the dictionary
        if key in self.DD:
            key_found = True
        self.assertEqual(key_found, False, "Key found in the dictionary")

    def test_pop_item(self):
        key = "Pop"
        value = "Item"
        self.DD[key] = value
        # Collect the value by popping the key
        pop_value = self.DD.pop(key)
        # Verify that the pop operation is correct
        self.assertEqual(value, pop_value, "Pop value is not returned correctly")

    def test_existing_value(self):
        key = "Testing"
        value = "Value"
        self.DD[key] = value
        value_found = False
        # Verify if the value is existing in the dictionary
        if value in self.DD.values():
            value_found = True
        self.assertEqual(value_found, True, "Value has not found in the dictionary")

    def test_not_existing_value(self):
        value = "Dummy Value"
        value_found = False
        # Verify if the value is existing in the dictionary
        if value in self.DD.values():
            value_found = True
        self.assertEqual(value_found, False, "Value found in the dictionary")

    def test_items(self):
        items = self.DD.items()
        for (key, value) in items:
            ret_value = self.DD[key]
            # Verify the value retrieved from the items with the specific key
            self.assertEqual(value, ret_value, "Value collected from the items is incorrect")

    def test_rename(self):
        key = "Hello"
        value = "Dictionary"
        self.DD[key] = value
        # Rename the above key
        new_key = "Renamed"
        self.DD.rename(key, new_key)
        ret_value = self.DD[new_key]
        # Verify that the key is properly renamed
        self.assertEqual(value, ret_value, "Rename key is not successful")


class TestDragonDictMultiNodeStress(unittest.TestCase):
    @classmethod
    def increment_values(self, client_id, ev, ddict, locks):
        """Each client will increment the value in all keys of the dictionary.
        A corresponding lock will be taken for each key while incrementing the value.

        :param client_id: id of the client
        :type client_id: int
        :param ev: event for the client
        :type ev: multiprocessing event object
        :type ddict: dragon dictionary object
        :param ddict: dragon distributed dictionary
        :type locks: dict of locks of multiprocessing value object for each key
        :type locks: dict
        """
        # Add key to the dictionary, and wait for the event
        ddict[client_id] = int(0)
        ev.wait()

        keys = ddict.keys()
        random.shuffle(keys)
        for key in keys:
            # Acquire the lock for the key and increment the associated value
            success = locks[key].acquire(block=True, timeout=None)
            assert success == True, "Could not acquire lock"
            value = int(ddict[key])
            value += 1
            ddict[key] = int(value)
            locks[key].release()

        # Close the client
        ddict.close()

    @classmethod
    def do_dict_workload(self, num_clients, ddict):
        """Test the workload on the dictionary, with each client incrementing the values
        in each key of the dictionary. A corresponding lock will be taken for each key
        while incrementing the value for synchronization purposes. Each key will have a
        multiprocessing value object to be acquired before incrementing the value.

        :param num_clients: total number of clients
        :type num_clients: int
        :type ddict: dragon dictionary object
        :param ddict: dragon distributed dictionary
        """
        procs = []
        locks = []
        event = mp.Event()
        for i in range(num_clients):
            locks.append(mp.Lock())

        for i in range(num_clients):
            client_proc = mp.Process(target=self.increment_values, args=(i, event, ddict, locks,))
            client_proc.start()
            procs.append(client_proc)

        wait = True
        while wait:
            if(len(ddict) == num_clients):
                wait = False

        # Wait is done
        event.set()
        for i in range(len(procs)):
            procs[i].join()

        # Clean all the processes
        for i in range(len(procs)):
            procs[i].kill()

    def test_dict_with_stress_load(self):
        managers_per_node = 1 # 2 Managers per node
        num_nodes = len(get_list()) # Collect the total number of nodes
        total_mem_size = 4*(1024*1024*1024) # 4 GB total size
        ddict = DragonDict(managers_per_node, num_nodes, total_mem_size)

        # Being conservative, increase to the max cpu_count later
        num_clients = max(2, mp.cpu_count() // 4)
        self.do_dict_workload(num_clients, ddict)

        # Verify the values are updated properly
        # Each value should be equal to the num_clients
        values = ddict.values()
        for value in values:
            self.assertEqual(value, num_clients, "Value is not equal to the number of clients in the dictionary")
        ddict.stop()


if __name__ == "__main__":
    mp.set_start_method("dragon")

    # Disable the non-deterministic behavior of the python hashing algorithm behavior
    hashseed = os.getenv('PYTHONHASHSEED')
    if not hashseed:
        os.environ['PYTHONHASHSEED'] = '0'
        os.execv(sys.executable, [sys.executable] + sys.argv)
    unittest.main()

test/multi-node/test_dict.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
`dragon.native.Lock` object.  The test scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.

The test is run with `dragon test_lock.py -f -v`
"""

import unittest
import pickle
import sys

import dragon
from dragon.globalservices.process import (
    multi_join,
    create,
    get_list,
    join,
)
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.infrastructure.parameters import this_process
from dragon.native.lock import Lock
from dragon.native.machine import cpu_count
import dragon.utils as du


def _burner(env_str):

    num_acq = 128

    # unpack the lock from the environment string
    bytes = du.B64.str_to_bytes(env_str)
    lock = pickle.loads(bytes)

    for _ in range(num_acq):
        success = lock.acquire(block=True, timeout=None)
        lock.release()  # burn it !
        assert success == True, "Could not acquire lock"


def _releaser(env_str):

    # unpack the lock from the environment string
    bytes = du.B64.str_to_bytes(env_str)
    rlock = pickle.loads(bytes)

    # try to release it
    try:
        rlock.release()
    except AssertionError:
        pass
    else:
        assert False, "Release should not have been successful"


def _acquirer(env_str):

    # unpack the lock from the environment string
    bytes = du.B64.str_to_bytes(env_str)
    lock = pickle.loads(bytes)

    lock.acquire(block=True, timeout=None)
    assert lock.is_mine(), "is_mine() broken"


class TestLockMultiNode(unittest.TestCase):
    def test_lock_basic(self):
        """Test that processes actually wait on the lock and can be released
        by any other process.
        """

        nprocs = max(2, cpu_count() // 8)  # all sleeping in a lock

        # get a lock
        lock = Lock(recursive=False)
        self.assertTrue(lock.acquire())  # head process owns the lock now
        self.assertFalse(lock.acquire(timeout=0))  # cannot acquire again

        # start some processes
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(lock, protocol=5))

        puids = []
        for _ in range(nprocs):
            the_descr = create(
                cmd,
                wdir,
                [__file__, "_acquirer", env_str],
                None,
                options=options,
            )
            puids.append(the_descr.p_uid)

        # all should be waiting, we still own the lock
        live_puids = get_list()
        live_puids.remove(this_process.my_puid)

        for puid in puids:
            self.assertTrue(puid in live_puids)

        # release one by one
        for _ in puids:
            lock.release()
            ready = multi_join(live_puids, timeout=None, join_all=False)
            self.assertTrue(len(ready[0]) == 1)  # only one has exited !
            ecode = ready[0][0][1]
            self.assertTrue(ecode == 0)
            puid = ready[0][0][0]
            live_puids.remove(puid)

        # cause a value error
        lock.release()  # last puid left the Lock in a locked state
        self.assertRaises(ValueError, lock.release)

    def test_recursive(self):
        """Test that recursive locks can be acquired multiple times and
        need to be released multiple times by the same process
        """

        # get a recursive lock
        rlock = Lock(recursive=True)

        # acquire it 25 times
        for _ in range(25):
            self.assertTrue(rlock.acquire())
        self.assertTrue(rlock._accesscount == 25)

        # start a processes that tries to release it
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(rlock, protocol=5))

        the_descr = create(
            cmd,
            wdir,
            [__file__, "_releaser", env_str],
            None,
            options=options,
        )
        puid = the_descr.p_uid

        # wait until it has tried, counter must still be 25
        ready = multi_join([puid], timeout=None, join_all=True)

        # release 25 times
        for _ in range(25):
            rlock.release()  # cannot raise an exception
        self.assertTrue(rlock._accesscount == 0)

        # counter is now 0
        self.assertRaises(AssertionError, rlock.release)

    def test_rlock_ownership(self):
        """Test that only the calling thread owns the lock."""

        # get a recursive lock
        rlock = Lock(recursive=True)

        # acquire it
        self.assertTrue(rlock.acquire())
        self.assertTrue(rlock.is_mine())

        # still mine after pickle ?
        bytes = pickle.dumps(rlock)
        rlock2 = pickle.loads(bytes)
        self.assertTrue(rlock2.is_mine())

        # release it
        rlock.release()
        self.assertFalse(rlock.is_mine())

        # still not mine after pickle ?
        bytes = pickle.dumps(rlock)
        rlock3 = pickle.loads(bytes)
        self.assertFalse(rlock3.is_mine())

        # start a processes that acquires the lock and leaves
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(rlock, protocol=5))

        the_descr = create(
            cmd,
            wdir,
            [__file__, "_acquirer", env_str],
            None,
            options=options,
        )
        puid = the_descr.p_uid

        ready = join(puid, timeout=None)
        self.assertTrue(ready == 0)  # process has exited

        # should not be mine and I cannot release it
        self.assertFalse(rlock.is_mine())
        self.assertRaises(AssertionError, rlock.release)

    def test_lock_contention(self):
        """Test lock correctness under contention."""

        # get a recursive lock
        lock = Lock(recursive=False)

        nprocs = max(2, cpu_count() // 8)

        # create env string
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(lock, protocol=5))

        # start processes
        puids = []
        for _ in range(nprocs):
            the_descr = create(
                cmd,
                wdir,
                [__file__, "_burner", env_str],
                None,
                options=options,
            )
            puids.append(the_descr.p_uid)

        # let them finish
        ready = multi_join(puids, timeout=None, join_all=True)

        # check they all have exited without error
        for puid, exit_code in ready[0]:
            self.assertTrue(exit_code == 0)


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "_acquirer":
        _acquirer(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "_releaser":
        _releaser(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "_burner":
        _burner(sys.argv[2])
    else:
        unittest.main()

test/multi-node/test_lock.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
`dragon.native.machine` module. The tests scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.

The test is run with `dragon test_machine.py -f -v`
"""

import unittest
import os

import dragon
from dragon.native.machine import cpu_count
import multiprocessing as mp


class TestMachineMultiNode(unittest.TestCase):
    """We need to test `cpu_count, or all other tests could be bogus."""

    def test_on_multi_node(self):

        ncpu_native = max(2, cpu_count())
        ncpu_mp = max(2, mp.cpu_count())
        ncpu_posix = int(os.sysconf("SC_NPROCESSORS_ONLN"))

        self.assertTrue(ncpu_mp > ncpu_posix, "You're not using a multi-node allocation.")
        self.assertTrue(ncpu_native > ncpu_posix, "You're not using a multi-node allocation.")


if __name__ == "__main__":
    mp.set_start_method("dragon")
    unittest.main()

test/multi-node/test_machine.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
PMOD MPI job launcher. The test runs 4 MPI jobs simulataneously and
waits for their completion.

The test is run with `dragon test_mpi_hello_world.py -f -v`
"""

import unittest
import dragon.channels as dch
import dragon.managed_memory as dmm
import dragon.pmod as djm
from dragon.utils import B64
from dragon.dtypes import DragonError
import multiprocessing as mp
import socket
import subprocess
import os

SUCCESS = DragonError.SUCCESS
FAILURE = DragonError.FAILURE

# tracks base PIDs currently in use by active MPI jobs
pid_dict = {}


class MPIJobParams:
    def __init__(self, ppn, nid, nnodes, job_id):
        self.ppn = ppn
        self.nid = nid
        self.nnodes = nnodes
        self.id = job_id


def allocate_mpi_pid_base(pid_bases, nranks):

    # reserve 0 for the single HSTA agent. this will not be needed
    # when we begin launching HSTA with PMOD
    min_pid = 1

    # handle the case where pid_bases is empty

    if not pid_bases:
        pid_bases.append((min_pid, nranks))
        return min_pid

    # loop over (base, size) pairs looking for a free segment of PIDS
    # with size at least nranks

    # (1) for a given (base, size) pair, check the segment starting at
    # the PID just past the previous segment

    prev_end = min_pid

    for base, size in pid_bases:
        if base - prev_end >= nranks:
            pid_bases.append((prev_end, nranks))
            pid_bases.sort()
            return prev_end
        else:
            prev_end = base + size

    # (2) check the final segment from the end of the last (base, size)
    # pair extending to the maximum PID

    max_pid = 510

    if (max_pid + 1) - prev_end >= nranks:
        pid_bases.append((prev_end, nranks))
        pid_bases.sort()
        return prev_end

    return -1


def free_mpi_pid_base(pid_bases, base, size):

    pid_bases.remove((base, size))


def node_shepherd(job_params, q_to_spawner):

    # get list of hostnames for all nodes in job

    hostname = socket.gethostname()
    q_to_spawner.put(hostname)
    hostname_list = q_to_spawner.get()
    base_pid = q_to_spawner.get()

    if job_params.nid == 0:
        print(f"Running MPI job with id {job_params.id} on the following hosts: {hostname_list}", flush=True)

    # create memory pool

    pool_size = 2**30
    # each job can have up to 2**20 nodes
    pool_uid = (2**20 * job_params.id) + job_params.nid
    pool_name = "test_pmod_mpi_hello_world-" + str(job_params.id)
    mpool = dmm.MemoryPool(pool_size, pool_name, pool_uid)

    # we are currently assuming an equal number of ranks per node
    # (but this isn't necessary, just convenient for this test)

    nranks = job_params.ppn * job_params.nnodes

    nidlist = []

    for nid in range(job_params.nnodes):
        for lrank in range(job_params.ppn):
            nidlist.append(nid)

    pmod = djm.PMOD(
        job_params.ppn, job_params.nid, job_params.nnodes, nranks, nidlist, hostname_list, job_params.id
    )

    child_handles = []

    os.environ["LD_PRELOAD"] = "libdragon.so"
    os.environ["MPICH_OFI_CXI_PID_BASE"] = str(base_pid)
    os.environ["DL_PLUGIN_RESILIENCY"] = "1"
    os.environ["_DRAGON_PALS_ENABLED"] = "1"

    for lrank in range(job_params.ppn):
        # create and serialize a channel that we can use to talk to this child process

        channel_uid = 2**20 * job_params.id + lrank + 42
        child_ch = dch.Channel(mpool, channel_uid)
        os.environ["DRAGON_PMOD_CHILD_CHANNEL"] = str(B64(child_ch.serialize()))

        # no way to dynamically search for a free port, since the first free port
        # may differ from node to node, depending on how MPI jobs are laid out. so
        # we simply use a unique offset from a arbitrary starting port

        port = 7777 + job_params.id
        os.environ["PMI_CONTROL_PORT"] = str(port)

        handle = subprocess.Popen(
            "./mpi_hello",
        )

        # now send the basic PMI data to the child process

        pmod.send_mpi_data(lrank, child_ch)

        child_handles.append(handle)

    passed = True

    for handle in child_handles:
        handle.wait()
        if handle.returncode != SUCCESS:
            passed = False

    q_to_spawner.put(passed)

    if job_params.nid == 0:
        print(f"child processes for job {job_params.id} are done!")

    return passed


def start_mpi_job(job_id, ppn, nnodes):

    process_handles = []
    mpi_pid_bases = []
    queue_list = []
    hostname_list = []

    # start a shepherd on each node in the job

    for nid in range(nnodes):
        job_params = MPIJobParams(ppn, nid, nnodes, job_id)

        q = mp.Queue()
        queue_list.append(q)

        p = mp.Process(
            target=node_shepherd,
            args=(
                job_params,
                q,
            ),
        )
        p.start()
        process_handles.append(p)

        hostname = q.get()
        hostname_list.append(hostname)

    # broadcast the hostname list back to each node

    for q in queue_list:
        q.put(hostname_list)

    # finally allocate and send a pid base to each nid

    hostname_and_q_list = zip(hostname_list, queue_list)

    for hostname, q in hostname_and_q_list:
        if hostname in pid_dict:
            allocated_bases = pid_dict[hostname]
        else:
            allocated_bases = []
            pid_dict[hostname] = allocated_bases

        pid_base = allocate_mpi_pid_base(allocated_bases, ppn)
        assert pid_base != -1, "unable to allocate PIDs for MPI job"

        mpi_pid_bases.append((hostname, pid_base, ppn))

        q.put(pid_base)

    return process_handles, queue_list, mpi_pid_bases


def wait_for_job(process_handles, queue_list, mpi_pid_bases):

    passed = True

    proc_and_q_list = zip(process_handles, queue_list)

    for p, q in proc_and_q_list:
        node_passed = q.get()
        passed = passed and node_passed
        p.join()

    for hostname, base, size in mpi_pid_bases:
        allocated_bases = pid_dict[hostname]
        free_mpi_pid_base(allocated_bases, base, size)

        if not allocated_bases:
            del pid_dict[hostname]

    return passed


class TestPMODMultiNode(unittest.TestCase):
    def test_four_jobs(self):
        job_id = 0
        ppn = 7
        nnodes = 3

        process_handles_0, q_list_0, mpi_pid_bases_0 = start_mpi_job(job_id, ppn, nnodes)

        job_id = 1
        ppn = 4
        nnodes = 2

        process_handles_1, q_list_1, mpi_pid_bases_1 = start_mpi_job(job_id, ppn, nnodes)

        job_id = 2
        ppn = 2
        nnodes = 3

        process_handles_2, q_list_2, mpi_pid_bases_2 = start_mpi_job(job_id, ppn, nnodes)

        job_id = 3
        ppn = 1
        nnodes = 4

        process_handles_3, q_list_3, mpi_pid_bases_3 = start_mpi_job(job_id, ppn, nnodes)

        passed = True

        job_passed = wait_for_job(process_handles_0, q_list_0, mpi_pid_bases_0)
        passed = passed and job_passed

        job_passed = wait_for_job(process_handles_1, q_list_1, mpi_pid_bases_1)
        passed = passed and job_passed

        job_passed = wait_for_job(process_handles_2, q_list_2, mpi_pid_bases_2)
        passed = passed and job_passed

        job_passed = wait_for_job(process_handles_3, q_list_3, mpi_pid_bases_3)
        passed = passed and job_passed

        os.system("./cleanup.sh")

        if not job_passed:
            exit(1)


if __name__ == "__main__":
    mp.set_start_method("dragon")
    unittest.main()

test/multi-node/test_mpi_hello_world.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
`multiprocessing.Pool` object.  The test scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.

The test is run with `dragon test_pool.py -f -v`
"""

import unittest
import time
import itertools

import dragon
import multiprocessing as mp
from multiprocessing import TimeoutError

import numpy as np

TIMEOUT1, TIMEOUT2, TIMEOUT3 = 0.82, 0.35, 1.4

def sqr(x, wait=0.0):
    time.sleep(wait)
    return x * x


def mul(x, y):
    return x * y


def raise_large_valuerror(wait):
    time.sleep(wait)
    raise ValueError("x" * 1024**2)


def identity(x):
    return x


class TestPoolMultiNode(unittest.TestCase):
    """These are multi-node versions of some of the Multiprocessing Pool unit
    tests."""

    @classmethod
    def setUpClass(cls):
        ncpu = max(2, mp.cpu_count() // 8)
        cls.pool = mp.Pool(ncpu)

    @classmethod
    def tearDownClass(cls):
        cls.pool.terminate()
        cls.pool.join()
        cls.pool = None

    def test_apply(self):
        papply = self.pool.apply
        self.assertEqual(papply(sqr, (5,)), sqr(5))
        self.assertEqual(papply(sqr, (), {"x": 3}), sqr(x=3))

    def test_map(self):
        pmap = self.pool.map
        self.assertEqual(pmap(sqr, list(range(10))), list(map(sqr, list(range(10)))))
        self.assertEqual(pmap(sqr, list(range(100)), chunksize=20), list(map(sqr, list(range(100)))))

    def test_starmap(self):
        psmap = self.pool.starmap
        tuples = list(zip(range(10), range(9, -1, -1)))
        self.assertEqual(psmap(mul, tuples), list(itertools.starmap(mul, tuples)))
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(psmap(mul, tuples, chunksize=20), list(itertools.starmap(mul, tuples)))

    def test_starmap_async(self):
        tuples = list(zip(range(100), range(99, -1, -1)))
        self.assertEqual(self.pool.starmap_async(mul, tuples).get(), list(itertools.starmap(mul, tuples)))

    def test_map_async(self):
        self.assertEqual(self.pool.map_async(sqr, list(range(10))).get(), list(map(sqr, list(range(10)))))

    def test_map_async_callbacks(self):
        call_args = []
        self.pool.map_async(int, ["1"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(1, len(call_args))
        self.assertEqual([1], call_args[0])
        self.pool.map_async(int, ["a"], callback=call_args.append, error_callback=call_args.append).wait()
        self.assertEqual(2, len(call_args))
        self.assertIsInstance(call_args[1], ValueError)

    def test_async(self):
        res = self.pool.apply_async(
            sqr,
            (
                7,
                TIMEOUT1,
            ),
        )
        start = time.monotonic()
        self.assertEqual(res.get(), 49)
        stop = time.monotonic()
        self.assertAlmostEqual(stop - start, TIMEOUT1, 1)

    def test_async_timeout(self):
        res = self.pool.apply_async(sqr, (6, TIMEOUT2 + 1.0))
        start = time.monotonic()
        try:
            res.get(timeout=TIMEOUT2)
            self.assertTrue(False)
        except TimeoutError:
            pass
        stop = time.monotonic()
        self.assertAlmostEqual(stop - start, TIMEOUT2, 1)

    def test_imap(self):
        it = self.pool.imap(sqr, list(range(10)))
        self.assertEqual(list(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap(sqr, list(range(10)))
        for i in range(10):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

        it = self.pool.imap(sqr, list(range(1000)), chunksize=100)
        for i in range(1000):
            self.assertEqual(next(it), i * i)
        self.assertRaises(StopIteration, it.__next__)

    def test_imap_unordered(self):
        it = self.pool.imap_unordered(sqr, list(range(10)))
        self.assertEqual(sorted(it), list(map(sqr, list(range(10)))))

        it = self.pool.imap_unordered(sqr, list(range(1000)), chunksize=100)
        self.assertEqual(sorted(it), list(map(sqr, list(range(1000)))))


class TestPoolScalingMultiNode(unittest.TestCase):
    def test_strong_scalability(self) -> None:
        """Trivial process throughput/scalability benchmark:
        https://parsl.readthedocs.io/en/stable/userguide/performance.html
        Strong Scaling
        """

        N = 5000

        maxcpus = np.flip(np.logspace(0, np.log10(max(2, mp.cpu_count() // 8)), num=16))

        params = [1 for _ in range(N)]

        for cpus in maxcpus:
            start = time.monotonic()
            with mp.Pool(int(cpus)) as pool:
                pool.map(time.sleep, params)
            stop = time.monotonic()
            print(f"{N=} {int(cpus)=} {stop-start=} {(stop-start)*cpus/max(maxcpus)=}", flush=True)

    def test_strong_scalability(self) -> None:
        """Trivial process throughput/scalability benchmark:
        https://parsl.readthedocs.io/en/stable/userguide/performance.html
        Weak Scaling
        """

        maxcpus = np.flip(np.logspace(0, np.log10(max(2, mp.cpu_count() // 8)), num=16))

        for cpus in maxcpus:
            N = int(10 * cpus)
            params = [1 for _ in range(N)]
            start = time.monotonic()
            with mp.Pool(int(cpus)) as pool:
                pool.map(time.sleep, params)
            stop = time.monotonic()
            print(f"{N=} {int(cpus)=} {stop-start=}", flush=True)


if __name__ == "__main__":
    mp.set_start_method("dragon")
    unittest.main()

test/multi-node/test_pool.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
`multiprocessing.Process` object. The test scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.

The test is run with `dragon test_process.py -f -v`
"""

import unittest
import time

import dragon
import multiprocessing as mp
from dragon.globalservices.process import query, multi_join, this_process


def inception(nnew: int, q: mp.Queue, ev1: mp.Event, ev2: mp.Event, sem: mp.Semaphore) -> None:

    ev1.wait(timeout=None)

    for _ in range(nnew):

        if not sem.acquire(timeout=0.1):  # try for a while
            break  # stop spawning

        p = mp.Process(target=inception, args=(nnew, q, ev1, ev2, sem))
        p.start()
        q.put(p.sentinel)

    ev2.wait(timeout=None)


class TestProcessMultiNode(unittest.TestCase):
    def test_inception(self) -> None:
        """Have processes spawn `nnew` processes, until `nchildren` have been
        spawned."""

        nchildren = max(2, mp.cpu_count())
        nnew = 4

        q = mp.Queue(maxsize=nchildren)
        sem = mp.Semaphore(value=nchildren)
        ev1 = mp.Event()
        ev2 = mp.Event()

        processes = []
        for _ in range(nnew):
            sem.acquire()
            p = mp.Process(target=inception, args=(nnew, q, ev1, ev2, sem))
            p.start()
            q.put(p.sentinel)
            processes.append(p)

        ev1.set()

        puids = []
        while not len(puids) == nchildren:
            child_puid = q.get(timeout=None)
            puids.append(child_puid)

        parents = puids + [this_process.my_puid]

        for puid in puids:
            descr = query(puid)
            self.assertTrue(descr.p_uid == puid)
            self.assertTrue(descr.p_p_uid in parents)

        ev2.set()

        mp.connection.wait(puids, timeout=None)

        start = time.monotonic()
        ready = multi_join(puids, timeout=None, join_all=True)  # should exit "immediately"
        stop = time.monotonic()

        # self.assertTrue(2 > (stop - start))  # performance of GS

        for puid, ecode in ready[0]:
            self.assertTrue(ecode == 0)
            self.assertTrue(puid in puids)

        for p in processes:
            self.assertTrue(p.exitcode == 0)


if __name__ == "__main__":
    mp.set_start_method("dragon")
    unittest.main()

test/multi-node/test_process.py
DragonHPC/dragon
"""This file contains multi-node tests for the process group classes.
"""
import os
import unittest
import time
import random
import signal

import dragon

from dragon.native.process_group import ProcessGroup
from dragon.native.process import TemplateProcess, Process
from dragon.native.machine import cpu_count, current

from dragon.globalservices.process import query as process_query, kill, signal as dragon_signal
from dragon.globalservices.node import query as node_query, get_list as node_get_list
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.infrastructure.policy import Policy


class TestProcessGroupMultiNode(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        """Set parameters"""

        # a few standard values
        cls.nproc = cpu_count() * 2
        cls.cmd = "sleep"
        cls.args = ("1000",)
        cls.cwd = os.getcwd()
        cls.options = ProcessOptions(make_inf_channels=True)
        cls.template = TemplateProcess(cls.cmd, args=cls.args, cwd=cls.cwd)

    @unittest.skip(f"CIRRUS-1163: Will fail until process descriptor contains a h_uid/host_id.")
    def test_placement_roundrobin(self):

        pg = ProcessGroup(restart=False)
        pg.add_process(self.nproc, self.template)
        pg.init()

        pg.start()

        self.assertTrue(pg.status == "Running")

        puids = pg.puids  # list is ordered by placement policy !

        huids = node_get_list()

        for i, puid in enumerate(puids):

            idx = i % len(huids)

            p = Process(None, ident=puid)
            self.assertTrue(
                p.node == huids[i], f"{p.node=} {huids[idx]=} {idx=} "
            )  # will fail until placement is fixed to h_uid

        pg.kill()

        self.assertTrue(pg.status == "Idle")

        pg.stop()

        self.assertTrue(pg.status == "Stop")

    def test_maintain_stress(self):
        """This test may take a minute."""

        # we will keep killing processes, 10% will keep exiting on their own.

        testtime = 10  # sec

        pg = ProcessGroup(ignore_error_on_exit=True, policy=Policy())

        count = 0
        for i in range(self.nproc):
            if i % 10 == 0:
                args = (f"{testtime/self.nproc * (i+1)}",)
                t = TemplateProcess("sleep", args=args, cwd=self.cwd)
                pg.add_process(1, t)
            else:
                count += 1
        args = ("10000000",)
        t = TemplateProcess("sleep", args=args, cwd=self.cwd)
        pg.add_process(count, t)

        pg.init()

        man_proc = pg._manager._proc
        self._update_interval_sec = 0.1

        pg.start()

        beg = time.monotonic()

        while time.monotonic() - beg < testtime:

            puids = pg.puids

            try:
                puid = puids[random.randint(0, len(puids))]
                kill(puid, sig=dragon_signal.SIGKILL)
            except Exception:  # maybe it disappeared already
                pass

            self.assertTrue(man_proc.is_alive)
            self.assertTrue(pg.status != "Error")

            time.sleep(testtime // 10)

        puids = pg.puids

        pg.kill(signal.SIGTERM)

        while not pg.status == "Idle":
            self.assertTrue(pg.status != "Error")

        self.assertTrue(man_proc.is_alive)

        for puid in puids:
            gs_info = process_query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

        self.assertTrue(pg.status == "Stop")

        man_proc.join()

    def test_walltime(self):

        wtime = 3

        # small number of procs, so that the startup overhead is minimal
        nproc = len(node_get_list())

        pg = ProcessGroup(walltime=wtime)
        pg.add_process(nproc, self.template)
        pg.init()

        pg.start()

        start = time.monotonic()
        while not pg.status == "Idle":
            self.assertFalse(pg.status == "Error")
            time.sleep(0.5)

        stop = time.monotonic()

        self.assertAlmostEqual(stop - start, wtime, 0)

        pg.stop()


if __name__ == "__main__":
    unittest.main()

test/multi-node/test_process_group.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance tests for the
`multiprocessing.Queue` object.  The test scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.
The test is run with `dragon test_queue.py -f -v`
"""

import unittest
import random

import dragon
import multiprocessing as mp

max_msg_size = 1024  # limit to 1KB


def _joiner(q, item):
    q.put(item)
    q.join()


class TestQueueMultiNode(unittest.TestCase):
    @classmethod
    def _writer(self, id, q, nitems, msg_list, qver):
        for i in range(nitems):
            msg_idx = id * nitems + i
            q.put((msg_idx, msg_list[i]))
            qver.put((msg_idx, msg_list[i]))
        q.close()
        qver.close()

    @classmethod
    def _reader(self, id, q, nitems):
        for i in range(nitems):
            q.get()

        q.close()

    def test_multi_producer_consumer(self):
        """Start 2 processes per cpu and have each of them send 16 messages to
        each other. Message sizes rotate between 1 byte and 1 Mbyte"""

        num_readers = num_writers = max(2, mp.cpu_count() // 8)

        msg_list = {}  # key is the msg identifier, value is the payload
        nitems = 42  # Seems like the right value.

        for i in range(nitems):
            msg_list[i] = random.randbytes(max_msg_size)

        q = mp.Queue()
        q_ver = mp.Queue(
            maxsize=nitems * num_writers
        )  # used to verify that all messages are gotten from the queue

        writers = []
        for i in range(num_writers):
            p = mp.Process(target=self._writer, args=(i, q, nitems, msg_list, q_ver))
            p.start()
            writers.append(p)

        readers = []
        reader_qlist = []
        for i in range(num_readers):
            p = mp.Process(target=self._reader, args=(i, q, nitems))
            p.start()
            readers.append(p)

        for p in writers:
            p.join()

        for p in readers:
            p.join()

        self.assertTrue(q_ver.qsize() == nitems * num_writers)

        while not q_ver.empty():
            _id, payload = q_ver.get()
            if _id in msg_list:
                del msg_list[_id]

        self.assertEqual(len(msg_list), 0)
        self.assertTrue(q_ver.empty())

        self.assertTrue(q.empty())

        q.close()
        q_ver.close()

    def test_joinable(self):
        """Test joinability of a multi-node queue"""

        q = mp.JoinableQueue()
        num_processes = max(2, mp.cpu_count() // 8)

        self.assertTrue(q._joinable == True)

        items = [f"JQueue-item-{i}" for i in range(num_processes)]

        processes = []
        for item in items:
            p = mp.Process(target=_joiner, args=(q, item))
            p.start()
            processes.append(p)

        for _ in range(num_processes):
            item = q.get()
            self.assertTrue(item in items)

        for p in processes:  # make sure they are all blocked
            p.join(timeout=0)
            self.assertTrue(p.exitcode == None)

        for p in processes:  # unblock workers
            q.task_done()

        for p in processes:
            p.join(timeout=None)
            self.assertTrue(p.exitcode == 0)


if __name__ == "__main__":
    mp.set_start_method("dragon", force=True)
    unittest.main()

test/multi-node/test_queue.py
DragonHPC/dragon
""" This file contains Dragon multi-node acceptance test for the
`dragon.native.Value` object.  The test scales with the total number of CPUs
reported by the allocation, i.e. it becomes tougher on larger allocations.

The test is run with `dragon test_value.py -f -v`
"""
import unittest
import os
import sys
import numpy as np

import dragon
import multiprocessing as mp


def process_function(
    event,
    value,
    queue,
    max_val,
):

    p = mp.current_process()
    mypuid = p.pid

    # each child process waits on parent process
    event.wait()

    val = -1

    while val < max_val:
        with value.get_lock():
            val = value.value
            if val <= max_val:
                value.value += 1
                queue.put(
                    (
                        mypuid,
                        val,
                    )
                )


class TestValueMultiNode(unittest.TestCase):
    def test_lock_fairness(
        self,
    ):
        """Tests that value can be incremented by multiple processes in a
        multinode environment.

        1. The Test uses `nproc = mp.cpu_count()` processes.
        2. Each child process gets handed the same `Value(int(0), lock=True)`, a
           Queue, an Event and the max_val.
        3. Each child process first waits in the Event.
        4. Once all processes have been started by the parent, the parent sets
           the event, so that all processes start working at the same time. This
           means the runtime sees a very sudden increase in activity, which is a
           test in itself.
        5. Each child process tries to acquire the lock and increase the value
           associated with the `Value`. This (step 6) is repeated until Value is equal to max_val.
        6. When the lock was acquired, the child process saves the current value
           from the Value object in a variable, and then increments the value in
           the Value object by 1. It then puts the saved value and its PUID into
           the queue and releases the lock.
        7. The parent process `get`s from the Queue `max_val` times without a
           timeout and checks that the values it is receiving from the Queue are
           monotonically increasing by 1 each time.
        8. The parent process counts how often it received a value from any of
           the puids. It prints the minimum, maximum, mean and median of the number
           of times.
        9. Finally the parent joins on all processes and check that their
           exit_code is 0, i.e. that they exited cleanly without an exception
        """

        procs = []
        puid_dict = {}

        # number of processes for the test
        nproc = max(2, mp.cpu_count() // 16)
        max_val = nproc * 8

        # each child process is handed the same value, queue, and event
        value = mp.Value(
            "i",
            0,
            lock=True,
        )
        queue = mp.Queue()
        event = mp.Event()

        for _ in range(nproc):
            p = mp.Process(
                target=process_function,
                args=(
                    event,
                    value,
                    queue,
                    max_val,
                ),
            )
            p.start()
            procs.append(p)
            puid_dict[p.pid] = 0

        # parent sets event for children processes to work at same time
        event.set()

        cnt = 0

        for _ in range(
            0,
            max_val,
        ):

            (
                puid,
                val,
            ) = queue.get()

            # check that the value from child process is valid
            self.assertTrue(val == cnt)
            cnt += 1

            # increment each time child puid returns value
            puid_dict[puid] += 1

        numpy_puids = np.array(list(puid_dict.values()))

        # pring minimum, maximum, mean, and median from puids_dict
        (puids_max, puids_min, puids_mean, puids_median,) = (
            np.max(numpy_puids),
            np.min(numpy_puids),
            np.mean(numpy_puids),
            np.median(numpy_puids),
        )

        print(f"Lock acquisition, {max_val/nproc} tries / proc, {nproc} processes:")
        print(f"    Maximum: {puids_max}, Minimum: {puids_min}, Mean: {puids_mean}, Median: {puids_median}")

        # check that the child process exited cleanly in parent process
        for p in procs:
            p.join()
            self.assertEqual(
                p.exitcode,
                0,
            )


if __name__ == "__main__":
    mp.set_start_method(
        "dragon",
        force=True,
    )
    unittest.main()

test/multi-node/test_value.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import pickle
import sys
import ctypes
import random
import string

import dragon
from dragon.globalservices.process import create, multi_join
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.native.array import Array, _TYPECODE_TO_TYPE, _SUPPORTED_TYPES
from dragon.native.queue import Queue
import dragon.utils as du


def create_array_assignment(type):

    if type is ctypes.c_char:
        return [b"0", b"1", bytes(str(random.choice(string.ascii_letters)).encode("utf-8"))]
    elif type is ctypes.c_wchar:
        return ["0", "1", str(random.choice(string.ascii_letters))]
    elif type in [ctypes.c_double, ctypes.c_float]:
        return [
            0.0,
            1.0,
            random.uniform(sys.float_info.min, sys.float_info.max),
            sys.float_info.min,
            sys.float_info.max,
        ]
    else:
        return [0, 1, random.randint(-sys.maxsize - 1, sys.maxsize), -sys.maxsize - 1, sys.maxsize]


def test_array(args):

    bytes = du.B64.str_to_bytes(args)
    array, array_assignment, parent_queue, child_queue = pickle.loads(bytes)

    # set array.size_or_initializer to array_assignment
    array[:] = array_assignment
    # put True into queue
    child_queue.put(True)
    # wait on the queue and exit cleanly
    assert parent_queue.get(timeout=None)


class TestArray(unittest.TestCase):
    def test_typedef_or_type(self):
        """typedef_or_type type checking"""
        self.assertRaises(AttributeError, lambda: Array(typecode_or_type=0, size_or_initializer=0))
        self.assertRaises(AttributeError, lambda: Array(typecode_or_type="v", size_or_initializer=0))
        self.assertRaises(AttributeError, lambda: Array(typecode_or_type=(), size_or_initializer=0))

    def test_array_assignment(self):
        """check array assignment"""
        self.assertEqual(Array("i", [1, 2, 3])[:], [1, 2, 3])
        arr = Array("i", 10)
        arr[:] = range(10)
        self.assertEqual(list(arr), list(range(10)))

    def test_ping_pong(self):
        """queue ping pong between 2 processes that tests array assignment in
        parent and child processes"""

        for type in _SUPPORTED_TYPES:

            # Create a list of values to be tested array_assignment: bytes, char, int, and float
            array_assignment = create_array_assignment(type)

            # Create a array of type
            array = Array(typecode_or_type=type, size_or_initializer=[])

            # Start a process, hand over v, a queue q and the array to be tested array_assignment
            cmd = sys.executable
            wdir = "."
            options = ProcessOptions(make_inf_channels=True)
            parent_queue, child_queue = Queue(joinable=False), Queue(joinable=False)
            puids = []
            env_str = du.B64.bytes_to_str(
                pickle.dumps([array, array_assignment, parent_queue, child_queue], protocol=5)
            )
            proc = create(
                cmd,
                wdir,
                [__file__, "test_array", env_str],
                None,
                options=options,
            )

            # wait on the queue _ = q.get(timeout=None)
            self.assertTrue(child_queue.get(timeout=None))

            # check each element in the array to make sure that the array assignment is correct
            self.assertEqual(list(array[:]), array_assignment)

            for i in range(len(array)):
                self.assertEqual(array[i], array_assignment[i])

            #put True into the queue q.put(True)
            parent_queue.put(True)

            puids.append(proc.p_uid)

        # join on the puid of the process
        ready = multi_join(puids)

        # check the exit code of the process from ready = multi_join(puids) from ready[0][0]
        self.assertEqual(ready[0][0][1], 0)


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "test_array":
        test_array(sys.argv[2])
    else:
        unittest.main()

test/native/test_array.py
DragonHPC/dragon
import os
import unittest
import time
import pickle
import sys

import dragon

from dragon.native.process import current
from dragon.native.queue import Queue
from dragon.native.barrier import Barrier, BrokenBarrierError

from dragon.globalservices.process import (
    join,
    create,
    multi_join,
)
from dragon.infrastructure.process_desc import ProcessOptions
import dragon.utils as du


def timeout_wait_barrier(args):
    bytes = du.B64.str_to_bytes(args)
    barrier = pickle.loads(bytes)[0]
    try:
        barrier.wait(timeout=2)
    except BrokenBarrierError:
        pass


def wait_barrier_forever(args):
    bytes = du.B64.str_to_bytes(args)
    barrier = pickle.loads(bytes)[0]
    barrier.wait(timeout=None)


def wait_with_broken_barrier_helper(args):
    bytes = du.B64.str_to_bytes(args)
    barrier = pickle.loads(bytes)[0]
    try:
        barrier.wait(timeout=None)
    except BrokenBarrierError:
        pass


def reset_barrier(args):
    bytes = du.B64.str_to_bytes(args)
    barrier = pickle.loads(bytes)[0]
    barrier.reset()


def abort_barrier(args):
    bytes = du.B64.str_to_bytes(args)
    barrier = pickle.loads(bytes)[0]
    barrier.abort()


def barrier_wait_helper(args):
    bytes = du.B64.str_to_bytes(args)
    barrier, queue = pickle.loads(bytes)
    queue.put(barrier.wait())


def finite_barrier_wait(args):
    bytes = du.B64.str_to_bytes(args)
    barrier, queue = pickle.loads(bytes)
    try:
        barrier.wait(timeout=1)
    except BrokenBarrierError:
        queue.put(True)


def barrier_wait_and_put(args):
    bytes = du.B64.str_to_bytes(args)
    barrier, queue = pickle.loads(bytes)
    try:
        barrier.wait()
        queue.put(True)
    except BrokenBarrierError:
        queue.put(False)


def my_action():
    puid = current()
    with open("puid_file.txt", "a") as f:
        f.writelines(str(puid))


class TestBarrier(unittest.TestCase):
    def test_requirement_1_1(self):
        """Requirement 1.1: If the method is called with parties < 2, it has to
        raise a ValueError.
        """
        self.assertRaises(ValueError, lambda: Barrier(parties=0))
        self.assertRaises(ValueError, lambda: Barrier(parties=-1))

    def test_requirement_1_2(self):
        """Requirement 1.2: If the method is called with action not being either
        callable or None, it has to raise a ValueError
        """
        self.assertRaises(ValueError, lambda: Barrier(action=10))

    def test_requirement_1_3(self):
        """Requirement 1.3: If the method is called with timeout <0, it has to
        raise a ValueError.
        """
        self.assertRaises(ValueError, lambda: Barrier(timeout=-0.1))

    def test_requirement_2_0(self):
        """If the wait is called with a timeout, it takes precedent of the one
        that was supplied to __init__
        """
        barrier = Barrier(timeout=1, parties=2)
        start = time.monotonic()
        try:
            barrier.wait(timeout=0)
        except BrokenBarrierError:
            pass
        stop = time.monotonic()
        self.assertAlmostEqual(0, stop - start, 2)

    def test_requirement_2_1(self):
        """If the wait is called with timeout <0, it has to assume raise a ValueError."""
        self.assertRaises(ValueError, lambda: Barrier().wait(timeout=-0.1))

    def test_requirement_2_2(self):
        """
        If wait is called by less than parties processes within timeout seconds
        from the first process calling the method, all processes who have called
        it and will call it until reset() was called, have to raise a
        BrokenBarrierError immediately.  Note also Req 2.4.
        """
        barrier = Barrier(parties=2)
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "timeout_wait_barrier", env_str],
            None,
            options=options,
        )
        self.assertRaises(BrokenBarrierError, lambda: barrier.wait(timeout=0))
        join(process_1.p_uid)

    def test_requirement_2_3_a(self):
        """
        If wait is called with timeout=None, the calling process has to block
        until the method has been called by parties processes, or abort() or
        reset() has been called by any other process.
        """
        barrier = Barrier(parties=2)
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_barrier_forever", env_str],
            None,
            options=options,
        )
        process_2 = create(
            cmd,
            wdir,
            [__file__, "wait_barrier_forever", env_str],
            None,
            options=options,
        )
        multi_join([process_1.p_uid, process_2.p_uid])

    def test_requirement_2_3_b(self):
        """
        If wait is called with timeout=None, the calling process has to block
        until the method has been called by parties processes, or abort() or
        reset() has been called by any other process.
        """
        barrier = Barrier()
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_with_broken_barrier_helper", env_str],
            None,
            options=options,
        )
        process_2 = create(
            cmd,
            wdir,
            [__file__, "reset_barrier", env_str],
            None,
            options=options,
        )
        multi_join([process_1.p_uid, process_2.p_uid])

    def test_requirement_2_3_c(self):
        """
        If wait is called with timeout=None, the calling process has to block
        until the method has been called by parties processes, or abort() or
        reset() has been called by any other process.
        """
        barrier = Barrier()
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_with_broken_barrier_helper", env_str],
            None,
            options=options,
        )
        process_2 = create(
            cmd,
            wdir,
            [__file__, "abort_barrier", env_str],
            None,
            options=options,
        )
        multi_join([process_1.p_uid, process_2.p_uid])

    def test_requirement_2_4(self):
        """
        If wait() is called after abort() has been called by any process, raise
        a BrokenBarrierError until reset() is called.
        """
        barrier = Barrier()
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
                protocol=5,
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_with_broken_barrier_helper", env_str],
            None,
            options=options,
        )
        process_2 = create(
            cmd,
            wdir,
            [__file__, "abort_barrier", env_str],
            None,
            options=options,
        )
        multi_join([process_1.p_uid, process_2.p_uid])
        self.assertRaises(BrokenBarrierError, lambda: barrier.wait())
        self.assertIsNone(barrier.reset(), "The barrier did not reset.")

    @unittest.skipIf(
        os.path.dirname(os.path.realpath(__file__)) != os.getcwd(),
        f"Pickling callables works only if test is executed from the file dir.",
    )
    def test_requirement_2_5(self):
        """
        If wait has been called parties times by different processes/threads
        within timeout seconds, one process/thread waiting in the method has to
        call action if it is a callable. If action is None, do nothing.
        """

        if os.path.exists("puid_file.txt"):
            os.remove("puid_file.txt")

        barrier = Barrier(parties=2, action=my_action)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ],
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_barrier_forever", env_str],
            None,
            options=options,
        )
        process_2 = create(
            cmd,
            wdir,
            [__file__, "wait_barrier_forever", env_str],
            None,
            options=options,
        )

        multi_join([process_1.p_uid, process_2.p_uid])

        puid_val = None
        with open("puid_file.txt", "r") as f:
            lines = f.readlines()
            puid_val = int(lines[0])
        self.assertTrue(puid_val in [process_1.p_uid, process_2.p_uid])

        if os.path.exists("puid_file.txt"):
            os.remove("puid_file.txt")
        os.chdir("..")

    def test_requirement_2_6(self):
        """
        If wait has been called parties times within timeout seconds,
        it has to return a unique integer on every process in range 0<=i<parties.
        """
        barrier = Barrier(action=None, parties=5, timeout=None)
        queue = Queue(joinable=False)
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps([barrier, queue], protocol=5))
        puids = []
        for _ in range(5):
            process_1 = create(
                cmd,
                wdir,
                [__file__, "barrier_wait_helper", env_str],
                None,
                options=options,
            )
            puids.append(process_1.p_uid)
        multi_join(puids)
        queue_vals = []
        for _ in range(5):
            queue_vals.append(queue.get())
        queue_vals = sorted(queue_vals)
        self.assertListEqual(
            [0, 1, 2, 3, 4], queue_vals, "The queue did not receive unique integers between 0 and 5"
        )

    def test_requirement_3_1(self):
        """
        If the reset has been called while processes are waiting in wait(), they
        have to raise a BrokenBarrierError immediately.
        """
        barrier = Barrier()
        queue = Queue(joinable=False)
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps([barrier, queue], protocol=5))
        process_1 = create(
            cmd,
            wdir,
            [__file__, "finite_barrier_wait", env_str],
            None,
            options=options,
        )
        process_2 = create(
            cmd,
            wdir,
            [__file__, "reset_barrier", env_str],
            None,
            options=options,
        )
        self.assertTrue(queue.get(), "The BrokenBarrierException has not occurred.")
        multi_join([process_1.p_uid, process_2.p_uid])

    def test_requirement_3_2(self):
        """
        After the reset has been called, processes calling wait() will not raise
        a BrokenBarrierError anymore. (Duplicate to 2.2 & 2.4)
        """
        barrier = Barrier(parties=5, timeout=None)
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(
            pickle.dumps(
                [
                    barrier,
                ]
            )
        )
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_with_broken_barrier_helper", env_str],
            None,
            options=options,
        )
        process_2 = create(
            cmd,
            wdir,
            [__file__, "abort_barrier", env_str],
            None,
            options=options,
        )
        multi_join([process_1.p_uid, process_2.p_uid])

        barrier.reset()

        queue = Queue(joinable=False)

        env_str = du.B64.bytes_to_str(pickle.dumps([barrier, queue]))
        process_4 = create(
            cmd,
            wdir,
            [__file__, "barrier_wait_and_put", env_str],
            None,
            options=options,
        )
        process_5 = create(
            cmd,
            wdir,
            [__file__, "barrier_wait_and_put", env_str],
            None,
            options=options,
        )
        multi_join([process_4.p_uid, process_5.p_uid])

        self.assertTrue(queue.get(), "The first barrier wait did not complete successfully.")
        self.assertTrue(queue.get(), "The second barrier wait did not complete successfully.")

    def test_requirements_5_6_7(self):
        """
        The class needs to implement an attribute Barrier.parties : int that is
        > 0 and holds the parties attribute from the __init__ method.

        The class needs to implement an attribute Barrier.n_waiting : int that
        is >= 0 and the number of waiting processes/threads.

        The class needs to implement an attribute Barrier.broken : int
        """
        barrier = Barrier()
        self.assertIsInstance(barrier.parties, int, "The parties are not an int.")
        self.assertIsInstance(barrier.n_waiting, int, "The number of waiting threads are not an int.")
        self.assertIsInstance(barrier.broken, int, "The broken barrier error is not returning an int")


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "wait_barrier_forever":
        wait_barrier_forever(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "barrier_wait_helper":
        barrier_wait_helper(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "timeout_wait_barrier":
        timeout_wait_barrier(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "finite_barrier_wait":
        finite_barrier_wait(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "barrier_wait_and_put":
        finite_barrier_wait(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "reset_barrier":
        reset_barrier(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "abort_barrier":
        abort_barrier(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "wait_with_broken_barrier_helper":
        wait_with_broken_barrier_helper(sys.argv[2])
    else:
        unittest.main()

test/native/test_barrier.py
DragonHPC/dragon
import unittest
from threading import Thread
import time
import pickle
import sys

import dragon
import dragon.native.event
from dragon.globalservices.process import (
    join,
    create,
    multi_join,
)
from dragon.infrastructure.process_desc import ProcessOptions
import dragon.utils as du


def wait_event_with_timeout(args):
    bytes = du.B64.str_to_bytes(args)
    event, timeout = pickle.loads(bytes)
    wait_val = event.wait(timeout=timeout)
    assert wait_val == True, "The process wait did not poll a message"


def set_event(event):
    bytes = du.B64.str_to_bytes(event)
    event = pickle.loads(bytes)
    event.set()
    assert event.is_set() == True, "The event is not set"


class TestEvent(unittest.TestCase):
    @staticmethod
    def set_event(event):
        event.set()
        assert event.is_set() == True, "The event is not set"

    @staticmethod
    def wait_event_with_timeout(event, timeout):
        wait_val = event.wait(timeout=timeout)
        assert wait_val == True, "The thread wait did not poll a message"

    @staticmethod
    def thread_forloop_wait_event_with_timeout(event, timeout):
        time.sleep(0.01)
        wait_val = event.wait(timeout=timeout)
        assert wait_val == True, "The thread wait did not poll a message"

    def test_process_event_none_timeout(self):
        """Requirement 1.1 for processes: If the method is called with timeout=None,
        the calling process blocks until another process calls set() on the same event class.
        The calling process then unblocks and the class return True.
        """
        event = dragon.native.event.Event()
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps([event, None], protocol=5))
        event.set()
        process_1 = create(
            cmd,
            wdir,
            [__file__, "wait_event_with_timeout", env_str],
            None,
            options=options,
        )
        env_str_2 = du.B64.bytes_to_str(pickle.dumps(event, protocol=5))
        process_2 = create(
            cmd,
            wdir,
            [__file__, "set_event", env_str_2],
            None,
            options=options,
        )
        join(process_1.p_uid, timeout=None)
        join(process_2.p_uid, timeout=None)
        self.assertTrue(event.wait(), "The event is not set by the second process.")

    def test_thread_event_none_timeout(self):
        """Requirement 1.1 for threads."""
        event = dragon.native.event.Event()
        t1 = Thread(target=self.wait_event_with_timeout, args=(event, None))
        t2 = Thread(target=self.set_event, args=(event,))
        event.set()
        t1.start()
        t2.start()
        t1.join()
        t2.join()
        self.assertTrue(event.wait(), "The event is not set by the second process.")

    def test_process_event_postive_timeout(self):
        """Requirement 1.2 for processes: If the method is called with timeout > 0 and no other process has called set, the calling process blocks until
        timeout number of seconds have elapsed,
        The class then returns False.
        """
        event = dragon.native.event.Event()
        start = time.monotonic()
        event_wait = event.wait(0.1)
        stop = time.monotonic()
        timeout = round(float(float(stop) - float(start)), 1)
        self.assertAlmostEqual(0.1, timeout, "The timeout is not 0.1")
        # This proves the event was not set and wait returns False as a result.
        self.assertFalse(event_wait, "The event is set")
        self.assertFalse(event.is_set(), "The event is set")

    def test_process_wait_event_with_timeout(self):
        """Requirement 1.3 for processes: The calling process does not block on a timeout of 0, but returns False immediately if set() has not been called before.
        It returns True immediately, if set() has been called before by any process.
        """
        event = dragon.native.event.Event()
        start = time.monotonic()
        wait_val = event.wait(timeout=0)
        stop = time.monotonic()
        timeout = int(int(stop) - int(start))
        self.assertAlmostEqual(0, timeout)
        self.assertFalse(wait_val, "The event is set")
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(event, protocol=5))
        the_descr = create(
            cmd,
            wdir,
            [__file__, "set_event", env_str],
            None,
            options=options,
        )
        join(the_descr.p_uid)
        self.assertEqual(event.is_set(), event.wait(), "The event is not set")

    def test_thread_wait_event_with_timeout(self):
        """Requirement 1.3 for threads."""
        event = dragon.native.event.Event()
        start = time.monotonic()
        wait_val = event.wait(timeout=0)
        stop = time.monotonic()
        timeout = int(int(stop) - int(start))
        self.assertAlmostEqual(0, timeout)
        self.assertFalse(wait_val, "The event is set")
        t1 = Thread(target=self.set_event, args=(event,))
        t1.start()
        t1.join()
        self.assertEqual(event.is_set(), event.wait(), "The event is not set")

    def test_process_event_negative_wait_timeout(self):
        """Requirement 1.4 for process: If the method is called with timeout<0 ,
        it has to behave as if it was called with timeout=0.
        """
        event = dragon.native.event.Event()
        start = time.monotonic()
        wait_val = event.wait(timeout=-1)
        stop = time.monotonic()
        timeout = int(int(stop) - int(start))
        self.assertAlmostEqual(0, timeout)
        self.assertFalse(wait_val, "The event is set")

    def test_process_multiple_sets(self):
        """Requirement 1.5 for processes: If the method is called after set()
        has been called by any process, it always returns True immediately.
        """
        event = dragon.native.event.Event()
        event.set()
        self.assertTrue(event.is_set(), "The event was not set")
        wait_val = event.wait()
        self.assertTrue(wait_val, "The event was not set")

    def test_process_multiple_wait_set(self):
        """Requirement 2.1 for processes: If the method is called, all processes
        blocking in a wait() call return True.
        """
        event = dragon.native.event.Event()
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps([event, None]))
        puids = []
        for _ in range(5):
            the_descr = create(
                cmd,
                wdir,
                [__file__, "wait_event_with_timeout", env_str],
                None,
                options=options,
            )
            puids.append(the_descr.p_uid)

        event.set()

        ready = multi_join(puids, timeout=10, join_all=True)

        for puid, ret_code in ready[0]:
            self.assertTrue(puid in puids)
            self.assertTrue(ret_code == 0)

    def test_thread_multiple_wait_set(self):
        """Requirement 2.1 for threads."""

        event = dragon.native.event.Event()

        threads = []
        for _ in range(5):
            thread = Thread(target=self.thread_forloop_wait_event_with_timeout, args=(event, None))
            threads.append(thread)
            thread.start()

        event.set()

        for thread in threads:
            join_val = thread.join()
            self.assertIsNone(join_val)

    def test_process_varied_timeouts(self):
        """Requirement 2.2 for processes: After the method was called, all
        processes calling wait() in the future return True immediately,
        regardless if the value of timeout was timeout=None, timeout<0,
        timeout=0 or timeout>0.
        """
        event = dragon.native.event.Event()
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(event, protocol=5))
        puids = []
        timeouts = [-1, 1, 0, None]
        for timeout in timeouts:
            env_str = du.B64.bytes_to_str(pickle.dumps([event, timeout], protocol=5))
            the_descr = create(
                cmd,
                wdir,
                [__file__, "wait_event_with_timeout", env_str],
                None,
                options=options,
            )
            puids.append(the_descr.p_uid)
        event.set()
        ready = multi_join(puids, timeout=None, join_all=True)

        for puid, ret_code in ready[0]:
            self.assertTrue(puid in puids)
            self.assertTrue(ret_code == 0)

    def test_thread_varied_timeouts(self):
        """Requirement 2.2 for threads."""
        event = dragon.native.event.Event()
        threads = []
        timeouts = [None, -1, 0, 0]
        for timeout in timeouts:
            thread = Thread(target=self.thread_forloop_wait_event_with_timeout, args=(event, timeout))
            threads.append(thread)
            thread.start()
        event.set()
        for thread in threads:
            join_val = thread.join()
            self.assertIsNone(join_val, "The thread did not join properly.")
        self.assertTrue(event.is_set(), "The event is not registering as set by the main thread.")

    def test_set(self):
        """Requirement 3.1: The method must return True only if set() was called
        before on the same Event by the same process or any other process.
        """

        event = dragon.native.event.Event()
        event.set()
        self.assertTrue(event.is_set())

    def test_set_false(self):
        """Requirement 3.2: The method must return False only if set() was not
        called before on the same Event by the same or any other process.
        """

        event = dragon.native.event.Event()
        self.assertFalse(event.is_set())

    def test_set_clear(self):
        """Requirement 4.1: If the method is called after a process has call
        set(), the result is that the event class behaves as if set() has not
        been called.
        """

        event = dragon.native.event.Event()

        event.set()
        self.assertTrue(event.is_set())

        event.clear()
        self.assertFalse(event.is_set())

    def test_clear_no_set(self):
        """Requirement 4.2: If the method is called and set() has not been
        called before by any processes, it is a no-op."""

        event = dragon.native.event.Event()
        event.clear()
        self.assertFalse(event.is_set())


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "set_event":
        set_event(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "wait_event_with_timeout":
        wait_event_with_timeout(sys.argv[2])
    else:
        unittest.main()

test/native/test_event.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import time
import pickle
import sys

import dragon
from dragon.globalservices.process import (
    multi_join,
    create,
    get_list,
    join,
)
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.infrastructure.parameters import this_process
from dragon.native.lock import Lock
import dragon.utils as du


def _burner(env_str):

    num_acq = 1024

    # unpack the lock from the environment string
    bytes = du.B64.str_to_bytes(env_str)
    lock = pickle.loads(bytes)

    for _ in range(num_acq):
        success = lock.acquire(block=True, timeout=None)
        assert success == True, "Could not acquire lock"
        lock.release()


def _releaser(env_str):

    # unpack the lock from the environment string
    bytes = du.B64.str_to_bytes(env_str)
    rlock = pickle.loads(bytes)

    # try to release it
    try:
        rlock.release()
    except AssertionError:
        pass
    else:
        assert False, "Release should not have been successful"


def _acquirer(env_str):

    # unpack the lock from the environment string
    bytes = du.B64.str_to_bytes(env_str)
    lock = pickle.loads(bytes)

    lock.acquire(block=True, timeout=None)
    assert lock.is_mine(), "is_mine() broken"


class TestLock(unittest.TestCase):
    def test_lock_basic(self):
        """Test that processes actually wait on the lock and can be released
        by any other process.
        """

        # get a lock
        lock = Lock(recursive=False)
        self.assertTrue(lock.acquire())  # head process owns the lock now
        self.assertFalse(lock.acquire(timeout=0))  # cannot acquire again

        # start some processes
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(lock, protocol=5))

        puids = []
        for _ in range(2):
            the_descr = create(
                cmd,
                wdir,
                [__file__, "_acquirer", env_str],
                None,
                options=options,
            )
            puids.append(the_descr.p_uid)

        # all should be waiting, we still own lock
        live_puids = get_list()
        live_puids.remove(this_process.my_puid)

        for puid in puids:
            self.assertTrue(puid in live_puids)

        # release one
        lock.release()
        ready = multi_join(puids, timeout=None)
        self.assertTrue(len(ready[0]) == 1)  # first has exited
        puids.remove(ready[0][0][0])  # wow ...

        # release two
        lock.release()
        ready = multi_join(puids, timeout=None, join_all=True)
        self.assertTrue(len(ready[0]) == 1)  # rest has exited

        # cause a value error
        lock.release()  # one release left ...
        self.assertRaises(ValueError, lock.release)

    def test_recursive(self):
        """Test that recursive locks can be acquired multiple times and
        need to be released multiple times by the same process
        """

        # get a recursive lock
        rlock = Lock(recursive=True)

        # acquire it 25 times
        for _ in range(25):
            self.assertTrue(rlock.acquire())
        self.assertTrue(rlock._accesscount == 25)

        # start a processes that tries to release it
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(rlock, protocol=5))

        the_descr = create(
            cmd,
            wdir,
            [__file__, "_releaser", env_str],
            None,
            options=options,
        )
        puid = the_descr.p_uid

        # wait until it has tried, counter must still be 25
        ready = multi_join([puid], timeout=None, join_all=True)

        # release 25 times
        for _ in range(25):
            rlock.release()  # cannot raise an exception
        self.assertTrue(rlock._accesscount == 0)

        # counter is now 0
        self.assertRaises(AssertionError, rlock.release)

    def test_rlock_ownership(self):
        """Test that only the calling thread owns the lock."""

        # get a recursive lock
        rlock = Lock(recursive=True)

        # acquire it
        self.assertTrue(rlock.acquire())
        self.assertTrue(rlock.is_mine())

        # still mine after pickle ?
        bytes = pickle.dumps(rlock)
        rlock2 = pickle.loads(bytes)
        self.assertTrue(rlock2.is_mine())

        # release it
        rlock.release()
        self.assertFalse(rlock.is_mine())

        # still not mine after pickle ?
        bytes = pickle.dumps(rlock)
        rlock3 = pickle.loads(bytes)
        self.assertFalse(rlock3.is_mine())

        # start a processes that acquires the lock and leaves
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(rlock, protocol=5))

        the_descr = create(
            cmd,
            wdir,
            [__file__, "_acquirer", env_str],
            None,
            options=options,
        )
        puid = the_descr.p_uid

        ready = join(puid, timeout=None)
        self.assertTrue(ready == 0)  # process has exited

        # should not be mine and I cannot release it
        self.assertFalse(rlock.is_mine())
        self.assertRaises(AssertionError, rlock.release)

    def test_lock_contention(self):
        """Test lock correctness under contention."""

        # get a recursive lock
        lock = Lock(recursive=False)

        num_workers = 5

        # create env string
        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(lock, protocol=5))

        # start processes
        puids = []
        for _ in range(num_workers):
            the_descr = create(
                cmd,
                wdir,
                [__file__, "_burner", env_str],
                None,
                options=options,
            )
            puids.append(the_descr.p_uid)

        # let them finish
        ready = multi_join(puids, timeout=None, join_all=True)

        # check they all have exited without error
        for puid, exit_code in ready[0]:
            self.assertTrue(exit_code == 0)


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "_acquirer":
        _acquirer(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "_releaser":
        _releaser(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "_burner":
        _burner(sys.argv[2])
    else:
        unittest.main()

test/native/test_lock.py
DragonHPC/dragon
import unittest
from dragon.native.logging import log_to_dragon, log_to_stderr


@unittest.skip(f"Dragon native logging interface not implemented yet (PE-41692)")
class TestDragonNativeMisc(unittest.TestCase):
    def test_get_logger():
        pass

    def test_log_to_stderr():
        pass


if __name__ == "__main__":
    unittest.main()

test/native/test_logging.py
DragonHPC/dragon
import unittest
import os
import dragon

from dragon.utils import host_id


class TestMachineSingle(unittest.TestCase):
    def test_multinode_cpu_count(self):
        """Test the number of logical CPUs in the single node case"""

        the_cpu_count = os.cpu_count()

        dragons_cpu_count = dragon.native.machine.cpu_count()

        # Since we don't have multiple nodes, the count should equal our CPU count
        self.assertTrue(dragons_cpu_count == the_cpu_count)

    def test_node_list(self):
        """Test that we can get a list of the registered nodes"""

        hlist = dragon.globalservices.node.get_list()
        self.assertEqual(1, len(hlist))

    def test_node_class(self):

        mynode = dragon.native.machine.current()
        self.assertTrue(mynode.h_uid == host_id())


if __name__ == "__main__":
    unittest.main()

test/native/test_machine.py
DragonHPC/dragon
import unittest
import time

import dragon  # DRAGON import before multiprocessing
from dragon.native.pool import Pool

def sqr(x, wait=0.0):
    time.sleep(wait)
    return x * x

def raising():
    raise KeyError("key")

class TestDragonNativePool(unittest.TestCase):
    """Unit tests for the Dragon Native Pool.
    """
    def test_async(self):

        pool = Pool(processes=4)
    
        test_timeout=1.2 
     
        res = pool.apply_async(sqr, (7, test_timeout,))
   
        start = time.monotonic()
        result = res.get() 
        stop = time.monotonic()
  
        self.assertEqual(result, 49)
        self.assertAlmostEqual(stop-start , test_timeout, 0)
      
        pool.close()
        pool.join()

    def test_async_timeout(self):
        
        pool = Pool(processes=4)
    
        test_timeout=1.2 
        res = pool.apply_async(sqr, (6, test_timeout + 1.0))
        get = res.get
        self.assertRaises(TimeoutError, get, timeout=test_timeout)
        
        pool.close()
        pool.join()

    def test_map_async(self):
       
        pool = Pool(processes=4)
        
        self.assertEqual(pool.map_async(sqr, list(range(10))).get(), list(map(sqr, list(range(10)))))
        self.assertEqual(pool.map_async(sqr, list(range(100)), chunksize=20).get(), list(map(sqr, list(range(100)))))
        
        pool.close()
        pool.join()

    # if max value of range is huge then we can wait in main thread for a long time before returning
    def test_terminate(self):
        pool = Pool(processes=4)
        _ = pool.map_async(time.sleep, [0.1 for i in range(100)], chunksize=1)
        pool.terminate()
        start = time.monotonic()
        pool.join()
        stop = time.monotonic()
  
        # Sanity check the pool didn't wait for all work to be done 
        self.assertLess(stop-start, 2.0)

    def test_empty_iterable(self):
        p = Pool(1)
        self.assertEqual(p.map_async(sqr, []).get(), [])
        p.close()
        p.join()


    def test_async_error_callback(self):
        p = Pool(2)

        scratchpad = [None]

        def errback(exc):
            scratchpad[0] = exc

        res = p.apply_async(raising, error_callback=errback)
        self.assertRaises(KeyError, res.get)
        self.assertTrue(scratchpad[0])
        self.assertIsInstance(scratchpad[0], KeyError)

        p.close()
        p.join()


if __name__ == "__main__":
    unittest.main()

test/native/test_pool.py
DragonHPC/dragon
import unittest
import sys
import time

import dragon

from dragon.native.queue import Queue
from dragon.native.process import Process, current, TemplateProcess
from dragon.infrastructure.parameters import this_process


class TestDragonNativeProcess(unittest.TestCase):
    def test_basic(self):

        exe = "sleep"
        args = ("10000",)

        p = Process(exe, args, ident="Apple")  # creates a new process
        p.start()
        self.assertTrue(p.is_alive)
        self.assertTrue(p.ident == "Apple")

        p2 = Process(None, None, ident="Apple")
        self.assertTrue(p.name == p2.name)
        self.assertTrue(p.puid == p2.puid)

        pp = p.parent()
        self.assertTrue(pp.puid == this_process.my_puid)
        me = current()
        self.assertTrue(me.puid == this_process.my_puid)

        p.kill()

        p2.join(timeout=None)
        self.assertTrue(p.is_alive == False)
        self.assertTrue(p2.is_alive == False)

        self.assertRaises(AttributeError, Process, "gobbledygook")

    def test_subclassing(self):
        class UserProcess(Process):
            pass

        self.assertRaises(NotImplementedError, UserProcess, "sleep")

    @classmethod
    def putter(cls, q):
        q.put(True)
        _ = q.get()

    @classmethod
    def raiser(cls):
        raise Exception("Bad Function - ignore")

    @classmethod
    def exiter(cls):
        sys.exit(42)

    @classmethod
    def sleeper(cls):
        time.sleep(1000000)

    def test_basic_python(self):

        q = Queue()

        pyproc = Process(self.putter, args=(q,), ident="Pear")
        pyproc.start()

        item = q.get()
        self.assertTrue(item == True)
        self.assertTrue(pyproc.is_alive == True)

        q.put(True)

        pyproc.join()

        self.assertTrue(pyproc.is_alive == False)
        self.assertTrue(pyproc.returncode == 0)

    def test_exception_handling(self):

        pyproc = Process(self.raiser)
        pyproc.start()
        pyproc.join()
        self.assertTrue(pyproc.returncode == 1)

    def test_exit_handling(self):

        pyproc = Process(self.exiter)
        pyproc.start()
        pyproc.join()
        self.assertTrue(pyproc.returncode == 42)

    def test_kill(self):

        pyproc = Process(self.sleeper)
        pyproc.start()
        self.assertTrue(pyproc.is_alive == True)
        pyproc.kill()
        pyproc.join(timeout=None)

        pyproc = Process(self.exiter)
        pyproc.start()
        pyproc.join()
        self.assertTrue(pyproc.returncode == 42)

    def test_templating_basic(self):

        exe = "sleep"
        args = ("10000",)

        templ = TemplateProcess(exe, args)

        p = Process.from_template(templ, ident="Banana")  # creates a new process

        p.start()
        self.assertTrue(p.is_alive)
        self.assertTrue(p.ident == "Banana")
        p.kill()
        p.join()

    def test_templating_python(self):

        q = Queue()

        templ = TemplateProcess(self.putter, args=(q,))
        func, args, kwargs = templ.get_original_python_parameters()
        self.assertTrue(callable(func))
        self.assertIsInstance(args[0], Queue)
        self.assertIsInstance(args, tuple)
        self.assertTrue(kwargs == {})

        self.assertTrue(templ.is_python)

        pyproc = Process.from_template(templ, ident="Pineapple")

        pyproc.start()

        item = q.get()
        self.assertTrue(item == True)
        self.assertTrue(pyproc.is_alive == True)

        q.put(True)

        pyproc.join()

        self.assertTrue(pyproc.is_alive == False)
        self.assertTrue(pyproc.returncode == 0)


if __name__ == "__main__":
    unittest.main()

test/native/test_process.py
DragonHPC/dragon
import unittest
import time
import os
import sys
import random
import signal

import dragon
from dragon.globalservices.process import ProcessError, query, kill, signal as dragon_signal
from dragon.infrastructure.process_desc import ProcessOptions

from dragon.native.process import Process, TemplateProcess
from dragon.native.event import Event
from dragon.native.queue import Queue
from dragon.native.process_group import ProcessGroup, DragonProcessGroupError

# we have to test every transition in the state diagram here.
# plus a few robustness tests


class TestDragonNativeProcessGroup(unittest.TestCase):
    """Unit tests for the Dragon Native ProcessGroup Manager.
    This class has to be very robust against processes disappearing.
    """

    @classmethod
    def setUpClass(cls):
        """Set parameters"""

        # a few standard values
        cls.nproc = 3
        cls.cmd = "sleep"
        cls.args = ("128",)
        cls.cwd = os.getcwd()
        cls.options = ProcessOptions(make_inf_channels=True)
        cls.template = TemplateProcess(cls.cmd, args=cls.args, cwd=cls.cwd)

    def test_init(self):

        pg = ProcessGroup()
        pg.add_process(self.nproc, self.template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        man_proc = pg._manager._proc

        self.assertTrue(man_proc.is_alive)

        pg.stop()

        man_proc.join(timeout=None)
        self.assertFalse(man_proc.is_alive)

    def test_start_stop(self):

        pg = ProcessGroup()
        pg.add_process(self.nproc, self.template)
        pg.init()

        man_proc = pg._manager._proc
        self.assertTrue(man_proc.is_alive)
        self.assertTrue(pg.status == "Idle")

        pg.start()

        self.assertTrue(pg.status == "Maintain")

        puids = pg.puids
        processes = [Process(None, ident=puid) for puid in puids]

        for p in processes:
            self.assertTrue(p.is_alive)

        pg.stop()

        man_proc.join(timeout=None)  # will the manager exit after stop ?
        self.assertFalse(man_proc.is_alive)

    @classmethod
    def event_quitter(cls, ev):
        ev.wait()

    def test_join_from_idle(self):

        ev = Event()

        template = TemplateProcess(self.event_quitter, args=(ev,), cwd=".")
        pg = ProcessGroup(restart=False)
        pg.add_process(self.nproc, template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        pg.start()

        self.assertTrue(pg.status == "Running")

        puids = pg.puids
        processes = [Process(None, ident=puid) for puid in puids]

        for p in processes:
            self.assertTrue(p.is_alive)

        ev.set()

        for p in processes:
            p.join()

        for p in processes:
            self.assertFalse(p.is_alive)

        while not pg.status == "Idle":
            self.assertTrue(pg.status == "Running")
            time.sleep(0.1)  # avoid race condition

        puids = pg.puids
        for puid in puids:
            self.assertIsNone(puid)

        pg.stop()

    def test_join_from_maintain(self):
        # test that join indeed moves to idle when the processes exit

        ev = Event()

        template = TemplateProcess(self.event_quitter, args=(ev,), cwd=".")
        pg = ProcessGroup(restart=True)
        pg.add_process(self.nproc, template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        pg.start()

        self.assertTrue(pg.status == "Maintain")

        puids = pg.puids
        processes = [Process(None, ident=puid) for puid in puids]
        self.assertRaises(TimeoutError, pg.join, 0)
        self.assertTrue(pg.status == "Running")

        for p in processes:
            self.assertTrue(p.is_alive)

        ev.set()

        for p in processes:
            p.join()

        # test autotransition
        while not pg.status == "Idle":
            self.assertTrue(pg.status == "Running")
            time.sleep(0.1)  # avoid race condition

        puids = pg.puids
        for puid in puids:
            self.assertIsNone(puid)
        pg.stop()

    @unittest.skip("This one is very slow for reasons I don't understand.")
    def test_join_with_timeout(self):

        ev = Event()

        template = TemplateProcess(self.event_quitter, args=(ev,), cwd=".")
        pg = ProcessGroup(restart=True)
        pg.add_process(3, self.template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        pg.start()

        self.assertTrue(pg.status == "Maintain")

        puids = pg.puids
        processes = [Process(None, ident=puid) for puid in puids]

        self.assertRaises(TimeoutError, pg.join, 0)

        self.assertTrue(pg.status == "Running")

        start = time.monotonic()
        self.assertRaises(TimeoutError, pg.join, 1)
        stop = time.monotonic()

        self.assertAlmostEqual(stop - start, 1, 0)

        ev.set()

        for p in processes:
            p.join()  # processes won't exit. Scheduling/spin wait issue ?

        start = time.monotonic()
        pg.join(timeout=None)
        stop = time.monotonic()
        self.assertAlmostEqual(stop - start, 0, 0)

        self.assertTrue(pg.status == "Idle")

        pg.stop()

    def test_stop_from_idle(self):

        pg = ProcessGroup()
        pg.add_process(self.nproc, self.template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        man_proc = pg._manager._proc

        self.assertTrue(man_proc.is_alive)

        pg.stop()

        man_proc.join(timeout=None)  # will the manager exit after stop ?

        self.assertFalse(man_proc.is_alive)

    @classmethod
    def _putters(cls, q):

        termflag = False

        def handler(signum, frame):
            nonlocal termflag
            termflag = True

        signal.signal(signal.SIGTERM, handler)

        q.put(True)

        while not termflag:
            time.sleep(1)

    def test_shutdown_from_maintain(self):

        q = Queue()

        template = TemplateProcess(self._putters, args=(q,), cwd=".")
        pg = ProcessGroup()
        pg.add_process(self.nproc, template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        pg.start()

        self.assertTrue(pg.status == "Maintain")

        puids = pg.puids

        for _ in puids:
            __ = q.get()  # make sure they've all executed code

        pg.kill(signal.SIGTERM)

        while not pg.status == "Idle":
            self.assertTrue(pg.status == "Running")
            time.sleep(0.1)

        for puid in pg.puids:
            self.assertIsNone(puid)

        for puid in puids:  # check if really dead
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

    def test_shutdown_from_join(self):

        q = Queue()

        template = TemplateProcess(self._putters, args=(q,), cwd=".")
        pg = ProcessGroup(restart=False)
        pg.add_process(self.nproc, template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        pg.start()

        self.assertTrue(pg.status == "Running")

        puids = pg.puids

        pg.kill(signal.SIGTERM)

        while not pg.status == "Idle":
            self.assertTrue(pg.status == "Running")
            time.sleep(0.1)

        for puid in pg.puids:
            self.assertIsNone(puid)

        for puid in puids:  # check if really dead
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

    def test_kill_from_maintain(self):

        pg = ProcessGroup()
        pg.add_process(self.nproc, self.template)
        pg.init()

        self.assertTrue(pg.status == "Idle")
        man_proc = pg._manager._proc
        self.assertTrue(man_proc.is_alive)

        pg.start()

        self.assertTrue(pg.status == "Maintain")

        puids = pg.puids

        pg.kill()

        self.assertTrue(pg.status == "Idle")

        for puid in pg.puids:
            self.assertIsNone(puid)

        for puid in puids:
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

    def test_kill_from_join(self):

        pg = ProcessGroup(restart=False)
        pg.add_process(self.nproc, self.template)
        pg.init()

        self.assertTrue(pg.status == "Idle")

        pg.start()

        self.assertTrue(pg.status == "Running")

        puids = pg.puids

        for puid in puids:
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.ACTIVE)

        pg.kill()

        self.assertTrue(pg.status == "Idle")

        for puid in puids:
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

    @classmethod
    def _failer(cls, ev):

        ev.wait()

        raise Exception("42 - ignore")

    def test_no_error_from_maintain(self):

        ev = Event()

        template = TemplateProcess(self._failer, args=(ev,), cwd=".")
        pg = ProcessGroup(restart=True)
        pg.add_process(self.nproc, template)
        pg.init()
        pg.start()
        self.assertTrue(pg.status == "Maintain")

        puids = pg.puids

        ev.set()

        cnt = 0
        while not pg.status == "Error":  # maintain should catch failing processes
            self.assertTrue(pg.status == "Maintain")
            time.sleep(0.1)
            cnt += 1
            if cnt == 32:  # good enough
                break
        else:
            cnt = -1

        self.assertTrue(cnt > -1)
        self.assertTrue(pg.status == "Maintain")

        pg.kill()

        while not pg.status == "Idle":
            self.assertTrue(pg.status == "Error")
            time.sleep(0.1)

        for puid in puids:
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

    def test_error_and_kill_from_join(self):

        ev = Event()

        template = TemplateProcess(self._failer, args=(ev,), cwd=".")
        pg = ProcessGroup(restart=False)
        pg.add_process(self.nproc, template)
        pg.init()
        pg.start()
        self.assertTrue(pg.status == "Running")

        puids = pg.puids

        ev.set()

        while not pg.status == "Error":
            self.assertTrue(pg.status == "Running")
            time.sleep(0.1)

        self.assertTrue(pg.status == "Error")

        pg.kill()

        while not pg.status == "Idle":
            self.assertTrue(pg.status == "Error")
            time.sleep(0.05)

        for puid in puids:
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

    def test_error_and_stop_from_join(self):

        ev = Event()

        template = TemplateProcess(self._failer, args=(ev,), cwd=".")
        pg = ProcessGroup(restart=False)
        pg.add_process(self.nproc, template)
        pg.init()

        man_proc = pg._manager._proc
        self.assertTrue(man_proc.is_alive)

        pg.start()
        self.assertTrue(pg.status == "Running")

        puids = pg.puids

        ev.set()

        while not pg.status == "Error":
            self.assertTrue(pg.status == "Running")
            time.sleep(0.1)

        self.assertTrue(pg.status == "Error")

        pg.stop()

        while not pg.status == "Stop":
            self.assertTrue(pg.status == "Error")
            time.sleep(0.05)

        for puid in puids:
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        man_proc.join()
        self.assertFalse(man_proc.is_alive)

    def test_ignore_error_during_shutdown(self):
        pg = ProcessGroup(ignore_error_on_exit=True)
        pg.add_process(self.nproc, self.template)
        pg.init()

        pg.start()
        self.assertTrue(pg.status == "Maintain")

        puids = pg.puids
        puids.reverse()

        pg.kill(signal=signal.SIGTERM)  # return immediately

        for puid in puids:  # force kill all processes
            try:
                kill(puid, sig=dragon_signal.SIGKILL)
            except Exception as e:
                pass

        while not pg.status == "Idle":
            self.assertFalse(pg.status == "Error")

        pg.stop()

    def test_bad_transitions(self):

        pg = ProcessGroup()
        pg.add_process(self.nproc, self.template)
        pg.init()

        self.assertTrue(pg.status == "Idle")
        self.assertRaises(DragonProcessGroupError, pg.kill, signal.SIGTERM)
        self.assertRaises(DragonProcessGroupError, pg.kill)

        pg.start()
        self.assertTrue(pg.status == "Maintain")
        self.assertRaises(DragonProcessGroupError, pg.start)

        self.assertRaises(TimeoutError, pg.join, 0)
        self.assertTrue(pg.status == "Running")
        self.assertRaises(DragonProcessGroupError, pg.start)
        self.assertRaises(DragonProcessGroupError, pg.stop)

        pg.kill()
        pg.stop()
        self.assertTrue(pg.status == "Stop")
        self.assertRaises(DragonProcessGroupError, pg.start)
        self.assertRaises(DragonProcessGroupError, pg.stop)
        self.assertRaises(DragonProcessGroupError, pg.kill)
        self.assertRaises(DragonProcessGroupError, pg.kill, signal.SIGTERM)


    @unittest.skip("Pending resolution on CIRRUS-1455")
    def test_maintain_stress(self):

        # we will keep killing processes that keep exiting.
        # the class has to maintain them and GS has to handle all of that.

        testtime = 3  # sec

        template = TemplateProcess("sleep", args=(f"{round(testtime/3)}",), cwd=".")

        pg = ProcessGroup()
        pg.add_process(64, template)
        pg.init()
        man_proc = pg._manager._proc
        self._update_interval_sec = 0.1

        pg.start()

        beg = time.monotonic()

        while time.monotonic() - beg < testtime:
            puids = pg.puids

            try:
                puid = puids[random.randint(0, len(puids))]
                kill(puid, sig=dragon_signal.SIGKILL)
            except Exception:  # maybe it disappeared already
                pass

            self.assertTrue(man_proc.is_alive)
            self.assertTrue(pg.status != "Error")

            time.sleep(0.2)

        puids = pg.puids

        pg.kill(signal=signal.SIGTERM)

        while not pg.status == "Idle":
            self.assertTrue(pg.status != "Error")

        self.assertTrue(man_proc.is_alive)

        for puid in puids:
            gs_info = query(puid)
            self.assertTrue(gs_info.state == gs_info.State.DEAD)

        pg.stop()

        man_proc.join()

    def test_walltime(self):

        wtime = 3

        pg = ProcessGroup(walltime=wtime)
        pg.add_process(3, self.template)
        pg.init()

        pg.start()
        start = time.monotonic()

        while not pg.status == "Idle":
            self.assertFalse(pg.status == "Error")
            time.sleep(0.5)

        stop = time.monotonic()

        self.assertAlmostEqual(stop - start, wtime, 0)

    @unittest.skip("Waiting for group implementation")
    def test_pmi_enabled(self):
        pass

    @unittest.skip("Waiting for group implementation")
    def test_multiple_groups(self):
        pass

    @unittest.skip("Waiting for group implementation")
    def test_multiple_clients(self):
        pass

    @unittest.skip("Waiting for group implementation")
    def test_multiple_groups_and_multiple_clients(self):
        pass


if __name__ == "__main__":
    unittest.main()

test/native/test_process_group.py
DragonHPC/dragon
import unittest
import sys
import pickle
import time

import dragon
from dragon.native.queue import Queue

from dragon.channels import Channel, MASQUERADE_AS_REMOTE
from dragon.managed_memory import MemoryPool

from dragon.globalservices.process import create as process_create, join as process_join
from dragon.globalservices.pool import query as pool_query, get_list as pool_get_list
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.infrastructure.channel_desc import ChannelOptions
from dragon.infrastructure.facts import default_pool_muid_from_index
import dragon.utils as du


def _putter(env_str):

    bytes = du.B64.str_to_bytes(env_str)
    q, item = pickle.loads(bytes)

    q.put(item)


def _joiner(env_str):

    bytes = du.B64.str_to_bytes(env_str)
    q, item = pickle.loads(bytes)

    q.put(item)

    q.join()


class TestQueue(unittest.TestCase):
    def test_simple(self):
        """Test basic functionality"""

        q = Queue(joinable=False)
        item = "Queue-item-0"

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps((q, item), protocol=5))

        the_descr = process_create(cmd, wdir, [__file__, "putter", env_str], None, options=options)

        new_item = q.get()

        self.assertTrue(new_item == item)

        process_join(the_descr.p_uid)

    def test_joinable(self):
        """Test that the queue can be joinable"""

        q = Queue(joinable=True)

        self.assertTrue(q._joinable == True)

        item = "JQueue-item-0"

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps((q, item), protocol=5))

        the_descr = process_create(cmd, wdir, [__file__, "joiner", env_str], None, options=options)

        new_item = q.get()
        self.assertTrue(new_item == item)
        time.sleep(0.1)

        ready = process_join(the_descr.p_uid, timeout=0)  # worker remains blocked
        self.assertTrue(ready == None)

        q.task_done()  # unblock worker

        ready = process_join(the_descr.p_uid, timeout=None)  # must cleanly exit
        self.assertTrue(ready == 0)


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "putter":
        _putter(sys.argv[2])
    elif len(sys.argv) > 1 and sys.argv[1] == "joiner":
        _joiner(sys.argv[2])
    else:
        unittest.main()

test/native/test_queue.py
DragonHPC/dragon
import os
import unittest
import time
import pickle
import sys
import shutil

import dragon

# from dragon.native.process import current
# from dragon.native.queue import Queue
# from dragon.native.barrier import Barrier, BrokenBarrierError
from dragon.native.process import Popen
from dragon.channels import Channel, ChannelExistsError
from dragon.utils import B64

class TestIORedirection(unittest.TestCase):
    def test_native_process_stdout(self):

        exe = shutil.which('echo')
        proc = Popen(executable=exe, args=['Hello World'], stdout=Popen.PIPE)

        result = ''
        try:
            while True:
                data = proc.stdout.recv()
                result += data
        except EOFError:
            pass

        self.assertEqual('Hello World\n', result)
        proc.stdout.close()

    def test_native_process_stdout_to_devnull(self):
        '''
            No output should appear on the terminal from running this test.
        '''
        exe = shutil.which('echo')
        proc = Popen(executable=exe, args=['Hello World'], stdout=Popen.DEVNULL)

        self.assertEqual(proc.stdout, None)

    def test_native_process_stderr(self):
        exe = shutil.which('cat')
        proc = Popen(executable=exe, args=['notafile'], stderr=Popen.PIPE)

        result = ''
        try:
            while True:
                data = proc.stderr.recv()
                result += data
        except EOFError:
            pass

        self.assertEqual('/usr/bin/cat: notafile: No such file or directory\n', result)
        proc.stderr.close()

    def test_native_process_stderr_to_devnull(self):
        '''
            No output should appear on the terminal from running this test.
        '''
        exe = shutil.which('cat')
        proc = Popen(executable=exe, args=['notafile'], stderr=Popen.DEVNULL)

        self.assertEqual(proc.stderr, None)

    def test_native_process_stdin(self):
        exe = shutil.which('cat')
        proc = Popen(executable=exe, args=['-'], stdin=Popen.PIPE, stdout=Popen.PIPE)

        x = 'Hello, How Are You?'
        proc.stdin.send(x)
        proc.stdin.close()

        result = ''
        try:
            while True:
                data = proc.stdout.recv()
                result += data
        except EOFError:
            pass

        self.assertEqual(x, result)
        proc.stdout.close()

    def test_native_process_stderr_to_stdout(self):
        exe = shutil.which('cat')
        proc = Popen(executable=exe, args=['notafile'], stderr=Popen.STDOUT, stdout=Popen.PIPE)

        result = ''
        try:
            while True:
                data = proc.stdout.recv()
                result += data
        except EOFError:
            pass

        self.assertEqual('/usr/bin/cat: notafile: No such file or directory\n', result)
        proc.stdout.close()

if __name__ == "__main__":
    unittest.main()

test/native/test_redirection.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import time
import pickle
import sys
import os

import dragon
import dragon.globalservices.channel
from dragon.globalservices.process import (
    multi_join,
    create,
    get_list,
    join,
)
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.native.semaphore import Semaphore
import dragon.utils as du


def _waiter(env_str):

    # unpack the semaphore
    bytes = du.B64.str_to_bytes(env_str)
    sem = pickle.loads(bytes)

    # acquire it
    sem.acquire(blocking=True, timeout=None)


class TestSemaphore(unittest.TestCase):
    def setUp(self):
        self.num_proc = 4

    def test_basic_lifecycle(self):

        sem = Semaphore(value=1)

        # init correct ?
        self.assertTrue(sem._bounded == False)
        self.assertTrue(sem._closed == False)
        self.assertTrue(sem._channel.capacity == 1)
        self.assertTrue(1 == sem.get_value())

        # simple acquire correct ?
        ret_val = sem.acquire()
        self.assertTrue(ret_val == True)
        self.assertTrue(0 == sem.get_value())
        self.assertTrue(sem.acquire(False) == False)  # cannot acquire below 0

        # blocking correct ?
        start = time.monotonic()
        ret_val = sem.acquire(blocking=False, timeout=None)
        end = time.monotonic()
        self.assertTrue(ret_val == False)
        self.assertAlmostEqual(0, end - start, 1)

        # timeout correct ?
        start = time.monotonic()
        ret_val = sem.acquire(blocking=True, timeout=0.05)
        end = time.monotonic()
        self.assertTrue(ret_val == False)
        self.assertAlmostEqual(0.05, end - start, 1)

        # can release ?
        sem.release(n=1)
        ret_val = sem.acquire(blocking=True, timeout=None)
        self.assertTrue(ret_val == True)
        ret_val = sem.acquire(blocking=False, timeout=None)
        self.assertTrue(ret_val == False)

        # can release many at once ?
        sem.release(n=16)
        for _ in range(16):
            ret_val = sem.acquire(blocking=True, timeout=None)
            self.assertTrue(ret_val == True)

        ret_val = sem.acquire(blocking=False, timeout=None)
        self.assertTrue(ret_val == False)

        # get_value correct ?
        value = sem.get_value()
        self.assertTrue(value == 0)

        sem.release(n=16)
        value = sem.get_value()
        self.assertTrue(value == 16)

    def test_serialization(self):

        value = 14041981
        sem = Semaphore(value=value)
        # pickle correct ?
        data = pickle.dumps(sem, protocol=5)
        sem2 = pickle.loads(data)

        self.assertTrue(sem2._channel.cuid == sem._channel.cuid)
        self.assertTrue(sem.get_value() == value)
        self.assertTrue(sem2.get_value() == value)

    def test_many_processes_lifecycle(self):

        sem = Semaphore(value=0)

        cmd = sys.executable
        wdir = "."
        options = ProcessOptions(make_inf_channels=True)
        env_str = du.B64.bytes_to_str(pickle.dumps(sem, protocol=5))

        puids = []
        for _ in range(0, self.num_proc):
            the_descr = create(
                cmd,
                wdir,
                [__file__, "waiter", env_str],
                None,
                options=options,
            )
            puids.append(the_descr.p_uid)

        live_puids = get_list()

        for puid in puids:
            self.assertTrue(puid in live_puids)  # all should be blocked

        sem.release()  # release one process

        object_tuples, proc_status = multi_join(puids, timeout=None)
        ready_puids = [i[0] for i in object_tuples]
        self.assertTrue(len(ready_puids) == 1)  # one ! should have returned

        value = sem.get_value()  # semaphore value should 0
        self.assertTrue(value == 0)

        puids.remove(ready_puids[0])

        live_puids = get_list()
        for puid in puids:
            self.assertTrue(puid in live_puids)  # all should be blocked

        # release them all
        sem.release(n=self.num_proc - 1)

        for puid in puids:  # should not hang
            join(puid)

    def test_is_bounded(self):

        bsem = Semaphore(value=21071985, bounded=True)

        self.assertTrue(bsem._bounded == True)
        self.assertRaises(ValueError, bsem.release)


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "waiter":
        _waiter(sys.argv[2])
    else:
        unittest.main()

test/native/test_semaphore.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import pickle
import sys
import ctypes
import random
import string

import dragon
from dragon.globalservices.process import create, multi_join
from dragon.infrastructure.process_desc import ProcessOptions
from dragon.native.value import Value, _TYPECODE_TO_TYPE, _SUPPORTED_TYPES
from dragon.native.queue import Queue
import dragon.utils as du


def create_value_assignment(type):

    if type is ctypes.c_char:
        return [b"0", b"1", bytes(str(random.choice(string.ascii_letters)).encode("utf-8"))]
    elif type is ctypes.c_wchar:
        return ["0", "1", str(random.choice(string.ascii_letters))]
    elif type in [ctypes.c_double, ctypes.c_float]:
        return [
            0.0,
            1.0,
            random.uniform(sys.float_info.min, sys.float_info.max),
            sys.float_info.min,
            sys.float_info.max,
        ]
    else:
        return [0, 1, random.randint(-sys.maxsize - 1, sys.maxsize), -sys.maxsize - 1, sys.maxsize]


def test_value(args):

    bytes = du.B64.str_to_bytes(args)
    value, parent_queue, child_queue, values_assignment = pickle.loads(bytes)

    # For every item in values_assignment
    for val in values_assignment:
        # set value.value to val
        value.value = val
        # put True into child queue
        child_queue.put(True)
        # wait on the parent queue and exit cleanly
        assert parent_queue.get(timeout=None)


class TestValue(unittest.TestCase):
    def test_requirement_1_1(self):
        """typedef_or_type type checking"""
        self.assertRaises(AttributeError, Value, (0))

    def test_requirement_1_2(self):
        """typedef_or_type type checking"""
        self.assertRaises(AttributeError, Value, (str("v")))

    def test_requirement_1_3(self):
        """typedef_or_type type checking"""
        self.assertRaises(AttributeError, Value, ())

    def test_ping_pong(self):
        """queue ping pong between 2 processes that tests value assignment in
        parent and child processes"""

        for type in _SUPPORTED_TYPES:

            # Create a list of values to be tested values_assignment: bytes, char, int, and float
            values_assignment = create_value_assignment(type)

            # Create a Value v of type type
            value = Value(typecode_or_type=type)

            # Start a process, hand over v, a queue q and the list of value to be tested values_assignment
            cmd = sys.executable
            wdir = "."
            options = ProcessOptions(make_inf_channels=True)
            # naming the queues based on which process puts true in the queue
            parent_queue, child_queue = Queue(joinable=False), Queue(joinable=False)
            puids = []
            env_str = du.B64.bytes_to_str(
                pickle.dumps([value, parent_queue, child_queue, values_assignment], protocol=5)
            )
            proc = create(
                cmd,
                wdir,
                [__file__, "test_value", env_str],
                None,
                options=options,
            )

            # For every item in the list of values to be tested
            for val in values_assignment:

                # wait on the child queue _ = q.get(timeout=None)
                self.assertTrue(child_queue.get(timeout=None))

                # assert that value.value is equal to the current item
                self.assertEqual(value.value, val)

                # put True into the parent queue q.put(True)
                parent_queue.put(True)

            puids.append(proc.p_uid)

        # join on the puid of the process
        ready = multi_join(puids)

        # check the exit code of the process from ready = multi_join(puids) from ready[0][0]
        self.assertEqual(ready[0][0][1], 0)


if __name__ == "__main__":

    if len(sys.argv) > 1 and sys.argv[1] == "test_value":
        test_value(sys.argv[2])
    else:
        unittest.main()

test/native/test_value.py
DragonHPC/dragon
from dragon.launcher.pmsgqueue import PMsgQueue
import dragon.infrastructure.messages as dmsg
import os

pid = os.getpid()

sendh = PMsgQueue("/test"+str(pid), write_intent=True)
recvh = PMsgQueue("/test"+str(pid), read_intent=True)
sendh.send('Bad Message')
sendh.reset()
recvh.reset()

sendh.send(dmsg.SHPingBE(tag=0, shep_cd="", be_cd="", gs_cd="", default_pd="", inf_pd="").serialize())

msg = recvh.recv()
msg = dmsg.parse(msg)
print(repr(msg))


recvh.close()
sendh.close(destroy=True)

test/pmsgq/ex.py
DragonHPC/dragon
#!/usr/bin/env python3

import dragon.mpbridge

import multiprocessing
import sys
import time

def by_hyenas(what_eaten):
    print('hyenas ate your {}'.format(what_eaten))


def main(flag):
    print('> hello from the head process')

    if '-u' == flag:
        print('> using mpbridge.DragonProcess')
        method = 'dragon'

        print(f'> setting start method to {method}')
        multiprocessing.set_start_method(method)
        print(f'> making process')
        p = dragon.mpbridge.DragonProcess(target=dragon.mpbridge.CloudyWrapFunc(by_hyenas),
                                         args=('pineapples',))

    elif '-s' == flag:
        method = 'spawn'
        print(f'> setting start method to {method}')
        multiprocessing.set_start_method(method)
        print(f'> making process')
        p = multiprocessing.Process(target=by_hyenas, args=('pineapples',))
    elif '-d' == flag:
        print('> using mpbridge.DragonProcess')
        method = 'impbridge'
        print(f'> setting start method to {method}')
        multiprocessing.set_start_method(method)
        print(f'> making process')
        p = dragon.mpbridge.DragonProcess(target=dragon.mpbridge.CloudyWrapFunc(by_hyenas),
                                           args=('pineapples',))

    else:
        raise NotImplementedError('derp')

    print(f'> starting process')
    p.start()
    print('> process is started\n')
    sys.stdout.flush()
    print('> trying to join()')
    p.join()
    print('> goodbye from the head process')


if __name__ == "__main__":
    if len(sys.argv) >= 2:
        flag = sys.argv[1]
    else:
        flag = '-s'

    # 'spawn' Process vs 'mpbridge' Process vs eventual dragon process
    assert flag in {'-s', '-u', '-d'}

    main(flag)

test/process/firstlight.py
DragonHPC/dragon
import dragon
import multiprocessing as mp

def f():
    print('Hello World')

def main():
    p = mp.Process(target=f, args=())
    p.start()
    p.join()

if __name__ == "__main__":
    main()

test/process/mphello.py
DragonHPC/dragon
import dragon
import dragon.native.process as nativeproc
import logging

LOG = logging.getLogger('NativeProcTest :')

def main():
    proc = nativeproc.Popen([],"hello.py",stdout=nativeproc.PIPE)

    try:
        while True:
            output = proc.stdout.recv()
            print("From the pipe:", output, end="", flush=True)
    except EOFError:
        pass
    except Exception as ex:
        print(f'Unexpected error: {repr(ex)}')

    del proc

if __name__=="__main__":
    main()
test/process/nativeprocess.py
DragonHPC/dragon
import time

def foo():
    print("Hello World")

if __name__ == "__main__":
    foo()

test/shepherd/proc1.py
DragonHPC/dragon
import time

def foo():
    print("Hello World")
    print("Doing some more")


if __name__ == "__main__":
    foo()

test/shepherd/proc2.py
DragonHPC/dragon
import time

def foo():
    print("Hello World")
    time.sleep(3)
    print("Doing some more")


if __name__ == "__main__":
    foo()

test/shepherd/proc3.py
DragonHPC/dragon
import time

def foo():
    s = input("Enter some text: ")
    print("Here is the response text.")
    print(s)


if __name__ == "__main__":
    foo()

test/shepherd/proc4.py
DragonHPC/dragon
import os

def foo():

    print(os.environ['new_env_var'])

if __name__ == "__main__":
    foo()

test/shepherd/procenv1.py
DragonHPC/dragon
import os

def foo():

    print(os.environ['SHELL'])

if __name__ == "__main__":
    foo()

test/shepherd/procenv2.py
DragonHPC/dragon
import pathlib
import sys
import os


# Needed by tests spawning a new python process by executable name
orig_path = os.getenv('PATH')
py_exe_path = str(pathlib.Path(sys.executable).parent)
os.putenv('PATH', py_exe_path + ':' + orig_path)


# Needed if no PYTHONPATH is set to point at dragon/src or dragon not installed in site-packages
dragon_path = pathlib.Path(__file__).parent / '..' / 'src'
support_path = pathlib.Path(__file__).parent
if not dragon_path.exists():
    # Sometimes needed by fixtures such as launcher/fe_server.py
    dragon_path = pathlib.Path(__file__).parent / '..' / '..' / 'src'
    support_path = pathlib.Path(__file__).parent / '..'

for addl_path in (support_path, dragon_path):
    addl_abspath = str(addl_path.absolute())
    if addl_abspath not in sys.path:
        sys.path.append(addl_abspath)

test/shim_dragon_paths.py
DragonHPC/dragon
#!/usr/bin/env python3

"""
Test script for launcher startup
"""

import dragon.globalservices.process_desc as proc
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.log_setup as log_setup
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.parameters as parms
import dragon.utils as du
import dragon.launcher.backend as backend
import dragon.launcher.frontend as frontend
import inspect
import logging
import multiprocessing as mp
import support.util as tsu
import unittest
import shim_dragon_paths

# required to get certain tests to run correctly
mp.set_start_method("spawn", force=True)


class StreamQueue:
    def __init__(self, *args, **kwargs):
        self.queue = mp.Queue(*args, **kwargs)

    def write(self, msg):
        self.queue.put(msg)

    def flush(self):
        pass

    def read(self):
        rv = ""

        while not self.queue.empty():
            rv += self.queue.get()

        return rv

    def readline(self):
        return self.read()

    def close(self):
        self.queue.close()
        self.queue.join_thread()


########################################################################
# This is the unit test file for the Launcher Dragon Run-time Services #
########################################################################
# The unit tests presented here are documented in comments in this file
# and further documentation can be found in the Dragon documentation
# under the "Dragon API and Infrastructure Testing" section. In addition,
# documentation on the Launcher design and functionality can be found
# under the section titled "The Launcher" in the "Internal Documentation"
# of the Dragon Runtime services.
########################################################################

def startup_and_run_frontend(launcher_queue, backend_queue, argv, stdout, stderr, name_addition=''):
    log_setup.setup_logging(basename='la_frontend_' + name_addition, the_level=logging.DEBUG)
    frontend.launcher_start(launcher_queue, backend_queue, argv=argv, stdout=stdout, stderr=stderr)


PDESC = proc.ProcessDescriptor
PRESP = dmsg.GSProcessCreateResponse
PRESPERR = dmsg.GSProcessCreateResponse.Errors


class FrontendTests(unittest.TestCase):
    def setUp(self) -> None:
        self.launcher_main_rh, self.launcher_main_wh = mp.Pipe()
        self.backend_main_rh, self.backend_main_wh = mp.Pipe()
        self.stdout = StreamQueue()
        self.stderr = StreamQueue()

    def tearDown(self) -> None:
        self.launcher_main_rh.close()
        self.launcher_main_wh.close()
        self.backend_main_rh.close()
        self.backend_main_wh.close()
        self.stdout.close()
        self.stderr.close()

    def start_launcher(self, args=None):
        if args is None:
            args = []

        test_name = self.__class__.__name__ + '_' + inspect.stack()[2][0].f_code.co_name

        self.proc = mp.Process(target=startup_and_run_frontend,
                               args=(self.launcher_main_rh, self.backend_main_wh, args, self.stdout, self.stderr),
                               kwargs={'name_addition': test_name},
                               daemon=False,
                               name='launcher_fe')

        self.proc.start()

    def do_bringup(self, argv=None):

        if argv is None:
            argv = []

        self.start_launcher(argv)
        self.launcher_main_wh.send(dmsg.GSIsUp(tag=0).serialize())

    def do_teardown(self, timeout=tsu.DEFAULT_TIMEOUT):
        # The teardown is initiated by an exit of the launcher. When this happens,
        # the launcher sends the GSTeardown message to initiate shutdown
        gs_teardown = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSTeardown, timeout=timeout)

        # Send the GSHalted
        self.launcher_main_wh.send(dmsg.GSHalted(tag=0).serialize())
        # Get the SHTeardown from the Launcher. All messages are routed
        # through the backend.
        la_broadcast = tsu.get_and_check_type(self.backend_main_rh, dmsg.LABroadcast, timeout=timeout)

        sh_teardown = dmsg.parse(la_broadcast.data)

        assert isinstance(sh_teardown, dmsg.SHTeardown), 'Expected SHTeardown, got: ' + repr(gs_teardown)

        # The rest of the single-node teardown sequence is carried out by the
        # backend code.

        logging.shutdown()

        # Wait for the launcher frontend to exit here so we know that
        # we have completely finished processing all messages before
        # trying to verify output in individual test cases below.
        self.proc.join(10)

        # If the frontend did not exit within 10 seconds, then something is wrong.
        self.assertTrue(self.proc.exitcode is not None)


    def test_bringup_teardown(self):
        # Tests normal bringup followed by teardown.

        self.do_bringup(argv=["-r", "launcher/file0.py"])

        self.do_teardown()

    def test_start_proc1(self):

        args = ["launcher/test.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        self.launcher_main_wh.send(dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                                              err=dmsg.GSProcessJoinResponse.Errors.SUCCESS).serialize())

        self.do_teardown(timeout=None)

    def test_start_proc2(self):

        args = ["-r", "launcher/file1.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        self.launcher_main_wh.send(dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                                              err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                                              ).serialize())

        gs_process_kill = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessKill)

        self.launcher_main_wh.send(dmsg.GSProcessKillResponse(tag=0, ref=gs_process_kill.tag,
                                                              err=dmsg.GSProcessKillResponse.Errors.UNKNOWN,
                                                              ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # Errors.SUCCESS
        # Errors.SUCCESS
        # Errors.SUCCESS
        # [Kill unsuccessful, error is: UNKNOWN]
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        # For future tests
        # ====================
        # use these two print statements to determine output, then assertEqual for the actual text as
        # it appears in this test.
        # print("stdout", stdout)
        # print("stderr", stderr)

        self.assertIn('SUCCESS', stdout)
        self.assertIn('True', stdout)
        self.assertIn('UNKNOWN', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_start_proc3(self):

        args = ["-r", "launcher/file1.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        self.launcher_main_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        gs_process_kill = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessKill)

        self.launcher_main_wh.send(
            dmsg.GSProcessKillResponse(tag=0, ref=gs_process_kill.tag,
                                       err=dmsg.GSProcessKillResponse.Errors.SUCCESS,
                                       ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()
        stderr = self.stderr.read()

        self.assertIn('Errors.SUCCESS', stdout)
        self.assertIn('0', stdout)
        self.assertIn('True', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_start_proc4(self):
        args = ["-r", "launcher/file2.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        msg = dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                         err=dmsg.GSProcessJoinResponse.Errors.SUCCESS)
        self.launcher_main_wh.send(msg.serialize())

        self.launcher_main_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.launcher_main_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                             data="Hello World to stderr\n").serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        stderr = self.stderr.read()

        self.assertIn('Hello World', stdout)
        self.assertIn('True', stdout)

        self.assertIn('Hello World', stderr)
        self.assertIn('stderr', stderr)

    def test_start_proc5(self):

        args = ["-r", "launcher/file3.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)
        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        msg = dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0",
                               fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value, data="Hello World\n")
        self.launcher_main_wh.send(msg.serialize())

        msg = dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                               data="Hello World to stderr\n")
        self.launcher_main_wh.send(msg.serialize())

        msg = dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                         err=dmsg.GSProcessJoinResponse.Errors.SUCCESS)
        self.launcher_main_wh.send(msg.serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # [stdout: p_uid=4] Hello World
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        # '''[stderr: p_uid=4] Hello World to stderr'''

        self.assertIn('SUCCESS', stdout)
        self.assertIn('True', stdout)
        self.assertIn('Hello World', stdout)
        self.assertIn('[Exiting Launcher]', stdout)

        self.assertIn('Hello World', stderr)
        self.assertIn('stderr', stderr)

    def test_start_proc6(self):
        args = ["-r", "launcher/file4.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)
        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        self.launcher_main_wh.send(dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0",
                                                    fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                                                    data="Hello World\n").serialize())

        self.launcher_main_wh.send(dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0",
                                                    fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                                                    data="Hello World to stderr\n").serialize())

        self.launcher_main_wh.send(dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0",
                                                    fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                                                    data="Stop\n").serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        self.launcher_main_wh.send(dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                                              err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                                              ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()
        # '''(4, 'SUCCESS', None)
        # Got pipe text: Hello World
        #
        # Got pipe text: Hello World to stderr
        #
        # Got pipe text: Stop
        #
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('4', stdout)
        self.assertIn('Hello World', stdout)
        self.assertIn('stderr', stdout)
        self.assertIn('Stop', stdout)
        self.assertIn('True', stdout)
        self.assertIn('[Exiting Launcher]', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_start_proc7(self):

        args = ["-r", "launcher/file1.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        self.launcher_main_wh.send(dmsg.GSUnexpected(tag=0, ref=gs_process_join.tag).serialize())

        gs_process_kill = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessKill)

        self.launcher_main_wh.send(dmsg.GSUnexpected(tag=0, ref=gs_process_kill.tag).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # Errors.SUCCESS
        # Errors.SUCCESS
        # Errors.SUCCESS
        # [Kill unsuccessful, error is: UNKNOWN]
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('SUCCESS', stdout)
        self.assertIn('True', stdout)
        self.assertIn('Kill unsuccessful', stdout)
        self.assertIn('[Exiting Launcher]', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_start_proc8(self):

        args = ["-r", "launcher/file8.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        self.launcher_main_wh.send(dmsg.GSProcessCreateResponse(tag=3, err=1, ref=gs_process_create.tag,
                                                                err_info="[Errno 13] Permission denied: "
                                                                         "'python3'").serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''[Create unsuccessful, error is: FAIL]
        # p_uid is None
        # (None, 'FAIL', "[Errno 13] Permission denied: 'python3'")
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('FAIL', stdout)
        self.assertIn('None', stdout)
        self.assertIn('Permission denied', stdout)
        self.assertIn('True', stdout)
        self.assertIn('[Exiting Launcher]', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_start_server(self):

        args = ["-r", "launcher/file5.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        self.launcher_main_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.launcher_main_wh.send(dmsg.LAPassThruBF(tag=0, data="PassThru Message").serialize())

        lapassthrufb = tsu.get_and_check_type(self.backend_main_rh, dmsg.LAPassThruFB)

        self.assertEqual(lapassthrufb.data, "PassThru Message")

        self.launcher_main_wh.send(dmsg.LAPassThruBF(tag=0, data="Another Message").serialize())

        lapassthrufb = tsu.get_and_check_type(self.backend_main_rh, dmsg.LAPassThruFB)

        self.assertEqual(lapassthrufb.data, "Another Message")

        self.launcher_main_wh.send(dmsg.LAPassThruBF(tag=0, data="Exit").serialize())

        self.launcher_main_wh.send(dmsg.LAServerModeExit(tag=0, ref=gs_process_create.tag,
                                                         err=dmsg.LAServerModeExit.Errors.SUCCESS).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()
        # '''[stdout: p_uid=4] Hello World
        # Server Mode Exited
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('Hello', stdout)
        self.assertIn('Server', stdout)
        self.assertIn('Exited', stdout)
        self.assertIn('True', stdout)
        self.assertIn('[Exiting Launcher]', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_proc_list(self):

        args = ["-r", "launcher/file6.py"]

        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        self.launcher_main_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.launcher_main_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                             data="Hello World to stderr\n").serialize())

        self.launcher_main_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Stop\n").serialize())

        gs_process_join = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessJoin)

        self.launcher_main_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        gs_process_list = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessList)

        self.launcher_main_wh.send(
            dmsg.GSProcessListResponse(tag=0, ref=gs_process_list.tag, err=dmsg.GSProcessListResponse.Errors.SUCCESS,
                                       plist=[('4', '4-name')]).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': None}}
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': None}}
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': 0}}
        # {'4': '4-name'}
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('python3', stdout)
        self.assertIn('exit_code', stdout)
        self.assertIn('4-name', stdout)
        self.assertIn('True', stdout)
        self.assertIn('[Exiting Launcher]', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_send(self):

        args = ["-r", "launcher/file7.py"]

        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.backend_main_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        sh_process_input = tsu.get_and_check_type(self.backend_main_rh, dmsg.SHFwdInput)

        self.launcher_main_wh.send(
            dmsg.SHFwdInputErr(tag=0, ref=sh_process_input.tag, err=dmsg.SHFwdInputErr.Errors.SUCCESS).serialize())

        self.assertEqual(sh_process_input.input, "Hello World")

        sh_process_input2 = tsu.get_and_check_type(self.backend_main_rh, dmsg.SHFwdInput)

        self.launcher_main_wh.send(
            dmsg.SHFwdInputErr(tag=0, ref=sh_process_input2.tag, err=dmsg.SHFwdInputErr.Errors.SUCCESS).serialize())

        self.assertEqual(sh_process_input2.input, "Goodbye")

        self.launcher_main_wh.send(dmsg.LADumpState(tag=0).serialize())

        self.launcher_main_wh.send(dmsg.LADumpState(tag=0, filename="ladumptest.log").serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': None}}
        # response from send (<Errors.SUCCESS: 0>, '')
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('4', stdout)
        self.assertIn('python3', stdout)
        self.assertIn('True', stdout)
        self.assertIn('Exit', stdout)

        self.assertEqual(stderr.strip(), '')


def startup_and_run_backend(incoming_channel_queue=None, incoming_frontend_queue=None, incoming_l0_queue=None,
                            frontend_queue=None, gs_queue=None, sh_queue=None, sh_l0_queue=None,
                            all_other_channels=None, name_addition=''):
    log_setup.setup_logging(basename='la_backend_' + name_addition, the_level=logging.DEBUG)

    backend.backend_launcher_start(incoming_channel_queue, incoming_frontend_queue, incoming_l0_queue,
                                   frontend_queue, gs_queue, sh_queue, sh_l0_queue, all_other_channels,
                                   mode=dfacts.TEST_MODE)


class BackendTests(unittest.TestCase):
    def setUp(self) -> None:
        self.launcher_main_rh, self.launcher_main_wh = mp.Pipe()
        self.backend_main_rh, self.backend_main_wh = mp.Pipe()
        self.main_queue_rh, self.main_queue_wh = mp.Pipe()
        self.from_frontend_queue_rh, self.from_frontend_queue_wh = mp.Pipe()
        self.l0_queue_rh, self.l0_queue_wh = mp.Pipe()
        self.to_frontend_queue_rh, self.to_frontend_queue_wh = mp.Pipe()
        self.gs_queue_rh, self.gs_queue_wh = mp.Pipe()
        self.sh_queue_rh, self.sh_queue_wh = mp.Pipe()
        self.sh_l0_queue_rh, self.sh_l0_queue_wh = mp.Pipe()
        self.all_other_channels_rh, self.all_other_channels_wh = mp.Pipe()
        self.proc = None

    def tearDown(self) -> None:
        self.launcher_main_rh.close()
        self.launcher_main_wh.close()
        self.backend_main_rh.close()
        self.backend_main_wh.close()
        self.main_queue_rh.close()
        self.main_queue_wh.close()
        self.from_frontend_queue_rh.close()
        self.from_frontend_queue_wh.close()
        self.l0_queue_rh.close()
        self.l0_queue_wh.close()
        self.to_frontend_queue_rh.close()
        self.to_frontend_queue_wh.close()
        self.gs_queue_rh.close()
        self.gs_queue_wh.close()
        self.sh_queue_rh.close()
        self.sh_queue_wh.close()
        self.sh_l0_queue_rh.close()
        self.sh_l0_queue_wh.close()
        self.all_other_channels_rh.close()
        self.all_other_channels_wh.close()
        if self.proc is not None:
            self.proc.join(10)
            # If the backend did not exit within 10 seconds, then something is wrong.
            self.assertTrue(self.proc.exitcode is not None)

    def start_launcher_backend(self):
        if self.proc is not None:
            self.proc.join(10)
            # If the backend did not exit within 10 seconds, then something is wrong.
            self.assertTrue(self.proc.exitcode is not None)

        test_name = self.__class__.__name__ + '_' + inspect.stack()[2][0].f_code.co_name

        self.proc = mp.Process(target=startup_and_run_backend,
                               args=(self.launcher_main_rh, self.from_frontend_queue_rh, self.l0_queue_rh,
                                     self.to_frontend_queue_wh, self.gs_queue_wh, self.sh_queue_wh, self.sh_l0_queue_wh,
                                     self.all_other_channels_wh),
                               kwargs={'name_addition': test_name},
                               daemon=False, name='backend')

        self.proc.start()

    def do_bringup(self):
        # Start up the launcher backend by itself.
        self.start_launcher_backend()

        # do the message bringup sequence.
        # Get the node index message from the backend.
        tsu.get_and_check_type(self.sh_l0_queue_rh, dmsg.BENodeIdxSH)

        # Send the BEPingSH as part of the single node startup sequence
        self.l0_queue_wh.send(
            dmsg.SHPingBE(tag=0, shep_cd="", be_cd="", gs_cd="", default_pd="", inf_pd="").serialize())

        # Then the Shepherd is sent the SHPingBE message
        tsu.get_and_check_type(self.sh_queue_rh, dmsg.BEPingSH)

        # Next, SHChannelsUp
        self.launcher_main_wh.send(
            dmsg.SHChannelsUp(tag=0, host_id='ubuntu', ip_addrs=['127.0.0.1'], shep_cd=du.B64.bytes_to_str(b''),
                              gs_cd=du.B64.bytes_to_str(b'')).serialize())

        # Next, GSIsUp
        self.launcher_main_wh.send(dmsg.GSIsUp(tag=0).serialize())

        # And, it comes out to the launcher front end.
        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.GSIsUp)

    def do_teardown(self):
        # The teardown is initiated by an exit of the launcher. When this happens,
        # the launcher sends the GSTeardown message to initiate shutdown
        self.launcher_main_wh.send(dmsg.GSTeardown(tag=0).serialize())

        # Check that it gets routed to Global Services
        tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSTeardown)

        # Send the GSHalted
        self.launcher_main_wh.send(dmsg.GSHalted(tag=0).serialize())

        # Check that it gets routed to Front End
        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.GSHalted)

        # Get the SHTeardown from the Launcher. All messages are routed
        # through the backend.
        self.launcher_main_wh.send(dmsg.LABroadcast(tag=0, data=dmsg.SHTeardown(tag=0).serialize()).serialize())

        # See that the SHTeardown gets routed to the Shepherd.
        tsu.get_and_check_type(self.sh_queue_rh, dmsg.SHTeardown)

        # Send the SHHaltBE to the backend only. The backend responds to this one.
        self.launcher_main_wh.send(dmsg.SHHaltBE(tag=0).serialize())

        # Check to see the BEHalted got written to the Shepherd stdin.
        tsu.get_and_check_type(self.sh_l0_queue_rh, dmsg.BEHalted)

        # Send the SHHalted Message
        self.l0_queue_wh.send(dmsg.SHHalted(tag=0).serialize())

        # At this point the Backend should exit.
        logging.shutdown()

    def test_bringup_teardown(self):
        self.do_bringup()
        self.do_teardown()

    def test_message_routes(self):
        self.do_bringup()

        self.launcher_main_wh.send(dmsg.GSProcessCreate(tag=1, p_uid=dfacts.LAUNCHER_PUID,
                                                        r_c_uid=dfacts.BASE_BE_CUID, exe="launcher/test.py",
                                                        args=[]).serialize())

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)
        proc_desc = proc.ProcessDescriptor(p_uid=4, name='TestProc', node=1, p_p_uid=0)
        msg = dmsg.GSProcessCreateResponse(tag=0, ref=gs_process_create.tag,
                                           err=dmsg.GSProcessCreateResponse.Errors.SUCCESS, desc=proc_desc)
        self.launcher_main_wh.send(msg.serialize())

        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.GSProcessCreateResponse)

        self.launcher_main_wh.send(dmsg.GSProcessKill(tag=1, t_p_uid=4, p_uid=dfacts.LAUNCHER_PUID,
                                                      r_c_uid=dfacts.BASE_BE_CUID, sig=9).serialize())

        tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessKill)

        self.launcher_main_wh.send(dmsg.GSProcessKillResponse(tag=0, ref=4,
                                                              err=dmsg.GSProcessKillResponse.Errors.UNKNOWN,
                                                              exit_code=0).serialize())

        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.GSProcessKillResponse)

        self.launcher_main_wh.send(dmsg.GSProcessJoin(tag=1, p_uid=dfacts.LAUNCHER_PUID, t_p_uid=4,
                                                      r_c_uid=dfacts.BASE_BE_CUID).serialize())

        tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.launcher_main_wh.send(dmsg.GSProcessJoinResponse(tag=0, ref=1,
                                                              err=dmsg.GSProcessJoinResponse.Errors.SUCCESS).serialize())

        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.GSProcessJoinResponse)

        self.launcher_main_wh.send(dmsg.GSProcessList(tag=0, r_c_uid=dfacts.BASE_BE_CUID,
                                                      p_uid=dfacts.LAUNCHER_PUID).serialize())

        gs_process_list = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessList)

        self.launcher_main_wh.send(dmsg.GSProcessListResponse(tag=0, ref=gs_process_list.tag,
                                                              err=dmsg.GSProcessListResponse.Errors.SUCCESS,
                                                              plist=[('4', '4-name')]).serialize())

        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.GSProcessListResponse)

        self.launcher_main_wh.send(dmsg.SHFwdInput(tag=0, p_uid=dfacts.LAUNCHER_PUID, t_p_uid=4,
                                                   r_c_uid=dfacts.BASE_BE_CUID,
                                                   input="Hello World").serialize())

        tsu.get_and_check_type(self.sh_queue_rh, dmsg.SHFwdInput)

        self.launcher_main_wh.send(dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0",
                                                    fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                                                    data="Hello World\n").serialize())

        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.SHFwdOutput)

        self.launcher_main_wh.send(dmsg.LAPassThruBF(tag=0, data="PassThru Message").serialize())

        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.LAPassThruBF)

        self.launcher_main_wh.send(dmsg.LAPassThruFB(tag=0, c_uid=5, data="hello world").serialize())

        tsu.get_and_check_type(self.all_other_channels_rh, dmsg.LAPassThruFB)

        self.launcher_main_wh.send(
            dmsg.LAServerModeExit(tag=0, ref=0, err=dmsg.LAServerModeExit.Errors.SUCCESS).serialize())

        tsu.get_and_check_type(self.to_frontend_queue_rh, dmsg.LAServerModeExit)

        self.do_teardown()


class CombinedTests(unittest.TestCase):
    def setUp(self) -> None:
        self.backend_channel_rh, self.backend_channel_wh = mp.Pipe()
        self.backend_to_frontend_rh, self.backend_to_frontend_wh = mp.Pipe()
        self.frontend_to_backend_rh, self.frontend_to_backend_wh = mp.Pipe()
        self.l0_queue_rh, self.l0_queue_wh = mp.Pipe()
        self.gs_queue_rh, self.gs_queue_wh = mp.Pipe()
        self.sh_queue_rh, self.sh_queue_wh = mp.Pipe()
        self.sh_l0_queue_rh, self.sh_l0_queue_wh = mp.Pipe()
        self.all_other_channels_rh, self.all_other_channels_wh = mp.Pipe()
        self.stdout = StreamQueue()
        self.stderr = StreamQueue()

    def tearDown(self) -> None:
        self.backend_channel_rh.close()
        self.backend_channel_wh.close()
        self.backend_to_frontend_rh.close()
        self.backend_to_frontend_wh.close()
        self.frontend_to_backend_rh.close()
        self.frontend_to_backend_wh.close()
        self.l0_queue_rh.close()
        self.l0_queue_wh.close()
        self.gs_queue_rh.close()
        self.gs_queue_wh.close()
        self.sh_queue_rh.close()
        self.sh_queue_wh.close()
        self.sh_l0_queue_rh.close()
        self.sh_l0_queue_wh.close()
        self.all_other_channels_rh.close()
        self.all_other_channels_wh.close()

    def start_combined(self, args=None):
        if args is None:
            args = []
        test_name = self.__class__.__name__ + '_' + inspect.stack()[2][0].f_code.co_name

        self.backend = mp.Process(target=startup_and_run_backend,
                                  args=(self.backend_channel_rh, self.frontend_to_backend_rh, self.l0_queue_rh,
                                        self.backend_to_frontend_wh, self.gs_queue_wh, self.sh_queue_wh,
                                        self.sh_l0_queue_wh, self.all_other_channels_wh),
                                  kwargs={'name_addition': test_name},
                                  daemon=False, name='backend')

        self.backend.start()

        self.frontend = mp.Process(target=startup_and_run_frontend,
                                   args=(self.backend_to_frontend_rh, self.frontend_to_backend_wh,
                                         args, self.stdout, self.stderr),
                                   kwargs={'name_addition': test_name},
                                   daemon=False, name='frontend')

        self.frontend.start()

    def do_bringup(self, argv=None):
        # This is executed to do the bring up of the Shepherd in test. Before this
        # executes in any test, the setUp method has already run. This creates
        # the Shepherd as a multiprocessing Process with all the connections into
        # it for unit testing.
        if argv is None:
            argv = []

        self.start_combined(argv)

        # To finish bringup the sequence of bringup messages are initiated and
        # checked with the Shepherd.
        # Get the node index message from the backend.
        tsu.get_and_check_type(self.sh_l0_queue_rh, dmsg.BENodeIdxSH)

        # Send the BEPingSH as part of the single node startup sequence
        self.l0_queue_wh.send(
            dmsg.SHPingBE(tag=0, shep_cd="", be_cd="", gs_cd="", default_pd="", inf_pd="").serialize())

        # First receive the SHPingBE message
        tsu.get_and_check_type(self.sh_queue_rh, dmsg.BEPingSH)

        # Send the SHChannelsUp to complete the Shepherd part of the sequence
        up_msg = dmsg.SHChannelsUp(tag=0, host_id='ubuntu', ip_addrs=['127.0.0.1'],
                                   shep_cd=du.B64.bytes_to_str(b''), gs_cd=du.B64.bytes_to_str(b''))
        self.backend_channel_wh.send(up_msg.serialize())

        # Send the GSIsUp to complete the sequence
        self.backend_channel_wh.send(dmsg.GSIsUp(tag=0).serialize())

    def do_teardown(self, timeout=tsu.DEFAULT_TIMEOUT):
        # The teardown is initiated by an exit of the launcher. When this happens,
        # the launcher sends the GSTeardown message to initiate shutdown
        tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSTeardown, timeout=timeout)

        # Send the GSHalted
        self.backend_channel_wh.send(dmsg.GSHalted(tag=0).serialize())

        # Get the SHTeardown from the Launcher. All messages are routed
        # through the backend.
        tsu.get_and_check_type(self.sh_queue_rh, dmsg.SHTeardown, timeout=timeout)

        # The rest of the single-node teardown sequence is carried out by the
        # backend code.

        # Send the SHHaltBE to the backend only. The backend responds to this one.
        self.backend_channel_wh.send(dmsg.SHHaltBE(tag=0).serialize())

        # Check to see the BEHalted got written to the Shepherd stdin.
        tsu.get_and_check_type(self.sh_l0_queue_rh, dmsg.BEHalted)

        # Send the SHHalted Message
        self.l0_queue_wh.send(dmsg.SHHalted(tag=0).serialize())

        logging.shutdown()

        # wait for the frontend to exit before returning
        # from the do_teardown so output can be verified by the
        # test case.
        self.frontend.join(10)
        # If the frontend did not exit within 10 seconds, then something is wrong.
        self.assertTrue(self.frontend.exitcode is not None)

    def test_bringup_teardown(self):
        # Tests normal bringup followed by teardown.

        self.do_bringup(argv=["-r", "launcher/file0.py"])

        self.do_teardown()

        stdout = self.stdout.read()
        # '''True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        # removes all whitespace, but order is enforced within the output
        self.assertIn('True', stdout)
        self.assertIn('Exiting', stdout)
        self.assertEqual('', stderr.strip())

    def test_start_proc1(self):
        args = ["launcher/test.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()
        # '''[Exiting Launcher]'''

        stderr = self.stderr.read()
        self.assertIn('Exiting', stdout)

        self.assertEqual('', stderr.strip())

    def test_start_proc2(self):
        args = ["-r", "launcher/file1.py"]

        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        gs_process_kill = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessKill)

        self.backend_channel_wh.send(
            dmsg.GSProcessKillResponse(tag=0, ref=gs_process_kill.tag,
                                       err=dmsg.GSProcessKillResponse.Errors.UNKNOWN,
                                       ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # Errors.SUCCESS
        # Errors.SUCCESS
        # Errors.SUCCESS
        # [Kill unsuccessful, error is: UNKNOWN]
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('SUCCESS', stdout)
        self.assertIn('Kill', stdout)
        self.assertIn('unsuccessful', stdout)
        self.assertIn('True', stdout)

        self.assertEqual('', stderr.strip())

    def test_start_proc3(self):
        args = ["-r", "launcher/file1.py"]

        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        gs_process_kill = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessKill)

        self.backend_channel_wh.send(
            dmsg.GSProcessKillResponse(tag=0, ref=gs_process_kill.tag,
                                       err=dmsg.GSProcessKillResponse.Errors.SUCCESS,
                                       ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # Errors.SUCCESS
        # Errors.SUCCESS
        # Errors.SUCCESS
        # Killed it
        # Exit Code is 0
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('SUCCESS', stdout)
        self.assertIn('Killed', stdout)
        self.assertIn('0', stdout)
        self.assertIn('True', stdout)
        self.assertEqual('', stderr.strip())

    def test_start_proc4(self):
        args = ["-r", "launcher/file2.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0",
                             fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0",
                             fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                             data="Hello World to stderr\n").serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        self.assertIn('Hello World', stdout)
        self.assertIn('True', stdout)

    def test_start_proc5(self):
        args = ["-r", "launcher/file3.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                             data="Hello World to stderr\n").serialize())

        self.backend_channel_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # [stdout: p_uid=4] Hello World
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        # '[stderr: p_uid=4] Hello World to stderr'

        self.assertIn('SUCCESS', stdout)
        self.assertIn('True', stdout)

        self.assertIn('Hello', stderr)

    def test_start_proc6(self):
        args = ["-r", "launcher/file4.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                             data="Hello World to stderr\n").serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Stop\n").serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # expecting:
        # (4, 'SUCCESS', None)
        # Got pipe text: Hello World
        #
        # Got pipe text: Hello World to stderr
        #
        # Got pipe text: Stop
        #
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('4', stdout)
        self.assertIn('Hello World', stdout)
        self.assertIn('stderr', stdout)
        self.assertIn('Stop', stdout)
        self.assertIn('True', stdout)
        self.assertIn('[Exiting Launcher]', stdout)

        self.assertEqual(stderr.strip(), '')

    def test_start_proc7(self):
        args = ["-r", "launcher/file1.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(dmsg.GSUnexpected(tag=0, ref=gs_process_join.tag).serialize())

        gs_process_kill = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessKill)

        self.backend_channel_wh.send(dmsg.GSUnexpected(tag=0, ref=gs_process_kill.tag).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # Errors.SUCCESS
        # Errors.SUCCESS
        # Errors.SUCCESS
        # [Kill unsuccessful, error is: UNKNOWN]
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('SUCCESS', stdout)
        self.assertIn('Kill', stdout)
        self.assertIn('unsuccessful', stdout)
        self.assertIn('UNKNOWN', stdout)
        self.assertIn('True', stdout)

        self.assertEqual('', stderr.strip())

    def test_start_server(self):
        args = ["-r", "launcher/file5.py"]
        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.backend_channel_wh.send(dmsg.LAPassThruBF(tag=0, data="PassThru Message").serialize())

        lapassthrufb = tsu.get_and_check_type(self.all_other_channels_rh, dmsg.LAPassThruFB)

        self.assertEqual(lapassthrufb.data, "PassThru Message")

        self.backend_channel_wh.send(dmsg.LAPassThruBF(tag=0, data="Another Message").serialize())

        lapassthrufb = tsu.get_and_check_type(self.all_other_channels_rh, dmsg.LAPassThruFB)

        self.assertEqual(lapassthrufb.data, "Another Message")

        self.backend_channel_wh.send(dmsg.LAPassThruBF(tag=0, data="Exit").serialize())

        self.backend_channel_wh.send(dmsg.LAServerModeExit(tag=0, ref=gs_process_create.tag,
                                                           err=dmsg.LAServerModeExit.Errors.SUCCESS).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()
        # '''[stdout: p_uid=4] Hello World
        # Server Mode Exited
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('Hello', stdout)
        self.assertIn('World', stdout)
        self.assertIn('Server', stdout)
        self.assertIn('Mode', stdout)
        self.assertIn('True', stdout)

        self.assertEqual('', stderr.strip())

    def test_proc_list(self):
        args = ["-r", "launcher/file6.py"]

        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_uid=4, name='TestProc', node=1, p_p_uid=0)
        msg = PRESP(tag=0, ref=gs_process_create.tag, err=PRESPERR.SUCCESS,
                    desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Hello World\n").serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDERR.value,
                             data="Hello World to stderr\n").serialize())

        self.backend_channel_wh.send(
            dmsg.SHFwdOutput(tag=0, p_uid="4", idx="0", fd_num=dmsg.SHFwdOutput.FDNum.STDOUT.value,
                             data="Stop\n").serialize())

        gs_process_join = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessJoin)

        self.backend_channel_wh.send(
            dmsg.GSProcessJoinResponse(tag=0, ref=gs_process_join.tag,
                                       err=dmsg.GSProcessJoinResponse.Errors.SUCCESS,
                                       ).serialize())

        gs_process_list = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessList)

        self.backend_channel_wh.send(
            dmsg.GSProcessListResponse(tag=0, ref=gs_process_list.tag, err=dmsg.GSProcessListResponse.Errors.SUCCESS,
                                       plist=[('4', '4-name')]).serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()
        # '''(4, 'SUCCESS', None)
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': None}}
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': None}}
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': 0}}
        # {'4': '4-name'}
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('python3', stdout)
        self.assertIn('test.py', stdout)
        self.assertIn('env', stdout)
        self.assertIn('0', stdout)
        self.assertIn('True', stdout)

        self.assertEqual('', stderr.strip())

    def test_send(self):
        args = ["-r", "launcher/file7.py"]

        self.do_bringup(argv=args)

        gs_process_create = tsu.get_and_check_type(self.gs_queue_rh, dmsg.GSProcessCreate)

        proc_desc = PDESC(p_p_uid=0, p_uid=4, name='TestProc', node=1)
        msg = PRESP(tag=0, ref=gs_process_create.tag,
                    err=PRESPERR.SUCCESS, desc=proc_desc)
        self.backend_channel_wh.send(msg.serialize())

        sh_fwd_input = tsu.get_and_check_type(self.sh_queue_rh, dmsg.SHFwdInput)

        self.backend_channel_wh.send(
            dmsg.SHFwdInputErr(tag=0, ref=sh_fwd_input.tag, err=dmsg.SHFwdInputErr.Errors.SUCCESS).serialize())

        self.assertEqual(sh_fwd_input.input, "Hello World")

        sh_fwd_input2 = tsu.get_and_check_type(self.sh_queue_rh, dmsg.SHFwdInput)

        self.backend_channel_wh.send(
            dmsg.SHFwdInputErr(tag=0, ref=sh_fwd_input2.tag, err=dmsg.SHFwdInputErr.Errors.FAIL,
                               err_info='Error Information').serialize())

        self.assertEqual(sh_fwd_input2.input, "Goodbye")

        self.backend_channel_wh.send(dmsg.LADumpState(tag=0).serialize())

        self.backend_channel_wh.send(dmsg.LADumpState(tag=0, filename="ladumptest.log").serialize())

        self.do_teardown(timeout=None)

        stdout = self.stdout.read()

        # '''(4, 'SUCCESS', None)
        # {'4': {'exe': 'python3', 'args': ['test.py'], 'env': {}, 'run_dir': '', 'user_name': '', 'create_rc': 0,
        # 'exit_code': None}}
        # response from send (<Errors.FAIL: 1>, 'Error Information')
        # True
        # [Exiting Launcher]'''

        stderr = self.stderr.read()

        self.assertIn('4', stdout)
        self.assertIn('SUCCESS', stdout)
        self.assertIn('True', stdout)
        self.assertEqual('', stderr.strip())


if __name__ == "__main__":
    unittest.main()

test/skipped_test_launcher.py
DragonHPC/dragon
#!/usr/bin/env python3


"""Test script for PMsgQueue class"""

import unittest
import dragon.infrastructure.messages as dmsg
import dragon.pmsgqueue as pq


class PMsgQueueUnitTest(unittest.TestCase):
    def test_create_send_recv(self):
        sendh = pq.PMsgQueue("/test", write_intent=True)
        recvh = pq.PMsgQueue("/test", read_intent=True)

        orig = dmsg.SHPingBE(tag=0, shep_cd="", be_cd="", gs_cd="", default_pd="", inf_pd="")
        sendh.send(orig.serialize())

        msg = recvh.recv()
        msg = dmsg.parse(msg)

        recvh.close()
        sendh.close(destroy=True)

        self.assertEqual(repr(msg), repr(orig))

    def test_bad_create(self):
        with self.assertRaises(pq.PMsgQueueOpenException):
            pq.PMsgQueue('test')

    def test_bad_create2(self):
        with self.assertRaises(pq.PMsgQueueConfigException):
            pq.PMsgQueue("test", write_intent=True)

    def test_bad_create3(self):
        with self.assertRaises(pq.PMsgQueueConfigException):
            pq.PMsgQueue("test", read_intent=True)

    def test_recv_send(self):
        recvh = pq.PMsgQueue("/test", read_intent=True)
        sendh = pq.PMsgQueue("/test", write_intent=True)

        orig = dmsg.SHPingBE(tag=0, shep_cd="", be_cd="", gs_cd="", default_pd="", inf_pd="")
        sendh.send(orig.serialize())

        msg = recvh.recv()
        msg = dmsg.parse(msg)

        sendh.close()
        recvh.close(destroy=True)

        self.assertEqual(repr(msg), repr(orig))

    def test_exceed_msg_size(self):
        sendh = pq.PMsgQueue("/test", write_intent=True)
        recvh = pq.PMsgQueue("/test", read_intent=True)

        orig = 'msg' * 10000
        with self.assertRaises(pq.PMsgQueueSendException):
            sendh.send(orig)

        recvh.close()
        sendh.close(destroy=True)

    def test_send_on_read_handle(self):
        sendh = pq.PMsgQueue("/test", write_intent=True)
        recvh = pq.PMsgQueue("/test", read_intent=True)

        orig = 'msg' * 10000
        with self.assertRaises(pq.PMsgQueueSendException):
            recvh.send(orig)

        recvh.close()
        sendh.close(destroy=True)

    def test_recv_on_send_handle(self):
        sendh = pq.PMsgQueue("/test", write_intent=True)
        recvh = pq.PMsgQueue("/test", read_intent=True)

        with self.assertRaises(pq.PMsgQueueReceiveException):
            sendh.recv()

        recvh.close()
        sendh.close(destroy=True)

    def test_send_on_closed_handle(self):
        sendh = pq.PMsgQueue("/test", write_intent=True)
        sendh.close()

        orig = dmsg.SHPingBE(tag=0, shep_cd="", be_cd="", gs_cd="", default_pd="", inf_pd="").serialize()

        with self.assertRaises(pq.PMsgQueueSendException):
            sendh.send(orig)

    def test_recv_on_closed_handle(self):
        recvh = pq.PMsgQueue("/test", read_intent=True)
        recvh.close()

        with self.assertRaises(pq.PMsgQueueReceiveException):
            recvh.recv()

    def test_timeout(self):
        recvh = pq.PMsgQueue("/test", read_intent=True)

        with self.assertRaises(pq.PMsgQueueTimeoutException):
            recvh.recv(timeout=1)

        recvh.close(destroy=True)


if __name__ == "__main__":
    unittest.main()

test/skipped_test_msgqueue.py
DragonHPC/dragon
"""Utility functions to support testing"""

import dragon.infrastructure.messages as dmsg


DEFAULT_TIMEOUT = 20


def get_and_parse(read_handle, timeout=DEFAULT_TIMEOUT):
    if hasattr(read_handle, "poll"):
        if read_handle.poll(timeout):
            the_msg = dmsg.parse(read_handle.recv())
        else:
            raise TimeoutError()
    else:
        the_msg = dmsg.parse(read_handle.recv())

    return the_msg


def get_and_check_type(read_handle, the_type, timeout=DEFAULT_TIMEOUT):
    the_msg = get_and_parse(read_handle, timeout=timeout)

    assert isinstance(the_msg, the_type), 'Wrong type. Expected: {}, got {}'.format(the_type, the_msg)

    return the_msg


def get_and_check_several(self, read_handle, expected_dict, timeout=DEFAULT_TIMEOUT):
    # self is the test object that is currently executing.

    rv = {x:[] for x in expected_dict}
    msgs_expected = sum(expected_dict.values())

    for k in range(msgs_expected):
        if read_handle.poll(timeout):
            the_msg = dmsg.parse(read_handle.recv())
        else:
            raise TimeoutError(f'Got {rv} messages so far but timed out while waiting for the rest in {expected_dict}.')

        msg_type = type(the_msg)
        self.assertIn(msg_type, expected_dict, f'Found an unexpected message {repr(the_msg)} while looking for one of these {expected_dict}')
        rv[msg_type].append((the_msg, k))

    for key in expected_dict:
        self.assertEqual(expected_dict[key], len(rv[key]),
            f'The message type {key} was expected {expected_dict[key]} times and was found {len(rv[key])} times.')

    return rv

def get_and_check_several_ignore_SHFwdOutput(self, read_handle, expected_dict, timeout=DEFAULT_TIMEOUT):
    # self is the test object that is currently executing.

    rv = {x:[] for x in expected_dict}
    rv[dmsg.SHFwdOutput] = []
    msgs_expected = sum(expected_dict.values())

    msgs_received = 0
    when_received = 0
    while msgs_received < msgs_expected:
        if read_handle.poll(timeout):
            the_msg = dmsg.parse(read_handle.recv())
        else:
            raise TimeoutError(f'Got {rv} messages so far but timed out while waiting for the rest in {expected_dict}.')

        msg_type = type(the_msg)
        if not isinstance(the_msg, dmsg.SHFwdOutput):
            self.assertIn(msg_type, expected_dict, f'Found an unexpected message {repr(the_msg)} while looking for one of these {expected_dict}')
            msgs_received+=1
        rv[msg_type].append((the_msg, when_received))
        when_received+=1

    for key in expected_dict:
        self.assertEqual(expected_dict[key], len(rv[key]),
            f'The message type {key} was expected {expected_dict[key]} times and was found {len(rv[key])} times.')

    return rv
test/support/util.py
DragonHPC/dragon
#!/usr/bin/env python3

import shim_dragon_paths
import subprocess
import os
from glob import glob

base_dir = os.getcwd()

# Directories to scan for tests in
test_dirs = [
    "broadcast",
    "channels_subtests",
    "hashtable",
    "pmod",
    "pmsgq",
    "utils"
]

if __name__ == '__main__':

    print()

    tests_run = 0
    tests_failed = []

    print()
    print('######## RUNNING C TESTS ########')
    for d in test_dirs:
        print()
        print(f"++++++++ {str(d).upper()} ++++++++")
        os.chdir(d) # Enter test dir
        subprocess.run(['make', 'clean'], capture_output=True)
        tests = glob('test_*.c') # Grab list of test files
        for t in tests:
            fname = os.path.splitext(os.path.basename(t))[0]
            print(f'Testing {fname} ...... ', end='', flush=True)

            tests_run += 1
            # If any of these fail do we want to skip to the next test or halt testing entirely?
            make_res = subprocess.run(['make', fname], capture_output=True) # Build test
            if make_res.returncode != 0:
                print("FAIL")
                print(make_res.stdout.decode('utf-8'))
                print(make_res.stderr.decode('utf-8'))
                tests_failed.append(d + " " + fname)
                continue
            t_out = subprocess.run(['./'+fname], capture_output=True) # Run test
            if t_out.returncode != 0:
                print("FAIL")
                print(t_out.stdout.decode('utf-8'))
                tests_failed.append(d + " --  " + fname)
            else:
                print("OK")

        os.chdir(base_dir) # Reset to base test dir

    total_passed = tests_run - len(tests_failed)
    print()
    print("######## C TESTS RESULTS ########")
    print(f'{total_passed} / {tests_run} PASSED')
    if len(tests_failed) != 0:
        print()
        print("!!!!!!!! Tests failed:")
        print()
        for t in tests_failed:
            print(t)

test/test_c_files.py
DragonHPC/dragon
#!/usr/bin/env python3

import shim_dragon_paths
import multiprocessing as mp
import unittest
from channels_subtests.test_basic_channels import ChannelCreateTest, ChannelTests

if __name__ == '__main__':
    mp.set_start_method("spawn", force=True)
    unittest.main()

test/test_channels.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest
from connection.connection_obj_tests import ConnectionTest
from connection.standalone_conn_test import TestStandaloneConn

if __name__ == "__main__":
    unittest.main()

test/test_connection.py
DragonHPC/dragon
#!/usr/bin/env python3

import os
import sys
import string
import random
import dragon
import unittest
import multiprocessing as mp

from dragon.data.distdictionary.dragon_dict import DragonDict

def do_client_ops(_dict, key, value):
    """Each client will do the operation of adding the sample key,value pair
    into the dictionary and again delete the same pair.

    :param _dict: dragon distributed dictionary
    :type _dict: dragon dictionary object
    :param key: Information about the SET operation to the manager
    :type key: Any hashable structure
    :param value: Value to be stored inside the manager
    :type value: Any structure that can be serialized
    """
    _dict[key] = value
    del _dict[key]
    # Close the client
    _dict.close()

def generate_keys():
    keys = list()
    # Generate the list of 100 keys to add to the dictionary
    dict_size = 100
    letters = string.ascii_letters

    for _ in range(dict_size):
        # each key is 20 characters long and characters can be repeated
        key = ''.join(random.choice(letters) for i in range(20))
        keys.append(key)
    return keys

def do_set_ops(_keys, _dict, value_size):
    """Each client will generate the value of given size, and perform
    the SET operation to all the keys of the dictionary

    :param _keys: list of all keys to be added to the dictionary
    :type _keys: list
    :param _dict: dragon distributed dictionary
    :type _dict: dragon dictionary object
    :param value_size: size of the value in bytes added to each dictionary key
    :type value_size: int
    """
    num_keys = len(_keys)
    letters = string.ascii_letters
    value = ''.join(random.choice(letters) for i in range(value_size))

    # Add each key to the dictionary
    for key in _keys:
        _dict[key] = value
    # Close the client
    _dict.close()

def do_del_ops(_keys, _dict):
    """Each client will delete the given list of keys in the dictionary
    
    :param _keys: list of all keys to be deleted from the dictionary
    :type _keys: list
    :param _dict: dragon distributed dictionary
    :type _dict: dragon dictionary object
    """
    for key in _keys:
        del _dict[key]
    # Close the client
    _dict.close()

def do_get_ops(_keys, _dict):
    """Each client will retrieve the values for the given list of keys in the dictionary
    
    :param _keys: list of all keys to be fetched from the dictionary
    :type _keys: list
    :param _dict: dragon distributed dictionary
    :type _dict: dragon dictionary object
    """
    for key in _keys:
        value = _dict[key]
    # Close the client
    _dict.close()


class TestDragonDictSingleNode(unittest.TestCase):
    @classmethod
    def setUpClass(self) -> None:
        # Create a dragon dictionary on a single node with multiple manager processes
        self.managers_per_node = 2 # 2 Managers per node
        self.num_nodes = 1 # Single Node
        self.total_mem_size = 2*(1024*1024*1024) # 2 GB total size
        self.DD = DragonDict(self.managers_per_node, self.num_nodes, self.total_mem_size)

    @classmethod
    def tearDownClass(self) -> None:
        self.DD.stop()

    def test_dict_params(self):
        # Verify the dictionary params
        num_nodes = self.DD._dist_dict.nodes
        managers_per_node = self.DD._dist_dict.managers_per_node
        num_managers = num_nodes * managers_per_node
        self.assertEqual(num_nodes, 1, "Number of dictionary nodes are incorrect")
        self.assertEqual(managers_per_node, 2, "Number of managers per node are incorrect")
        self.assertEqual(num_managers, 2, "Number of dictionary managers are incorrect")

    def test_set_and_get(self):
        key = "Who are you Dragon?"
        value = "A Fire Breathing Monster!"
        self.DD[key] = value
        # Compare if the fetched values are correct
        self.assertEqual(self.DD[key], value, "Retrieved value from the dict is incorrect")

    def test_setup_and_close_client(self):
        # Create a client process and pass the dictionary
        client_proc = mp.Process(target=do_client_ops, args=(self.DD, "Dragon", "Dictionary"))
        client_proc.start()
        client_proc.join()

    def test_set_ops(self):
        keys = generate_keys()
        # Create client procs that does dictionary operations
        num_clients = 8
        value_size = 64
        procs = []
        for i in range(num_clients):
            client_proc = mp.Process(target=do_set_ops, args=(keys, self.DD, value_size))
            client_proc.start()
            procs.append(client_proc)

        for i in range(len(procs)):
            procs[i].join()
            procs[i].kill()

    def test_dictionary_length(self):
        num_kv_pairs = 100
        # Assign key value pairs to the dictionary
        for i in range(0, num_kv_pairs):
            key = "Hello" + "." + str(i)
            value = "Dragon" + "." + str(i)
            self.DD[key] = value

        # Calculate the length of the dictionary
        dict_length = len(self.DD)
        # Verify that length of dictionary is equal to number of kv pairs
        self.assertEqual(dict_length, num_kv_pairs, "Calculated length of the dictionary is incorrect")

    def test_set_and_del_item(self):
        key = "Hello"
        value = "Dictionary"
        self.DD[key] = value
        # Delete operation should be successful
        del self.DD[key]

    def test_del_item_with_no_key(self):
        # Delete operation should be successful, without key present also
        # Log the information that the key is not present.
        letters = string.ascii_letters
        key = ''.join(random.choice(letters) for i in range(10))
        del self.DD[key]

    def test_del_ops(self):
        keys = list()
        num_kv_pairs = 100
        # Collect the set of keys for the dictionary
        for i in range(0, num_kv_pairs):
            key = "Hello" + "." + str(i)
            keys.append(key)

        # Two client processes delete keys in parallel
        # This might conduct extra deletes with keys deleted by other proc before
        num_clients = 2
        procs = []
        for i in range(num_clients):
            client_proc = mp.Process(target=do_del_ops, args=(keys, self.DD))
            client_proc.start()
            procs.append(client_proc)

        for i in range(len(procs)):
            procs[i].join()
            procs[i].kill()

    def test_keys(self):
        # Collect the keys from the dictionary
        keys = self.DD.keys()
        num_keys = len(keys)
        dict_length = len(self.DD)
        # Verify that length of dictionary is equal to number of kv pairs
        self.assertEqual(dict_length, num_keys, "Calculated length of the dictionary is incorrect")

    def test_get_ops(self):
        keys = self.DD.keys()
        # Create client procs that does dictionary operations
        num_clients = 8
        procs = []
        for i in range(num_clients):
            client_proc = mp.Process(target=do_get_ops, args=(keys, self.DD))
            client_proc.start()
            procs.append(client_proc)

        for i in range(len(procs)):
            procs[i].join()
            procs[i].kill()

    def test_existing_key(self):
        key = "Testing"
        value = "Key"
        self.DD[key] = value
        key_found = False
        # Verify if the key is present in the dictionary
        if key in self.DD:
            key_found = True
        self.assertEqual(key_found, True, "Key has not found in the dictionary")

    def test_not_existing_key(self):
        key = "Test"
        key_found = False
        # Verify if the key is present in the dictionary
        if key in self.DD:
            key_found = True
        self.assertEqual(key_found, False, "Key found in the dictionary")

    def test_pop_item(self):
        key = "Pop"
        value = "Item"
        self.DD[key] = value
        # Collect the value by popping the key
        pop_value = self.DD.pop(key)
        # Verify that the pop operation is correct
        self.assertEqual(value, pop_value, "Pop value is not returned correctly")

    def test_existing_value(self):
        key = "Testing"
        value = "Value"
        self.DD[key] = value
        value_found = False
        # Verify if the value is existing in the dictionary
        if value in self.DD.values():
            value_found = True
        self.assertEqual(value_found, True, "Value has not found in the dictionary")

    def test_not_existing_value(self):
        value = "Dummy Value"
        value_found = False
        # Verify if the value is existing in the dictionary
        if value in self.DD.values():
            value_found = True
        self.assertEqual(value_found, False, "Value found in the dictionary")

    def test_items(self):
        items = self.DD.items()
        for (key, value) in items:
            ret_value = self.DD[key]
            # Verify the value retrieved from the items with the specific key
            self.assertEqual(value, ret_value, "Value collected from the items is incorrect")

    def test_rename(self):
        key = "Hello"
        value = "Dictionary"
        self.DD[key] = value
        # Rename the above key
        new_key = "Renamed"
        self.DD.rename(key, new_key)
        ret_value = self.DD[new_key]
        # Verify that the key is properly renamed
        self.assertEqual(value, ret_value, "Rename key is not successful")


if __name__ == "__main__":
    mp.set_start_method("dragon")

    # Disable the non-deterministic behavior of the python hashing algorithm behavior
    hashseed = os.getenv('PYTHONHASHSEED')
    if not hashseed:
        os.environ['PYTHONHASHSEED'] = '0'
        os.execv(sys.executable, [sys.executable] + sys.argv)
    unittest.main()

test/test_dict.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest
from globalservices.test_refcounting import TestGSRefcounting
from globalservices.single_internal import SingleInternal
from globalservices.single_process_msg import SingleProcMsgChannels
from globalservices.process_api import SingleProcAPIChannels
from globalservices.group_api import GSGroupAPI

if __name__ == "__main__":
    unittest.main()

test/test_globalservices.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest
from infrastructure.env_parameter_tests import LaunchParameterTest
from infrastructure.newline_stream_wrapper_test import NewlineStreamWrapperTest


if __name__ == "__main__":
    unittest.main()

test/test_infrastructure.py
DragonHPC/dragon
#!/usr/bin/env python3

"""Simple gs-shep integration tests

These tests are a 'first light' connecting global services and the shepherd and
the global services API.
"""
import shim_dragon_paths

import inspect
import logging
import multiprocessing
import threading
import os
import pickle
import signal
import sys
import time
import unittest
import time
import uuid

import dragon.channels as dch
import dragon.managed_memory as dmm
import dragon.utils as dutils

import dragon.globalservices.api_setup as dapi
import dragon.globalservices.process as dproc
import dragon.globalservices.channel as dgchan
import dragon.globalservices.pool as dgpool

import dragon.infrastructure.process_desc as pdesc
import dragon.infrastructure.channel_desc as dgchan_desc
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.util as dutil

import dragon.localservices.local_svc as dsls
import dragon.dlogging.util as dlog

import support.util as tsu

# Dummy executable for testing:
# `slow_echo.sh time msg` sleeps for time seconds then prints msg.
# this should be the path from where the test driver (this script) is run.
import pathlib


# make sure we use a safe pool uid for creating a pool
# that will never overlap with infra/default/user pool uids
CUSTOM_POOL_MUID = 2 ** 60

path_to_slow_echo = pathlib.Path(__file__).parent / 'slow_echo.sh'
if not path_to_slow_echo.exists():
    path_to_slow_echo = pathlib.Path(__file__).parent / 'integration' / 'slow_echo.sh'
SLOW_ECHO = str(path_to_slow_echo)


def mk_sleep_msg(seconds, tag, p_uid, r_c_uid):
    return dmsg.GSProcessCreate(tag=tag, exe='/bin/sleep', args=[str(seconds)],
                                p_uid=p_uid, r_c_uid=r_c_uid)


def mk_slow_echo_msg(seconds, msg, tag, p_uid, r_c_uid):
    return dmsg.GSProcessCreate(tag=tag, exe=SLOW_ECHO, args=[str(seconds), msg],
                                p_uid=p_uid, r_c_uid=r_c_uid)


def start_ls(shep_stdin_queue, shep_stdout_queue, env_update, name_addition=''):
    dlog.setup_logging(basename='ls_' + name_addition, level=logging.DEBUG)
    log = logging.getLogger('ls-gs integration test channels, ls starts gs')
    log.info('--------------------ls started-----------------------')

    dparm.this_process.mode = dfacts.SINGLE_MODE
    os.environ.update(dparm.this_process.env())
    log.info('starting')
    try:
        dsls.single(ls_stdin=shep_stdin_queue, ls_stdout=shep_stdout_queue,
                    gs_args=[sys.executable, '-c', dfacts.GS_SINGLE_LAUNCH_CMD],
                    gs_env=env_update)
        log.info('exited')
    except RuntimeError as rte:
        log.exception('runtime failure\n')
        shep_stdout_queue.send(dmsg.AbnormalTermination(tag=dsls.get_new_tag(),
                                                        err_info=f'{rte!s}').serialize())
    finally:
        shep_stdin_queue.close()
        shep_stdout_queue.close()

    logging.shutdown()


class GSLS(unittest.TestCase):
    """ Global Services / Local Services (Shepherd) Integration tests
    """

    SLEEPY_HEAD_TIME = 60  # seconds, 60 seconds

    def setUp(self) -> None:
        self.tag = 0
        self.mode = dfacts.SINGLE_MODE

        self.shep_stdin_rh, self.shep_stdin_wh = multiprocessing.Pipe(duplex=False)
        self.shep_stdout_rh, self.shep_stdout_wh = multiprocessing.Pipe(duplex=False)

        self.pool_name = 'gslsinttest_' + os.environ.get('USER', str(os.getuid())) + str(uuid.uuid4())
        self.pool_size = 2 ** 30
        self.pool_uid = CUSTOM_POOL_MUID
        self.test_pool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid)

        # need a test pool and test channels to hold the return channel
        # from GS to the test bench.  Picking cuid for these high enough
        # that they won't be used in the tests

        self.gs_return_cuid = 2 ** 64 - CUSTOM_POOL_MUID
        self.proc_gs_return_chan = dch.Channel(self.test_pool, self.gs_return_cuid)
        self.proc_gs_return_rh = dconn.Connection(inbound_initializer=self.proc_gs_return_chan)

        self.chatty_teardown = False

    def next_tag(self):
        tmp = self.tag
        self.tag += 1
        return tmp

    def start_duts(self):
        test_name = self.__class__.__name__ + '_' + inspect.stack()[1][0].f_code.co_name

        # Make a list of channels to pass thru the environment to GS for testing purposes.
        test_chan_list = [(self.gs_return_cuid, self.test_pool.serialize(),
                           self.proc_gs_return_chan.serialize(), False)]

        val = pickle.dumps(test_chan_list)
        env_update = {dfacts.GS_TEST_CH_EV: dutils.B64.bytes_to_str(val),
                      dfacts.GS_LOG_BASE: 'gs_' + test_name}

        start_args = (self.shep_stdin_rh, self.shep_stdout_wh, env_update)

        self.ls_dut = multiprocessing.Process(target=start_ls, args=start_args,
                                              kwargs={'name_addition': test_name},
                                              name='local svc')

        self.ls_dut.start()

        self.shep_stdin_wh.send(dmsg.BENodeIdxSH(tag=self.next_tag(), node_idx=0).serialize())

        ping_be_msg = tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHPingBE)

        env2b = dutils.B64.str_to_bytes

        self.inf_pool = dmm.MemoryPool.attach(env2b(ping_be_msg.inf_pd))

        self.shep_input_chan = dch.Channel.attach(env2b(ping_be_msg.shep_cd))
        self.shep_input_wh = dconn.Connection(outbound_initializer=self.shep_input_chan)

        self.bela_input_chan = dch.Channel.attach(env2b(ping_be_msg.be_cd))
        self.bela_input_rh = dconn.Connection(inbound_initializer=self.bela_input_chan)

        self.gs_input_chan = dch.Channel.attach(env2b(ping_be_msg.gs_cd))
        self.gs_input_wh = dconn.Connection(outbound_initializer=self.gs_input_chan)

        dapi.test_connection_override(test_gs_input=self.gs_input_wh,
                                      test_gs_return=self.proc_gs_return_rh,
                                      test_gs_return_cuid=self.gs_return_cuid,
                                      test_shep_input=self.shep_input_wh)

        self.shep_input_wh.send(dmsg.BEPingSH(tag=self.next_tag()).serialize())

        sh_msg = tsu.get_and_check_type(self.bela_input_rh, dmsg.SHChannelsUp)
        self.assertIsInstance(sh_msg, dmsg.SHChannelsUp)
        self.assertEqual(dutils.host_id(), sh_msg.host_id)
        if self.chatty_teardown:
            print(f"info from SHChannelsUp: host_id={sh_msg.host_id}, hostname={sh_msg.host_name}")
            sys.stdout.flush()

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSIsUp)

        self.starting_shm = dutil.survey_dev_shm()

    def tearDown(self) -> None:
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self.next_tag()).serialize())
        if self.chatty_teardown:
            print('sent GSTeardown')
            sys.stdout.flush()

        tsu.get_and_check_several_ignore_SHFwdOutput(self, self.bela_input_rh, {dmsg.GSHalted: 1})
        if self.chatty_teardown:
            print('got GSHalted')
            sys.stdout.flush()

        self.shep_input_wh.send(dmsg.SHTeardown(tag=self.next_tag()).serialize())
        if self.chatty_teardown:
            print('sent SHTeardown')
            sys.stdout.flush()

        tsu.get_and_check_several_ignore_SHFwdOutput(self, self.bela_input_rh, {dmsg.SHHaltBE: 1})
        if self.chatty_teardown:
            print('got SHHaltBE')
            sys.stdout.flush()

        self.shep_input_chan.detach()
        self.bela_input_chan.detach()
        self.gs_input_chan.detach()

        self.proc_gs_return_chan.destroy()

        self.inf_pool.detach()

        self.shep_stdin_wh.send(dmsg.BEHalted(tag=self.next_tag()).serialize())

        tsu.get_and_check_several_ignore_SHFwdOutput(self, self.shep_stdout_rh, {dmsg.SHHalted: 1})
        self.shep_stdin_rh.close()
        self.shep_stdin_wh.close()
        self.shep_stdout_rh.close()
        self.shep_stdout_wh.close()

        self.ls_dut.join()
        self.test_pool.destroy()
        dutil.compare_dev_shm(self.starting_shm)

    def kill_sleepy_head(self):
        msg = dmsg.GSProcessKill(tag=self.next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                                 r_c_uid=dfacts.BASE_BE_CUID,
                                 sig=signal.SIGKILL, t_p_uid=self.head_puid)
        self.gs_input_wh.send(msg.serialize())

        # Since the backend can still be receiving messages from LS,
        # discard its SHFwdOutput messages that still may be in the channel
        # queue
        msg = tsu.get_and_parse(self.bela_input_rh)
        while (isinstance(msg, dmsg.SHFwdOutput)):
            msg = tsu.get_and_parse(self.bela_input_rh)

        if isinstance(msg, dmsg.GSProcessKillResponse):
            tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit)
        else:
            self.assertIsInstance(msg, dmsg.GSHeadExit)
            tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessKillResponse)

    def start_sleepy_head(self, timeout):
        msg = mk_sleep_msg(timeout, self.next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                           r_c_uid=dfacts.BASE_BE_CUID)

        self.gs_input_wh.send(msg.serialize())

        resp = tsu.get_and_check_type(self.bela_input_rh, dmsg.GSProcessCreateResponse)

        self.assertEqual(resp.err, dmsg.GSProcessCreateResponse.Errors.SUCCESS)

        self.head_puid = resp.desc.p_uid
        dparm.this_process.my_puid = self.head_puid

    def test_proc_direct(self):

        self.start_duts()
        echotxt = 'Hello World'

        msg = dmsg.GSProcessCreate(tag=self.next_tag(), exe='/bin/echo',
                                args=[echotxt], p_uid=dfacts.LAUNCHER_PUID,
                                r_c_uid=dfacts.BASE_BE_CUID)

        self.gs_input_wh.send(msg.serialize())

        # Expect a GSPCR, a SHFwdOutput, and a GSHeadExit message where
        # the output could be in any order but the PCR will precede
        # the HeadExit.

        msgs = tsu.get_and_check_several_ignore_SHFwdOutput(self, self.bela_input_rh,
                                        {dmsg.GSProcessCreateResponse: 1,
                                        dmsg.GSHeadExit: 1})

        pcr_msg, _ = msgs[dmsg.GSProcessCreateResponse][0]
        self.assertEqual(pcr_msg.err, dmsg.GSProcessCreateResponse.Errors.SUCCESS)

        if len(msgs[dmsg.SHFwdOutput]) > 0:
            output_msg, _ = msgs[dmsg.SHFwdOutput][0]
            self.assertEqual(echotxt + '\n', output_msg.data)

    def test_sleepy_exit(self):
        self.start_duts()

        self.start_sleepy_head(1)

        tsu.get_and_check_type(self.bela_input_rh, dmsg.GSHeadExit, timeout=5)

    def test_sleepy_echo_head(self):
        self.start_duts()

        output = 'hyenas are awesome'
        msg = mk_slow_echo_msg(1, output, self.next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                               r_c_uid=dfacts.BASE_BE_CUID)
        self.gs_input_wh.send(msg.serialize())

        msgs = tsu.get_and_check_several_ignore_SHFwdOutput(self, self.bela_input_rh, {dmsg.GSProcessCreateResponse: 1,
                                                              dmsg.GSHeadExit: 1})

        create_msg, create_order = msgs[dmsg.GSProcessCreateResponse][0]
        self.assertEqual(create_msg.err, dmsg.GSProcessCreateResponse.Errors.SUCCESS)

        if len(msgs[dmsg.SHFwdOutput]) > 0:
            output_msg, _ = msgs[dmsg.SHFwdOutput][0]
            self.assertEqual(output + '\n', output_msg.data)


    def test_head_kill(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)
        self.kill_sleepy_head()

    def test_api_create(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)
        res = dproc.create(exe='/bin/sleep', run_dir='', args=['1'], env={}, user_name='hproc')
        self.assertEqual(res.name, 'hproc')
        self.kill_sleepy_head()

    @unittest.skip('derp')
    def test_api_create_output(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        res = dproc.create(exe=SLOW_ECHO, run_dir='', args=['10', 'hyena'],
                           env={}, user_name='hproc')
        self.assertEqual(res.name, 'hproc')

        resp = tsu.get_and_check_type(self.bela_input_rh, dmsg.SHFwdOutput, 3)
        self.assertEqual('hyena\n', resp.data)
        self.kill_sleepy_head()

    def test_api_join(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        res = dproc.create(exe='/bin/sleep', run_dir='', args=['1'], env={})

        dproc.join(res.p_uid)

        self.kill_sleepy_head()

    # will print warnings if the pool isn't cleaned up in
    # teardown.  Would like a better way to do this but
    # TODO: refactor the test so it fails properly if this does not happen
    def test_pool_and_cleanup(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        _ = dgpool.create(2 ** 20)

        self.kill_sleepy_head()

    def test_list(self):
        self.start_duts()

        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        res = dproc.create(exe='/bin/sleep', run_dir='', args=['1'], env={})
        procs = dproc.get_list()
        self.assertEqual([res.p_uid-1, res.p_uid], procs)

        self.kill_sleepy_head()

    def test_inf_channel_basic(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        res = dproc.create(exe='python3', run_dir='', env={},
                           args=['globalservices/managed_proc_test_target.py'])
        dproc.join(res.p_uid)
        dead_pd = dproc.query(res.p_uid)
        self.kill_sleepy_head()

        self.assertEqual(dead_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(dead_pd.ecode, 0)

    def test_inf_channel_gsr(self):
        self.start_duts()

        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        res = dproc.create(exe='python3', run_dir='',
                           args=['globalservices/managed_proc_test_target.py', '--gs-ret'],
                           env={},
                           options=pdesc.ProcessOptions(make_inf_channels=True))

        dproc.join(res.p_uid)

        dead_pd = dproc.query(res.p_uid)
        gsr_desc = dgchan.wait(dead_pd.gs_ret_cuid, timeout=3)

        self.kill_sleepy_head()

        self.assertEqual(dead_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(dead_pd.ecode, 0)
        self.assertEqual(gsr_desc.state, gsr_desc.State.DEAD)

    def test_inf_channels_nullarg(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        res = dproc.create(exe='python3', run_dir='',
                           args=['globalservices/managed_proc_test_target.py', '--args'],
                           env={},
                           options=pdesc.ProcessOptions(make_inf_channels=True))

        dproc.join(res.p_uid)
        dead_pd = dproc.query(res.p_uid)
        self.kill_sleepy_head()

        self.assertEqual(dead_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(dead_pd.ecode, 0)

    def test_inf_channels_nullarg_twice(self):
        self.start_duts()

        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        res = dproc.create(exe=sys.executable, run_dir='',
                           args=['globalservices/managed_proc_test_target.py', '--args'],
                           env={},
                           options=pdesc.ProcessOptions(make_inf_channels=True))

        res2 = dproc.create(exe=sys.executable, run_dir='',
                            args=['globalservices/managed_proc_test_target.py', '--args'],
                            env={},
                            options=pdesc.ProcessOptions(make_inf_channels=True))

        dproc.join(res.p_uid)
        dproc.join(res2.p_uid)

        dead_pd = dproc.query(res.p_uid)
        dead_pd2 = dproc.query(res2.p_uid)

        self.kill_sleepy_head()

        self.assertEqual(dead_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(dead_pd.ecode, 0)
        self.assertEqual(dead_pd2.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(dead_pd2.ecode, 0)

    def test_inf_channels_immediate_arg(self):
        self.start_duts()

        # Checks that the args got there.
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        test_script_args = ['globalservices/managed_proc_test_target.py',
                            '--args',
                            '23',
                            '45',
                            '17']

        the_argdata = pickle.dumps((23, 45, 17))
        res = dproc.create_with_argdata(exe=sys.executable, run_dir='',
                                        args=test_script_args,
                                        env={}, argdata=the_argdata)

        dproc.join(res.p_uid)

        dead_pd = dproc.query(res.p_uid)

        self.kill_sleepy_head()

        self.assertEqual(dead_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(dead_pd.ecode, 0)

    def test_inf_channels_jumbo_arg(self):
        self.start_duts()

        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)
        arg_size = 2 ** 25

        test_script_args = ['globalservices/jumbo_arg_test_target.py',
                            str(arg_size)]

        res = dproc.create_with_argdata(exe=sys.executable, run_dir='',
                                        args=test_script_args,
                                        env={}, argdata=bytearray(arg_size))

        dproc.join(res.p_uid)

        dead_pd = dproc.query(res.p_uid)

        self.kill_sleepy_head()

        self.assertEqual(dead_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(dead_pd.ecode, 0)

    def test_connection_and_api(self):
        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        first_cn = 'bob'
        second_cn = 'fred'

        my_muid = dfacts.default_pool_muid_from_index(0)

        dgchan.create(my_muid, first_cn)
        dgchan.create(my_muid, second_cn)

        first_args = ['connection/conn_arg_test_target.py', 'first']
        second_args = ['connection/conn_arg_test_target.py', 'second']

        delivered_args = (first_cn, second_cn)
        the_argdata = pickle.dumps(delivered_args)
        first_pd = dproc.create_with_argdata(exe=sys.executable, run_dir='',
                                             args=first_args,
                                             env={}, argdata=the_argdata)

        delivered_args = (second_cn, first_cn)
        the_argdata = pickle.dumps(delivered_args)
        second_pd = dproc.create_with_argdata(exe=sys.executable, run_dir='',
                                              args=second_args,
                                              env={}, argdata=the_argdata)

        dproc.join(first_pd.p_uid)
        dproc.join(second_pd.p_uid)

        # refresh the descriptors
        first_pd = dproc.query(first_pd.p_uid)
        second_pd = dproc.query(second_pd.p_uid)

        dgchan.destroy(first_cn)
        dgchan.destroy(second_cn)

        self.kill_sleepy_head()

        self.assertEqual(first_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(first_pd.ecode, 0)

        self.assertEqual(second_pd.state, pdesc.ProcessDescriptor.State.DEAD)
        self.assertEqual(second_pd.ecode, 0)

    def test_channel_refcount(self):
        """PE-38654 Test ref counted channels and channel parameters
        """

        self.start_duts()
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        # Create a refcounted user channel
        m_uid = dfacts.default_pool_muid_from_index(0)

        the_options = dgchan_desc.ChannelOptions(ref_count=True)
        the_name = "test_channel_refcount-"+os.environ.get('USER', str(os.getuid()))

        descr = dgchan.create(m_uid, the_name)

        # get identifiers
        c_uid = descr.c_uid
        c_name= descr.name

        # Functions for the threads
        def call_get_refcnt(name, c_uid):
            dgchan.get_refcnt(name)
            dgchan.get_refcnt(c_uid)

        def call_release_refcnt(name, c_uid):
            dgchan.release_refcnt(name)
            dgchan.release_refcnt(c_uid)

        # Petal to the metal
        maxproc = 16

        # Create a few threads and have them run get_refcnt()
        threads = [ threading.Thread(target=call_get_refcnt, args=(c_name, c_uid)) for i in range(0,maxproc)]

        for th in threads:
            th.start()

        for th in threads:
            th.join()

        # See what's going on with the channel
        new_descr = dgchan.query(c_uid, inc_refcnt=False)

        self.assertEqual(descr.sdesc, new_descr.sdesc) # at the least ...

        # Create a bunch of threads have them run release_refcnt()
        threads = [ threading.Thread(target=call_release_refcnt, args=(c_name, c_uid)) for i in range(0,maxproc)]

        for th in threads:
            th.start()

        for th in threads:
            th.join()

        # check that channel status is DEAD, but give Dragon some time.
        final_descr = dgchan.wait(c_uid, timeout=3)

        self.assertEqual(final_descr.state, final_descr.State.DEAD)

        # cleanup
        self.kill_sleepy_head()


def start_multi_ls(shep_stdin_queue, shep_stdout_queue, env_update, name_addition=''):
    dparm.this_process.mode = dfacts.CRAYHPE_MULTI_MODE
    dparm.this_process.debug = '1'
    os.environ.update(dparm.this_process.env())
    try:
        dsls.multinode(ls_stdin=shep_stdin_queue, ls_stdout=shep_stdout_queue,
                       gs_args=[sys.executable, '-c', dfacts.GS_MULTI_LAUNCH_CMD],
                       gs_env=env_update, fname='LS_' + name_addition,
                       ta_args=[sys.executable, '-c', dfacts.TA_MULTI_LAUNCH_CMD])
    except RuntimeError as rte:
        shep_stdout_queue.send(dmsg.AbnormalTermination(tag=dsls.get_new_tag(),
                                                        err_info=f'{rte!s}').serialize())
    finally:
        shep_stdin_queue.close()
        shep_stdout_queue.close()


class GSLSMulti(unittest.TestCase):
    """ Global Services / Local Services (Shepherd) Integration tests
    where GS process talks to two or more LS processes on the same node
    """

    SLEEPY_HEAD_TIME = 60  # seconds, 60 seconds

    def setUp(self) -> None:
        self.nls = 3 # number of LS processes
        self.tag = 0
        self.mode = dfacts.CRAYHPE_MULTI_MODE
        self.names = {} # helper dict that holds all the necessary names for connections for the different LS processes

        for i in range(self.nls):
            self.names[f'shep_stdin_rh_{i}'], self.names[f'shep_stdin_wh_{i}'] = multiprocessing.Pipe(duplex=False)
            self.names[f'shep_stdout_rh_{i}'], self.names[f'shep_stdout_wh_{i}'] = multiprocessing.Pipe(duplex=False)

        # need a test pool and test channels to hold the return channel
        # from GS to the test bench.  Picking cuid for these high enough
        # that they won't be used in the tests
        self.pool_name = 'gslsinttest_' + os.environ.get('USER', str(os.getuid())) + str(uuid.uuid4())
        self.pool_size = 2 ** 30
        self.pool_uid = CUSTOM_POOL_MUID
        self.test_pool = dmm.MemoryPool(self.pool_size, self.pool_name, self.pool_uid)

        self.gs_return_cuid = 2 ** 64 - CUSTOM_POOL_MUID
        self.proc_gs_return_chan = dch.Channel(self.test_pool, self.gs_return_cuid)
        self.proc_gs_return_rh = dconn.Connection(inbound_initializer=self.proc_gs_return_chan)

        self.chatty_teardown = False

        self.gs_cd = None

    def next_tag(self):
        tmp = self.tag
        self.tag += 1
        return tmp

    def start_multi_duts(self):
        for i in range(self.nls):
            test_name = f'{i}_{self.__class__.__name__}_{inspect.stack()[1][0].f_code.co_name}'
            if i == 0: # primary LS
                # Make a list of channels to pass thru the environment to GS for testing purposes.
                test_chan_list = [(self.gs_return_cuid, self.test_pool.serialize(),
                                   self.proc_gs_return_chan.serialize(), False)]
                val = pickle.dumps(test_chan_list)
                env_update = {dfacts.GS_TEST_CH_EV: dutils.B64.bytes_to_str(val),
                              dfacts.GS_LOG_BASE: 'gs_' + test_name}
                start_args = (self.names[f'shep_stdin_rh_{i}'], self.names[f'shep_stdout_wh_{i}'], env_update)
            else:
                start_args = (self.names[f'shep_stdin_rh_{i}'], self.names[f'shep_stdout_wh_{i}'], None)

            self.names[f'ls_dut_{i}'] = multiprocessing.Process(target=start_multi_ls, args=start_args,
                                                                kwargs={'name_addition': test_name},
                                                                name=f'local_svc_{i}')
            self.names[f'ls_dut_{i}'].start()

        if self.chatty_teardown:
            print(f'The {self.nls} local services have started', flush=True)

        chs_up = []
        for i in range(self.nls):
            if i == 0:
                self.names[f'shep_stdin_wh_{i}'].send(dmsg.BENodeIdxSH(tag=self.next_tag(), node_idx=i, primary=True).serialize())
            else:
                self.names[f'shep_stdin_wh_{i}'].send(dmsg.BENodeIdxSH(tag=self.next_tag(), node_idx=i, primary=False).serialize())

            ping_be_msg = tsu.get_and_check_type(self.names[f'shep_stdout_rh_{i}'], dmsg.SHPingBE)
            if self.chatty_teardown:
                print(f'got SHPingBE for LS {i}', flush=True)

            env2b = dutils.B64.str_to_bytes

            self.names[f'shep_input_chan_{i}'] = dch.Channel.attach(env2b(ping_be_msg.shep_cd))
            self.names[f'shep_input_wh_{i}'] = dconn.Connection(outbound_initializer=self.names[f'shep_input_chan_{i}'])


            self.names[f'bela_input_chan_{i}'] = dch.Channel.attach(env2b(ping_be_msg.be_cd))
            self.names[f'bela_input_rh_{i}'] = dconn.Connection(inbound_initializer=self.names[f'bela_input_chan_{i}'])

            if i == 0:
                self.gs_cd = ping_be_msg.gs_cd
                self.gs_input_chan = dch.Channel.attach(env2b(ping_be_msg.gs_cd))
                self.gs_input_wh = dconn.Connection(outbound_initializer=self.gs_input_chan)

                dapi.test_connection_override(test_gs_input=self.gs_input_wh,
                                              test_gs_return=self.proc_gs_return_rh,
                                              test_gs_return_cuid=self.gs_return_cuid,
                                              test_shep_input=self.names[f'shep_input_wh_{i}'])

            self.names[f'shep_input_wh_{i}'].send(dmsg.BEPingSH(tag=self.next_tag()).serialize())

            msg = tsu.get_and_check_type(self.names[f'bela_input_rh_{i}'], dmsg.SHChannelsUp)
            assert isinstance(msg, dmsg.SHChannelsUp)
            self.assertEqual(dutils.host_id(), msg.host_id)
            if self.chatty_teardown:
                print(f'got SHChannelsUp from LS {i}', flush=True)
                print(f'Info from SHChannelsUp: host_id={msg.host_id}, hostname={msg.host_name}', flush=True)

            chs_up.append(msg)

        la_ch_info = dmsg.LAChannelsInfo(tag=self.next_tag(), nodes_desc=chs_up,
                                         gs_cd=self.gs_cd, num_gw_channels=0)
        for i in range(self.nls):
            # communicate gs_cd to all local services process so that they
            # can communicate with gs
            self.names[f'shep_input_wh_{i}'].send(la_ch_info.serialize())
            if self.chatty_teardown:
                print(f'sent LACHannelsInfo to LS {i}', flush=True)

            # we need to receive TAUp
            tsu.get_and_check_type(self.names[f'bela_input_rh_{i}'], dmsg.TAUp)
            if self.chatty_teardown:
                print(f'received TAUp from LS {i}')
                sys.stdout.flush()

        tsu.get_and_check_type(self.names['bela_input_rh_0'], dmsg.GSIsUp)
        if self.chatty_teardown:
            print('received GSIsUp')
            sys.stdout.flush()

        self.starting_shm = dutil.survey_dev_shm()


    def tearDown(self) -> None:
        self.gs_input_wh.send(dmsg.GSTeardown(tag=self.next_tag()).serialize())
        if self.chatty_teardown:
            print('sent GSTeardown')
            sys.stdout.flush()

        tsu.get_and_check_type(self.names['bela_input_rh_0'], dmsg.GSHalted)
        if self.chatty_teardown:
            print('got GSHalted')
            sys.stdout.flush()

        self.names['shep_input_wh_0'].send(dmsg.SHHaltTA(tag=self.next_tag()).serialize())
        if self.chatty_teardown:
            print('sent SHHaltTA')
            sys.stdout.flush()

        tsu.get_and_check_type(self.names['bela_input_rh_0'], dmsg.TAHalted)
        if self.chatty_teardown:
            print('got TAHalted')
            sys.stdout.flush()

        for i in range(self.nls):
            self.names[f'shep_input_wh_{i}'].send(dmsg.SHTeardown(tag=self.next_tag()).serialize())
            if self.chatty_teardown:
                print(f'sent SHTeardown to LS {i}')
                sys.stdout.flush()

            tsu.get_and_check_type(self.names[f'bela_input_rh_{i}'], dmsg.SHHaltBE)
            if self.chatty_teardown:
                print(f'got SHHaltBE from LS {i}')
                sys.stdout.flush()

        for i in range(self.nls):
            self.names[f'shep_input_chan_{i}'].detach()
            self.names[f'bela_input_chan_{i}'].detach()

            self.names[f'shep_stdin_wh_{i}'].send(dmsg.BEHalted(tag=self.next_tag()).serialize())
            if self.chatty_teardown:
                print(f'sent BEHalted to LS {i}')
                sys.stdout.flush()

        self.gs_input_chan.detach()
        self.proc_gs_return_chan.destroy()

        for i in range(self.nls):
            self.names[f'shep_stdin_rh_{i}'].close()
            self.names[f'shep_stdin_wh_{i}'].close()
            self.names[f'shep_stdout_rh_{i}'].close()
            self.names[f'shep_stdout_wh_{i}'].close()

        for i in range(self.nls):
            self.names[f'ls_dut_{i}'].join()

        self.test_pool.destroy()
        dutil.compare_dev_shm(self.starting_shm)

    def kill_sleepy_head(self):
        msg = dmsg.GSProcessKill(tag=self.next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                                 r_c_uid=dfacts.BASE_BE_CUID,
                                 sig=signal.SIGKILL, t_p_uid=self.head_puid)
        self.gs_input_wh.send(msg.serialize())

        # Since the backend can still be receiving messages from LS,
        # discard its SHFwdOutput messages that still may be in the channel
        # queue
        msg = tsu.get_and_parse(self.names['bela_input_rh_0'])
        while (isinstance(msg, dmsg.SHFwdOutput)):
            msg = tsu.get_and_parse(self.names['bela_input_rh_0'])

        if isinstance(msg, dmsg.GSProcessKillResponse):
            tsu.get_and_check_type(self.names['bela_input_rh_0'], dmsg.GSHeadExit)
        else:
            self.assertIsInstance(msg, dmsg.GSHeadExit)
            tsu.get_and_check_type(self.names['bela_input_rh_0'], dmsg.GSProcessKillResponse)

    def start_sleepy_head(self, timeout):
        msg = mk_sleep_msg(timeout, self.next_tag(), p_uid=dfacts.LAUNCHER_PUID,
                           r_c_uid=dfacts.BASE_BE_CUID)

        self.gs_input_wh.send(msg.serialize())

        resp = tsu.get_and_check_type(self.names['bela_input_rh_0'], dmsg.GSProcessCreateResponse)

        self.assertEqual(resp.err, dmsg.GSProcessCreateResponse.Errors.SUCCESS)

        self.head_puid = resp.desc.p_uid
        dparm.this_process.my_puid = self.head_puid

    @unittest.skip("The teardown behavior is incorrect since SHHaltTA is sent to each node, not just primary.")
    def test_gs_ls_comm(self):
        self.start_multi_duts()
        if self.chatty_teardown:
            print('Everything is set up. Ready to fire up a few processes.', flush=True)
        self.start_sleepy_head(self.SLEEPY_HEAD_TIME)

        # create processes
        # GS will assign the creation of 2 processes to each LS
        # in a cyclic order
        procs = []
        for i in range(2 * self.nls):
            res = dproc.create(exe='/bin/sleep', run_dir='', args=['1'], env={})
            procs.append(res.p_uid)

        dproc.multi_join(procs, join_all=True)

        self.kill_sleepy_head()


if __name__ == "__main__":
    multiprocessing.set_start_method('spawn')
    unittest.main()

test/test_integration_shep_gs.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest
import logging

from launcher.test_launch_options import LaunchSelectionTest, LaunchOptionsTest
from launcher.test_network_config import FileNetworkConfigTest, SlurmNetworkConfigTest, PbsPalsNetworkConfigTest
from launcher.test_network_config import SSHNetworkConfigTest
from launcher.test_signal_handling import SigIntTest
from launcher.test_frontend_bringup import FrontendBringUpTeardownTest
from launcher.test_backend_bringup import BackendBringUpTeardownTest


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    unittest.main()
test/test_launcher.py
DragonHPC/dragon
#!/usr/bin/env python3
"""Test script for single node shepherd message handling"""

import logging
import os
import socket
import subprocess
import sys
import unittest

from dragon.cli import console_script_args
import dragon.channels as dch
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.facts as dfacts
import dragon.launcher.util as dlutil
import dragon.launcher.pmsgqueue as dpq
import dragon.dlogging.util as dlog
from dragon.utils import B64

import support.util as tsu

DEFAULT_TIMEOUT = 1

get_msg = tsu.get_and_parse


def run_ls(test_env=None):
    dlog.setup_logging(basename="ls_multi_", level=logging.DEBUG, force=True)
    log = logging.getLogger('run ls')
    log.info('starting ls from run')

    if test_env is None:
        the_env = os.environ
    else:
        the_env = dict(os.environ)
        the_env.update(test_env)

    try:
        ls = subprocess.Popen(console_script_args(dfacts.PROCNAME_LS), env=the_env)
    except RuntimeError as rte:
        # shep_stdout_queue.send(dmsg.AbnormalTermination(tag=dsls.get_new_tag(),
        #                                                err_info=f'{rte!s}').serialize())
        print("something here")
        raise rte
    return ls


class MultiLS(unittest.TestCase):
    DEFAULT_TEST_POOL_SIZE = 2 ** 20

    def setUp(self) -> None:
        # Create posix msg queue for launcher comm
        _user = os.environ.get('USER', str(os.getuid()))
        self.hostname = socket.gethostname()
        self.ip_addr = socket.gethostbyname(self.hostname)
        self.pq_stdin = dpq.PMsgQueue(f'/{_user}_{self.hostname}_la_stdout', write_intent=True)
        self.pq_stdout = dpq.PMsgQueue(f'/{_user}_{self.hostname}_la_stdin', read_intent=True)

        self.dragon_logger = dlog.setup_dragon_logging(node_index=0)
        self.logger_sdesc = self.dragon_logger.serialize()

        self.channels_connected = False
        self.ls = None
        self.tag = 0

    def tearDown(self) -> None:
        self.pq_stdin.close(destroy=True)
        self.pq_stdout.close(destroy=True)

        if (self.channels_connected):
            self.ls_channel.detach()
            self.la_channel.detach()

        if self.ls is not None:
            self.ls.wait()

        self.dragon_logger.destroy()

    def next_tag(self):
        tmp = self.tag
        self.tag += 1
        return tmp

    def do_bringup(self):
        sh_ping_be_msg = self.do_BENodeIdxSH_SHPingBE(test_env={dfacts.TRANSPORT_TEST_ENV: '1'},
                                                      node_idx=0)
        self.connect_to_ls_channels(sh_ping_be_msg)
        sh_channels_up = self.do_BEPingSH_SHChannelsUp()
        self.do_TAUp(sh_channels_up)

    def do_teardown(self):
        self.do_SHHaltTA()
        self.do_SHTeardown()

    def do_BENodeIdxSH_SHPingBE(self, test_env=None, node_idx=0):

        self.ls = run_ls(test_env=test_env)

        # Send BENodeIdxSH as if I were launcher backend
        self.pq_stdin.send(dmsg.BENodeIdxSH(tag=self.next_tag(), node_idx=node_idx, ip_addrs=[self.ip_addr],
                                            host_name=self.hostname, primary=(node_idx == 0),
                                            logger_sdesc=B64(self.logger_sdesc)).serialize())
        return tsu.get_and_check_type(self.pq_stdout, dmsg.SHPingBE)

    def connect_to_ls_channels(self, sh_ping_be):

        self.ls_channel = dch.Channel.attach(B64.from_str(sh_ping_be.shep_cd).decode())
        self.la_channel = dch.Channel.attach(B64.from_str(sh_ping_be.be_cd).decode())

        # IN test transport mode this is null
        try:
            self.gs_cd = sh_ping_be.gs_cd
        except TypeError:
            self.gs_cd = None

        self.ls_queue = dconn.Connection(outbound_initializer=self.ls_channel, policy=dparm.POLICY_INFRASTRUCTURE)
        self.la_queue = dconn.Connection(inbound_initializer=self.la_channel, policy=dparm.POLICY_INFRASTRUCTURE)

        self.channels_connected = True

    def do_BEPingSH_SHChannelsUp(self):

        self.ls_queue.send(dmsg.BEPingSH(tag=self.next_tag()).serialize())
        return tsu.get_and_check_type(self.la_queue, dmsg.SHChannelsUp)

    def do_TAUp(self, sh_channels_msg):

        num_gw_channels = dfacts.DRAGON_DEFAULT_NUM_GW_CHANNELS_PER_NODE

        # Send LAChannelsInfo in a test environment or to all
        la_ch_info = dmsg.LAChannelsInfo(tag=self.next_tag(), nodes_desc=[sh_channels_msg],
                                         gs_cd=self.gs_cd, num_gw_channels=num_gw_channels)
        self.ls_queue.send(la_ch_info.serialize())
        tsu.get_and_check_type(self.la_queue, dmsg.TAUp)

    def do_SHHaltTA(self):

        self.ls_queue.send(dmsg.SHHaltTA(tag=self.next_tag()).serialize())
        tsu.get_and_check_type(self.la_queue, dmsg.TAHalted)

    def do_SHTeardown(self):

        self.ls_queue.send(dmsg.SHTeardown(tag=self.next_tag()).serialize())
        tsu.get_and_check_type(self.la_queue, dmsg.SHHaltBE)
        self.pq_stdin.send(dmsg.BEHalted(tag=self.next_tag()).serialize())


    def test_bringup_teardown_no_gs_primary(self):
        '''Transport test mode with no global services and is primary node

        Teardown of LS will begin with SHProcessExit receipt and
        transmission of TAHalted
        '''
        log = logging.getLogger("test_bringup_teardown_no_gs_primary")

        self.do_bringup()

        proc_create_msg = dlutil.mk_shproc_echo_msg()
        self.ls_queue.send(proc_create_msg.serialize())
        tsu.get_and_check_type(self.la_queue, dmsg.SHProcessCreateResponse)

        # Should get a forwarded output
        tsu.get_and_check_type(self.la_queue, dmsg.SHFwdOutput)
        tsu.get_and_check_type(self.la_queue, dmsg.SHProcessExit)

        self.do_teardown()


    def test_bringup_abnormal_termination(self):
        log = logging.getLogger("test_bringup_abnormal_termination")

        self.do_bringup()

        self.ls_queue.send(dmsg.AbnormalTermination(tag=self.next_tag()).serialize())
        tsu.get_and_check_type(self.la_queue, dmsg.AbnormalTermination)

        self.do_teardown()


    def test_bringup_bad_message(self):
        log = logging.getLogger("test_bringup_bad_message")

        self.do_bringup()

        self.ls_queue.send("crap")
        log.info("sent bad message to ls")
        tsu.get_and_check_type(self.la_queue, dmsg.AbnormalTermination)
        log.info("received AbnormalTermination from ls")

        self.do_teardown()


    def test_bringup_unexpected_message(self):
        log = logging.getLogger("test_bringup_unexpected_message")

        self.do_bringup()

        # send a message that we're not supposed to
        # LS should send AbnormalTermination msg to launcher be
        self.ls_queue.send(dmsg.GSProcessCreate(tag=self.next_tag(), exe='/bin/echo',
                                                args=['Hello World'], p_uid=dfacts.LAUNCHER_PUID,
                                                r_c_uid=dfacts.BASE_BE_CUID).serialize())
        log.info("sent unexpected message to ls")
        tsu.get_and_check_type(self.la_queue, dmsg.AbnormalTermination)
        log.info("received AbnormalTermination from ls")

        self.do_teardown()


if __name__ == "__main__":
    unittest.main()

test/test_ls_multi.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest

from mpbridge.test_pipe_with_process import SendReceiveTest
from mpbridge.test_mpbridge_context_wait import TestDragonContextWait
from mpbridge.test_mpbridge_basic import TestStartMethod
from mpbridge.test_queue import TestQueue
from mpbridge.test_condition import TestCondition
from mpbridge.test_lock import TestDragonLocks
from mpbridge.test_process import TestMPBridgeProcess
from mpbridge.test_api import TestMultiprocessingAPI, TestMultiprocessingInternalPatching
# from mpbridge.test_barrier import TestBarrier


def setUpModule():
    import dragon
    import multiprocessing
    multiprocessing.set_start_method('dragon')


if __name__ == "__main__":
    unittest.main()

test/test_mpbridge.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest

from native.test_process_group import TestDragonNativeProcessGroup
from native.test_process import TestDragonNativeProcess
from native.test_semaphore import TestSemaphore
from native.test_machine import TestMachineSingle
from native.test_value import TestValue
from native.test_array import TestArray
from native.test_barrier import TestBarrier
from native.test_event import TestEvent
from native.test_lock import TestLock
from native.test_queue import TestQueue
from native.test_redirection import TestIORedirection


if __name__ == "__main__":
    unittest.main()

test/test_native.py
DragonHPC/dragon
#!/usr/bin/env python3
"""Test script for single node shepherd message handling"""

import shim_dragon_paths
import inspect
import multiprocessing as mp
import logging
import os
import signal
import sys
import time
import unittest

import dragon.channels as dch
import dragon.managed_memory as dmm

import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.connection as dconn
import dragon.infrastructure.facts as dfacts
import dragon.infrastructure.parameters as parms
import dragon.localservices.local_svc as dsls
import dragon.dlogging.util as dlog
import dragon.utils as du

import support.util as tsu

DEFAULT_TIMEOUT = 1

get_msg = tsu.get_and_parse


def run_ls(shep_stdin_queue, shep_stdout_queue, name_addition=''):
    dlog.setup_logging(basename='ls_' + name_addition, level=logging.DEBUG)
    log = logging.getLogger('outer runner')
    parms.this_process.mode = dfacts.SINGLE_MODE
    os.environ.update(parms.this_process.env())
    log.info('starting ls from runner')
    try:
        dsls.single(ls_stdin=shep_stdin_queue, ls_stdout=shep_stdout_queue)
    except RuntimeError as rte:
        log.exception('runner')
        shep_stdout_queue.send(dmsg.AbnormalTermination(tag=dsls.get_new_tag(),
                                                        err_info=f'{rte!s}').serialize())


# TODO: pending pool attach lifecycle solution
ATTACH_POOLS = True


class SingleLS(unittest.TestCase):
    DEFAULT_TEST_POOL_SIZE = 2 ** 20

    def setUp(self) -> None:
        self.shep_stdin_rh, self.shep_stdin_wh = mp.Pipe(duplex=False)
        self.shep_stdout_rh, self.shep_stdout_wh = mp.Pipe(duplex=False)
        self.mode = dfacts.SINGLE_MODE
        self.tag = 0

    def tearDown(self) -> None:
        self.shep_stdout_rh.close()
        self.shep_stdin_wh.close()

        self.shep_ch.detach()
        self.be_ch.detach()
        self.gs_ch.detach()

        if ATTACH_POOLS:
            self.def_pool.detach()
            self.inf_pool.detach()

    def next_tag(self):
        tmp = self.tag
        self.tag += 1
        return tmp

    def do_bringup(self):
        test_name = self.__class__.__name__ + '_' + inspect.stack()[1][0].f_code.co_name
        self.proc = mp.Process(target=run_ls,
                               args=(self.shep_stdin_rh, self.shep_stdout_wh),
                               kwargs={'name_addition': test_name}, daemon=True,
                               name='local_svc')
        self.proc.start()

        self.shep_stdin_wh.send(dmsg.BENodeIdxSH(tag=self.next_tag(), node_idx=0).serialize())

        msg = tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHPingBE)

        if ATTACH_POOLS:
            default_pd = du.B64.str_to_bytes(msg.default_pd)
            self.def_pool = dmm.MemoryPool.attach(default_pd)

            inf_pd = du.B64.str_to_bytes(msg.inf_pd)
            self.inf_pool = dmm.MemoryPool.attach(inf_pd)
        else:
            print(f'shep_cd: {msg.shep_cd}')

        shep_cd = du.B64.str_to_bytes(msg.shep_cd)
        self.shep_ch = dch.Channel.attach(shep_cd)
        self.shep_main_wh = dconn.Connection(outbound_initializer=self.shep_ch)

        be_cd = du.B64.str_to_bytes(msg.be_cd)
        self.be_ch = dch.Channel.attach(be_cd)
        self.be_main_rh = dconn.Connection(inbound_initializer=self.be_ch)

        gs_cd = du.B64.str_to_bytes(msg.gs_cd)
        self.gs_ch = dch.Channel.attach(gs_cd)
        self.gs_main_rh = dconn.Connection(inbound_initializer=self.gs_ch)

        self.shep_main_wh.send(dmsg.BEPingSH(tag=0).serialize())

        tsu.get_and_check_type(self.be_main_rh, dmsg.SHChannelsUp)

        self.shep_main_wh.send(dmsg.GSPingSH(tag=0).serialize())

        tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPingGS)

    def do_teardown(self):
        self.shep_main_wh.send(dmsg.GSHalted(tag=self.next_tag()).serialize())
        count = 0
        msg = get_msg(self.be_main_rh)

        # keep ignoring everything received other than GSHalted
        while not isinstance(msg, dmsg.GSHalted) and count < 100:
            msg = get_msg(self.be_main_rh)
            count += 1

        self.shep_main_wh.send(dmsg.SHTeardown(tag=self.next_tag()).serialize())

        count = 0
        msg = get_msg(self.be_main_rh)

        # keep ignoring everything received other than SHHaltBE
        while not isinstance(msg, dmsg.SHHaltBE) and count < 100:
            msg = get_msg(self.be_main_rh)
            count += 1

        self.shep_stdin_wh.send(dmsg.BEHalted(tag=self.next_tag()).serialize())
        tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHHalted)

        self.proc.join()
        if not ATTACH_POOLS:
            time.sleep(10)

    def do_abnormal_teardown(self):
        # when the BE receives AbnormalTermination, then the SHTeardown proceeds as usual
        # and LS expects SHTeardown in order to shut down
        tsu.get_and_check_type(self.be_main_rh, dmsg.AbnormalTermination)
        self.shep_main_wh.send(dmsg.SHTeardown(tag=self.next_tag()).serialize())
        tsu.get_and_check_type(self.be_main_rh, dmsg.SHHaltBE)
        self.shep_stdin_wh.send(dmsg.BEHalted(tag=self.next_tag()).serialize())
        tsu.get_and_check_type(self.shep_stdout_rh, dmsg.SHHalted)

        self.proc.join()

    def _make_pool(self, tag, p_uid, r_c_uid, size, m_uid, name):
        self.shep_main_wh.send(
            dmsg.SHPoolCreate(tag=tag, p_uid=p_uid, r_c_uid=r_c_uid, size=size,
                              m_uid=m_uid, name=name).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolCreateResponse)
        self.assertEqual(dmsg.SHPoolCreateResponse.Errors.SUCCESS, res.err)

    def _destroy_pool(self, tag, p_uid, r_c_uid, m_uid):
        # Destroy the pool
        self.shep_main_wh.send(dmsg.SHPoolDestroy(tag=tag, p_uid=p_uid,
                                                  r_c_uid=r_c_uid, m_uid=m_uid).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolDestroyResponse)
        self.assertEqual(dmsg.SHPoolDestroyResponse.Errors.SUCCESS, res.err)

    def test_bringup_teardown(self):
        # Tests normal bringup followed by teardown.
        self.do_bringup()
        self.do_teardown()

    # def test_bringup_teardown2(self):
    #     # for diagnosing pool lifecycle issues
    #     self.do_bringup()
    #     self.do_teardown()

    def test_process_fwdoutput(self):
        self.do_bringup()
        target_puid = 17777

        the_tag = self.next_tag()
        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                                    args=['shepherd/proc2.py'],
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)

        self.assertEqual(process_create_resp.ref, the_tag)

        process_output = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput, timeout=5)

        if process_output == "Hello World\n":
            process_output = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput, timeout=5)
            self.assertEqual(process_output.data, "Doing some more\n")
        else:
            self.assertEqual(process_output.data, "Hello World\nDoing some more\n")

        tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessExit)
        self.do_teardown()

    def test_process_fwdinput(self):
        self.do_bringup()

        target_puid = 17777
        the_tag = self.next_tag()
        self.shep_main_wh.send(
            dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                 args=['shepherd/proc4.py'], p_uid=dfacts.GS_PUID,
                                 r_c_uid=dfacts.GS_INPUT_CUID, t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)

        self.assertEqual(process_create_resp.ref, the_tag)

        self.shep_main_wh.send(
            dmsg.SHFwdInput(tag=self.next_tag(), p_uid=dfacts.GS_PUID, t_p_uid=target_puid,
                            r_c_uid=dfacts.GS_INPUT_CUID,
                            input="Hi There\n", confirm=True).serialize())

        tsu.get_and_check_type(self.gs_main_rh, dmsg.SHFwdInputErr)

        output = ''
        try:
            msg = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput)
            output += msg.data
        except TimeoutError:
            pass

        self.assertTrue((output in "Enter some text: Here is the response text.\nHi There\n") and (len(output) != 0))

        self.do_teardown()

    def test_sh_pool_new_destroy(self):
        self.do_bringup()

        test_muid = 123456
        self.shep_main_wh.send(
            dmsg.SHPoolCreate(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                              r_c_uid=dfacts.GS_INPUT_CUID, size=self.DEFAULT_TEST_POOL_SIZE,
                              m_uid=test_muid, name='single_shep_test').serialize())

        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolCreateResponse)
        self.assertEqual(dmsg.SHPoolCreateResponse.Errors.SUCCESS, res.err)

        self.shep_main_wh.send(dmsg.SHPoolDestroy(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                                                  r_c_uid=dfacts.GS_INPUT_CUID,
                                                  m_uid=test_muid).serialize())

        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolDestroyResponse)
        self.assertEqual(dmsg.SHPoolDestroyResponse.Errors.SUCCESS, res.err)

        self.do_teardown()

    def test_sh_pool_invalid_destroy(self):
        self.do_bringup()
        test_muid = 123456
        self.shep_main_wh.send(dmsg.SHPoolDestroy(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                                                  r_c_uid=dfacts.GS_INPUT_CUID,
                                                  m_uid=test_muid).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolDestroyResponse)
        self.assertEqual(dmsg.SHPoolDestroyResponse.Errors.FAIL, res.err)

        self.do_teardown()

    def test_sh_pool_duplicate_id(self):
        self.do_bringup()
        test_muid = 123456

        self.shep_main_wh.send(
            dmsg.SHPoolCreate(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                              r_c_uid=dfacts.GS_INPUT_CUID,
                              size=self.DEFAULT_TEST_POOL_SIZE, m_uid=test_muid,
                              name='single_shep_test').serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolCreateResponse)
        self.assertEqual(dmsg.SHPoolCreateResponse.Errors.SUCCESS, res.err)

        self.shep_main_wh.send(
            dmsg.SHPoolCreate(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                              r_c_uid=dfacts.GS_INPUT_CUID,
                              size=self.DEFAULT_TEST_POOL_SIZE, m_uid=test_muid,
                              name='single_shep_test').serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolCreateResponse)
        self.assertEqual(dmsg.SHPoolCreateResponse.Errors.FAIL, res.err)

        # Destroy the pool
        self.shep_main_wh.send(dmsg.SHPoolDestroy(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                                                  r_c_uid=dfacts.GS_INPUT_CUID,
                                                  m_uid=test_muid).serialize())

        # Check we succeeded so we know if other tests fail after this
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHPoolDestroyResponse)
        self.assertEqual(dmsg.SHPoolDestroyResponse.Errors.SUCCESS, res.err)

        self.do_teardown()

    def test_sh_channel_create_destroy(self):
        self.do_bringup()
        test_muid = 1234
        test_cuid = 4567

        self._make_pool(self.next_tag(), dfacts.GS_PUID, dfacts.GS_INPUT_CUID,
                        self.DEFAULT_TEST_POOL_SIZE, test_muid, 'single_shep_test')

        self.shep_main_wh.send(dmsg.SHChannelCreate(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID, m_uid=test_muid,
                                                    c_uid=test_cuid).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHChannelCreateResponse)
        self.assertEqual(dmsg.SHChannelCreateResponse.Errors.SUCCESS, res.err)

        self.shep_main_wh.send(dmsg.SHChannelDestroy(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                                                     r_c_uid=dfacts.GS_INPUT_CUID,
                                                     c_uid=test_cuid).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHChannelDestroyResponse)
        self.assertEqual(dmsg.SHChannelDestroyResponse.Errors.SUCCESS, res.err)

        self._destroy_pool(self.next_tag(), dfacts.GS_PUID, dfacts.GS_INPUT_CUID, test_muid)

        self.do_teardown()

    def test_sh_channel_create_duplicate_uid(self):
        self.do_bringup()
        test_muid = 1234
        test_cuid = 4567

        self._make_pool(self.next_tag(), dfacts.GS_PUID,
                        dfacts.GS_INPUT_CUID, self.DEFAULT_TEST_POOL_SIZE,
                        test_muid, 'single_shep_test')

        # Make a channel
        self.shep_main_wh.send(dmsg.SHChannelCreate(tag=self.next_tag(),
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    m_uid=test_muid, c_uid=test_cuid).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHChannelCreateResponse)
        self.assertEqual(dmsg.SHChannelCreateResponse.Errors.SUCCESS, res.err)

        self.shep_main_wh.send(dmsg.SHChannelCreate(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    m_uid=test_muid, c_uid=test_cuid).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHChannelCreateResponse)
        self.assertEqual(dmsg.SHChannelCreateResponse.Errors.FAIL, res.err)

        self.shep_main_wh.send(dmsg.SHChannelDestroy(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                                                     r_c_uid=dfacts.GS_INPUT_CUID,
                                                     c_uid=test_cuid).serialize())
        res = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHChannelDestroyResponse)
        self.assertEqual(dmsg.SHChannelDestroyResponse.Errors.SUCCESS, res.err)

        self._destroy_pool(self.next_tag(), dfacts.GS_PUID, dfacts.GS_INPUT_CUID, test_muid)

        self.do_teardown()

    def test_noprocess_fwdinput(self):
        self.do_bringup()
        target_puid = 177

        self.shep_main_wh.send(
            dmsg.SHFwdInput(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                            confirm=True, t_p_uid=target_puid,
                            r_c_uid=dfacts.GS_INPUT_CUID, input='Hi There\n').serialize())

        tsu.get_and_check_type(self.gs_main_rh, dmsg.SHFwdInputErr)

        self.do_teardown()

    def test_process_errexit(self):
        self.do_bringup()
        the_tag = self.next_tag()
        target_puid = 177
        self.shep_main_wh.send(
            dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                 args=['missing_prog.py'], p_uid=dfacts.GS_PUID,
                                 r_c_uid=dfacts.GS_INPUT_CUID,
                                 t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)

        self.assertEqual(process_create_resp.ref, the_tag)

        process_exit = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessExit)

        self.assertEqual(process_exit.exit_code, 2)

        tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput)

        self.do_teardown()

    def test_dump_state(self):
        self.do_bringup()
        self.shep_main_wh.send(dmsg.SHDumpState(tag=self.next_tag()).serialize())
        self.do_teardown()

    def test_dump_state_to_file(self):
        self.do_bringup()
        self.shep_main_wh.send(dmsg.SHDumpState(tag=self.next_tag(),
                                                filename='shep_dump.log').serialize())
        self.do_teardown()

    def test_process_kill(self):
        self.do_bringup()
        target_puid = 177
        the_tag = self.next_tag()
        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                                    args=['shepherd/proc2.py'],
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)

        self.assertEqual(process_create_resp.ref, the_tag)

        process_output = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput, timeout=5)

        self.shep_main_wh.send(
            dmsg.SHProcessKill(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                               t_p_uid=target_puid, r_c_uid=dfacts.GS_INPUT_CUID,
                               sig=signal.SIGKILL).serialize())

        if process_output.data == "Hello World\n":
            process_output = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput, timeout=5)
            self.assertEqual(process_output.data, "Doing some more\n")
        else:
            self.assertEqual(process_output.data, "Hello World\nDoing some more\n")

        msg = tsu.get_and_parse(self.gs_main_rh)
        self.assertTrue(isinstance(msg, dmsg.SHProcessExit) or isinstance(msg, dmsg.SHProcessKillResponse))
        self.do_teardown()

    def test_process_kill_wait(self):
        self.do_bringup()
        target_puid = 177
        the_tag = self.next_tag()
        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                                    args=['shepherd/proc3.py'],
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)

        self.assertEqual(process_create_resp.ref, the_tag)

        self.shep_main_wh.send(
            dmsg.SHProcessKill(tag=self.next_tag(), p_uid=dfacts.GS_PUID,
                               t_p_uid=target_puid, r_c_uid=dfacts.GS_INPUT_CUID,
                               sig=signal.SIGKILL).serialize())

        msg = tsu.get_and_parse(self.gs_main_rh)
        self.assertTrue(isinstance(msg, dmsg.SHProcessExit) or isinstance(msg, dmsg.SHProcessKillResponse))
        self.do_teardown()

    def test_bringup_abnormal_termination(self):
        self.do_bringup()
        self.shep_main_wh.send(dmsg.AbnormalTermination(tag=self.next_tag()).serialize())
        self.do_abnormal_teardown()

    def test_bringup_bad_message(self):
        self.do_bringup()
        self.shep_main_wh.send("crap")
        self.do_abnormal_teardown()

    def test_process_create(self):
        self.do_bringup()
        the_tag = self.next_tag()
        target_puid = 177
        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                                    args=['shepherd/proc1.py'],
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)

        self.assertEqual(process_create_resp.ref, the_tag)
        self.shep_main_wh.send(dmsg.SHDumpState(tag=self.next_tag()).serialize())
        process_output = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput)
        self.assertEqual(process_output.data, "Hello World\n")
        tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessExit, 5)

        self.do_teardown()

    def test_process_create_with_env(self):
        self.do_bringup()
        target_puid = 177
        the_tag = self.next_tag()
        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                                    args=['shepherd/procenv1.py'],
                                                    env={'new_env_var': 'env_value'},
                                                    p_uid=dfacts.GS_PUID, r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)
        self.assertEqual(process_create_resp.ref, the_tag)

        self.shep_main_wh.send(dmsg.SHDumpState(tag=0).serialize())

        process_output = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput)
        self.assertEqual(process_output.data, "env_value\n")

        process_done = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessExit, 5)
        self.assertEqual(process_done.exit_code, 0)

        self.do_teardown()

    def test_process_create_with_env2(self):
        self.do_bringup()
        target_puid = 177
        the_tag = self.next_tag()
        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=the_tag, exe=sys.executable,
                                                    args=['shepherd/procenv2.py'],
                                                    env={'SHELL': 'shell_value'},
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=target_puid).serialize())

        process_create_resp = tsu.get_and_check_type(self.gs_main_rh,
                                                     dmsg.SHProcessCreateResponse)
        self.assertEqual(process_create_resp.ref, the_tag)

        self.shep_main_wh.send(dmsg.SHDumpState(tag=self.next_tag()).serialize())

        process_output = tsu.get_and_check_type(self.be_main_rh, dmsg.SHFwdOutput)
        self.assertEqual(process_output.data, "shell_value\n")

        process_done = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessExit, 5)
        self.assertEqual(process_done.exit_code, 0)

        self.do_teardown()

    def test_process_create_dup(self):
        self.do_bringup()

        tag1 = self.next_tag()
        test_puid = 177

        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=tag1, exe=sys.executable,
                                                    args=['shepherd/proc1.py'],
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=test_puid).serialize())

        tag2 = self.next_tag()
        self.shep_main_wh.send(dmsg.SHProcessCreate(tag=tag2, exe=sys.executable,
                                                    args=['shepherd/proc1.py'],
                                                    p_uid=dfacts.GS_PUID,
                                                    r_c_uid=dfacts.GS_INPUT_CUID,
                                                    t_p_uid=test_puid).serialize())

        process_create_resp1 = tsu.get_and_check_type(self.gs_main_rh, dmsg.SHProcessCreateResponse)

        if process_create_resp1.ref == tag1:
            process_create_resp2 = tsu.get_and_check_type(self.gs_main_rh,
                                                          dmsg.SHProcessCreateResponse)
            self.assertEqual(process_create_resp2.ref, tag2)
            self.assertEqual(process_create_resp2.err, dmsg.SHProcessCreateResponse.Errors.FAIL)
        else:
            self.assertEqual(process_create_resp1.ref, tag2)
            self.assertEqual(process_create_resp1.err, dmsg.SHProcessCreateResponse.Errors.FAIL)
            process_create_resp2 = tsu.get_and_check_type(self.gs_main_rh,
                                                          dmsg.SHProcessCreateResponse)
            self.assertEqual(process_create_resp2.ref, tag1)

        self.do_teardown()


if __name__ == "__main__":
    unittest.main()

test/test_shepherd.py
DragonHPC/dragon
#!/usr/bin/env python3
import shim_dragon_paths
import unittest

from transport.test_tcp_transport import SingleNodeTransportBench

if __name__ == "__main__":
    unittest.main()

test/test_transport.py
DragonHPC/dragon
#!/usr/bin/env python3

import shim_dragon_paths
import unittest
from utils.test_basic_mempool import (
    MemPoolCreateTest,
    MemoryPoolAllocTest,
    MemoryPoolAllocNoSetupTests,
    MemoryPoolAttachTests,
)

from utils.test_mempool import MemPoolTest
from utils.test_pyheap import PHeapTest
from utils.test_logging import LoggingTest, TestLogHandler
from utils.test_logging import TestLoggingSubprocesses

# TEMPORARY HACK: somehow the use of parameterized in these tests
# seem to be interfering with unittests's discovery on import.
# This brutal method solves it for the time being but:
# TODO: reorganize the tests overall
with open("utils/test_locks.py") as fh:
    exec(fh.read())

if __name__ == "__main__":
    unittest.main()

test/test_utils.py
DragonHPC/dragon
from asyncio import StreamReader, StreamWriter
from ipaddress import ip_address
import unittest
from unittest.mock import MagicMock, patch

from dragon.transport.tcp import messages
from dragon.transport.tcp import transport


class AddressIPv4TestCase(unittest.IsolatedAsyncioTestCase):

    IP = ip_address('127.0.0.1')
    NETLOC = ['127.0.0.1', '[127.0.0.1]']

    LOOPBACK = transport.LOOPBACK_ADDRESS_IPv4

    Hello = messages.Hello

    def test_loopback(self):
        self.assertEqual(self.LOOPBACK, transport.Address(self.IP, None))

    def test_from_netloc(self):
        addr = transport.Address(self.IP, 8888)
        for s in self.NETLOC:
            self.assertEqual(self.LOOPBACK, transport.Address.from_netloc(s))
            self.assertEqual(addr, transport.Address.from_netloc(f'{s}:8888'))

    def test___str__(self):
        s = self.NETLOC[0]
        self.assertEqual(s, str(self.LOOPBACK))
        self.assertEqual(f'{s}:8888', str(transport.Address(self.IP, 8888)))

    def test_hello(self):
        self.assertRaises(AssertionError, self.LOOPBACK.hello)
        hello = transport.Address(self.IP, 8888).hello()
        self.assertIsInstance(hello, self.Hello)
        self.assertEqual(hello.host, self.IP)
        self.assertEqual(hello.port, 8888)

    def test_from_hello(self):
        addr = transport.Address.from_hello(self.Hello(self.IP, 8888))
        self.assertEqual(addr.host, self.IP)
        self.assertEqual(addr.port, 8888)

    # XXX Note that write_message() and read_message() must be patched from
    # XXX transport since they are imported there.
    @patch.object(transport, 'write_message')
    @patch.object(transport, 'read_message')
    async def _do_handshake(self, read_message, write_message):
        my_addr = transport.Address(self.IP, 8888)
        peer_addr = transport.Address(self.IP, 9999)
        read_message.return_value = peer_addr.hello()

        reader = MagicMock(spec=StreamReader)
        writer = MagicMock(spec=StreamWriter)

        addr = await my_addr.do_handshake(reader, writer)

        # Verify correct peer address is returned
        self.assertEqual(addr, peer_addr)

        # Verify mocks were called correctly
        write_message.assert_called_once_with(writer, my_addr.hello())
        read_message.assert_awaited_once_with(reader)


class AddressIPv6TestCase(AddressIPv4TestCase):

    IP = ip_address('::1')
    NETLOC = ['[::1]']

    LOOPBACK = transport.LOOPBACK_ADDRESS_IPv6

    Hello = messages.Hello6


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_address.py
DragonHPC/dragon
import unittest


class AgentTestCase(unittest.IsolatedAsyncioTestCase):
    pass


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_agent.py
DragonHPC/dragon
import unittest

from test_transport import TestMessages


class ClientTestCase(TestMessages, unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    async def test_run(self):
        raise NotImplementedError

    @unittest.skip
    def test_open(self):
        raise NotImplementedError

    @unittest.skip
    def test_close(self):
        raise NotImplementedError

    @unittest.skip
    async def test_recv(self):
        raise NotImplementedError

    @unittest.skip
    def test_process(self):
        raise NotImplementedError

    @unittest.skip
    async def test_wait_for_response(self):
        raise NotImplementedError

    @unittest.skip
    async def test_handle_unsupported_response(self):
        raise NotImplementedError

    @unittest.skip
    async def test_handle_error_response(self):
        raise NotImplementedError

    @unittest.skip
    async def test_handle_send_response(self):
        raise NotImplementedError

    @unittest.skip
    async def test_handle_recv_response(self):
        raise NotImplementedError

    @unittest.skip
    async def test_handle_event_response(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_client.py
DragonHPC/dragon
import unittest

from dragon.transport.tcp import errno


class ErrnoTestCase(unittest.TestCase):

    def test_get_errno(self):
        # TODO Add tests for Dragon exceptions that have the `lib_err` attribute
        self.assertEqual(errno.get_errno(Exception()), errno.DRAGON_FAILURE)
        self.assertEqual(errno.get_errno(ValueError()), errno.DRAGON_INVALID_ARGUMENT)
        self.assertEqual(errno.get_errno(NotImplementedError()), errno.DRAGON_NOT_IMPLEMENTED)
        self.assertEqual(errno.get_errno(TimeoutError()), errno.DRAGON_TIMEOUT)


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_errno.py
DragonHPC/dragon
import unittest


class GatewayMessaageTestCase(unittest.TestCase):

    @unittest.skip
    def test_request(self):
        raise NotImplementedError

    @unittest.skip
    def test_complete(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_gateway_message.py
DragonHPC/dragon
import asyncio
import base64
import unittest
from unittest.mock import MagicMock, patch

from dragon.transport.tcp.io import CodableIO, FixedBytesIO
from dragon.transport.tcp import transport


class FixedBytesIOTestCase(unittest.IsolatedAsyncioTestCase):

    @classmethod
    def setUpClass(cls):
        cls.types = [
            type(f'FixedBytesIO{size}', (FixedBytesIO,), dict(size=size))
            for size in (16, 64, 256, 1024)
        ]

    async def test_read(self):
        for cls in self.types:
            data = b'\xff' * cls.size
            r, w = await transport.create_pipe_streams()
            try:
                w.write(data)
                await w.drain()
            finally:
                await transport.close_writer(w)
            _data = await cls().read(r)
            self.assertEqual(_data, data)

    async def test_incomplete_read(self):
        for cls in self.types:
            data = b'\xff' * (cls.size - 1)
            r, w = await transport.create_pipe_streams()
            try:
                w.write(data)
                await w.drain()
            finally:
                await transport.close_writer(w)
            with self.assertRaises(asyncio.IncompleteReadError):
                await cls().read(r)

    async def test_write(self):
        for cls in self.types:
            data = b'\xff' * cls.size
            r, w = await transport.create_pipe_streams()
            try:
                cls().write(w, data)
                await w.drain()
            finally:
                await transport.close_writer(w)
            _data = await r.readexactly(cls.size)
            self.assertEqual(_data, data)

    async def test_incomplete_write(self):
        for cls in self.types:
            data = b'\xff' * (cls.size - 1)
            r, w = await transport.create_pipe_streams()
            try:
                with self.assertRaises(AssertionError):
                    cls().write(w, data)
            finally:
                await transport.close_writer(w)


class CodableIOTestCase(unittest.IsolatedAsyncioTestCase):

    @classmethod
    def setUpClass(cls):
        cls.cls = type('Base64CodableIO', (CodableIO,), dict(
            encode=staticmethod(base64.standard_b64encode),
            decode=staticmethod(base64.standard_b64decode),
        ))
        cls.DATA = b'\xffTEST DATA\xff'
        cls.ENCDATA = base64.standard_b64encode(cls.DATA)

    @patch.object(FixedBytesIO, 'read')
    async def test_read(self, read):
        read.return_value = self.ENCDATA
        reader = MagicMock(spec=asyncio.StreamReader)
        data = await self.cls().read(reader)
        self.assertEqual(data, self.DATA)
        read.assert_called_once_with(reader)

    @patch.object(FixedBytesIO, 'write')
    def test_write(self, write):
        writer = MagicMock(spec=asyncio.StreamWriter)
        self.cls().write(writer, self.DATA)
        write.assert_called_once_with(writer, self.ENCDATA)


class StructIOTestCase(unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    async def test_read(self):
        raise NotImplementedError

    @unittest.skip
    def test_write(self):
        raise NotImplementedError


class VariableBytesIOTestCase(unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    async def test_read(self):
        raise NotImplementedError

    @unittest.skip
    def test_write(self):
        raise NotImplementedError


class VariableTestIOTestCase(unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    async def test_read(self):
        raise NotImplementedError

    @unittest.skip
    def test_write(self):
        raise NotImplementedError


class IPv4AddressIOTestCase(unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    async def test_read(self):
        raise NotImplementedError

    @unittest.skip
    def test_write(self):
        raise NotImplementedError


class IPv6AddressIOTestCase(IPv4AddressIOTestCase):
    pass


class EnumIOTestCase(unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    async def test_read(self):
        raise NotImplementedError

    @unittest.skip
    def test_write(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_io.py
DragonHPC/dragon
from contextlib import asynccontextmanager
from ipaddress import ip_address
import math
import unittest
from uuid import uuid4

from dragon.transport.tcp import messages
from dragon.transport.tcp import transport
from dragon.transport.tcp.util import seconds_remaining


class TransportMessageTestCase(unittest.TestCase):

    @unittest.skip
    def test_init_subclass(self):
        raise NotImplementedError


class RequestTestCase(unittest.TestCase):

    def test_deadline_getter(self):
        req = messages.Request(seqno=None, timeout=1.0, channel_sd=None)
        self.assertEqual(req.deadline, req.timestamp + req.timeout)

    def test_deadline_getter_none(self):
        req = messages.Request(seqno=None, timeout=math.inf, channel_sd=None)
        self.assertEqual(req.deadline, math.inf)

    def test_deadline_setter_none(self):
        req = messages.Request(seqno=None, timeout=1.0, channel_sd=None)
        req.deadline = math.inf
        self.assertEqual(req.timeout, math.inf)

    def test_deadline_setter_increment_timeout(self):
        req = messages.Request(seqno=None, timeout=1.0, channel_sd=None)
        ts = req.timestamp  # preserve original timestamp
        req.deadline += 1.0
        # check that original timeout was increased by 1.0s
        self.assertEqual(req.deadline - ts, 2.0)

    def test_timeout_scales_to_zero(self):
        req = messages.Request(seqno=None, timeout=1.0, channel_sd=None)
        timeout, _ = seconds_remaining(req.deadline - 1.1)
        self.assertEqual(timeout, 0)


class TransmittableTestCase(unittest.IsolatedAsyncioTestCase):

    def setUp(self):
        if not (hasattr(self, 'msg') and hasattr(self, 'data')):
            self.skipTest('Missing "msg" and "data" attributes')
        self.assertIsInstance(self.msg, messages.TransportMessage)
        self.assertIsInstance(self.data, bytes)

    async def asyncSetUp(self):
        self.reader, self.writer = await transport.create_pipe_streams()

    @asynccontextmanager
    async def drain_and_close_writer(self):
        try:
            yield self.writer
            await self.writer.drain()
        finally:
            await transport.close_writer(self.writer)

    async def test_read(self):
        async with self.drain_and_close_writer() as w:
            w.write(self.data)
        req = await messages.read_message(self.reader)
        self.assertEqual(req, self.msg)

    async def test_write(self):
        async with self.drain_and_close_writer() as w:
            await messages.write_message(w, self.msg)
        data = await self.reader.read()
        self.assertEqual(data, self.data)


class HelloTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.Hello(ip_address('127.0.0.1'), 8888)
        cls.data = b'\x40\x7f\x00\x00\x01"\xb8'


class Hello6TestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.Hello6(ip_address('::1'), 8888)
        cls.data = b'\x60\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x01"\xb8'


class SendRequestTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.SendRequest(
            seqno=1, timeout=0.5, channel_sd=b'channel desc',
            return_mode=messages.SendReturnMode.WHEN_BUFFERED, sendhid=uuid4(),
            payload=b'payload',
        )
        cls.data = b'\x01\x00\x00\x00\x00\x00\x00\x00\x01?\x00\x00\x00\x00\x0cchannel desc\x02' + cls.msg.sendhid.bytes + b'\x00\x00\x00\x00\x00\x00\x00\x07payload'

class SendMemoryRequestTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.SendMemoryRequest(
            seqno=1, timeout=0.5, channel_sd=b'channel desc',
            return_mode=messages.SendReturnMode.WHEN_BUFFERED, sendhid=uuid4(),
            payload=b'payload', mem_sd=b'memory desc',
        )
        cls.data = b'\x02\x00\x00\x00\x00\x00\x00\x00\x01?\x00\x00\x00\x00\x0cchannel desc\x02' + cls.msg.sendhid.bytes + b'\x00\x00\x00\x00\x00\x00\x00\x07payload\x00\x0bmemory desc'


class RecvRequestTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.RecvRequest(seqno=1, timeout=0.5, channel_sd=b'channel desc')
        cls.data = b'\x03\x00\x00\x00\x00\x00\x00\x00\x01?\x00\x00\x00\x00\x0cchannel desc'


class EventRequestTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.EventRequest(seqno=1, timeout=0.5, channel_sd=b'channel desc', mask=3)
        cls.data = b'\x04\x00\x00\x00\x00\x00\x00\x00\x01?\x00\x00\x00\x00\x0cchannel desc\x00\x03'

class ErrorResponseTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.ErrorResponse(seqno=1, errno=42, text='error message')
        cls.data = b'\xff\x00\x00\x00\x00\x00\x00\x00\x01\x00*\x00\x00\x00\x00\x00\x00\x00\rerror message'


class SendResponseTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.SendResponse(seqno=1)
        cls.data = b'\xfe\x00\x00\x00\x00\x00\x00\x00\x01'


class RecvResponseTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.RecvResponse(seqno=1, payload=b'payload')
        cls.data = b'\xfc\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x07payload'


class EventResponseTestCase(TransmittableTestCase):

    @classmethod
    def setUpClass(cls):
        cls.msg = messages.EventResponse(seqno=1, errno=0, result=42)
        cls.data = b'\xfb\x00\x00\x00\x00\x00\x00\x00\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00*'


class MessagesTestCase(unittest.IsolatedAsyncioTestCase):

    # XXX Not sure of a more effective test compared to all the
    # XXX TransmittableTestCase subclasses.

    @unittest.skip
    def test__get_io_annotations(self):
        raise NotImplementedError

    @unittest.skip
    async def test_read_message(self):
        raise NotImplementedError

    @unittest.skip
    async def test_write_message(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_messages.py
DragonHPC/dragon
import asyncio
import unittest
from unittest.mock import AsyncMock, MagicMock

from dragon.transport.tcp.server import Server
from dragon.transport.tcp import transport
from dragon.transport.tcp import messages

from test_transport import TestMessages


class ServerTestCase(TestMessages, unittest.IsolatedAsyncioTestCase):

    ADDR = transport.Address.from_netloc('127.0.0.1:7575')

    def setUp(self):
        self.transport = transport.StreamTransport(self.ADDR)

    def tearDown(self):
        del self.transport

    async def asyncSetUp(self):
        self.transport.start()
        await self.transport.wait_serving()

    async def asyncTearDown(self):
        try:
            await self.transport.stop()
        except AssertionError:
            pass

    async def test_run_read_send_request(self):
        # Mock Transport.read_request() to read fake request
        req = self.SendRequest()
        addr = transport.Address.from_netloc('127.0.0.1:8888')
        self.transport.read_request = AsyncMock(side_effect=[(req, addr), asyncio.CancelledError])

        # Mock Server._ensure_process_task() to prevent a process task from
        # being created.
        server = Server(self.transport)
        server._ensure_process_task = MagicMock()

        await server.run()

        key = (req.channel_sd, req.sendhid)
        server._ensure_process_task.assert_called_once_with(key)

        _req, _addr = server._requests[key].get_nowait()
        self.assertIs(req, _req)
        self.assertIs(addr, _addr)

    async def test_run_read_recv_request(self):
        # Mock Transport.read_request() to read fake request
        req = self.RecvRequest()
        addr = transport.Address.from_netloc('127.0.0.1:8888')
        self.transport.read_request = AsyncMock(side_effect=[(req, addr), asyncio.CancelledError])

        # Mock Server._ensure_process_task() to prevent a process task from
        # being created.
        server = Server(self.transport)
        server._ensure_process_task = MagicMock()

        await server.run()

        key = (req.channel_sd, messages.RecvRequest)
        server._ensure_process_task.assert_called_once_with(key)

        _req, _addr = server._requests[key].get_nowait()
        self.assertIs(req, _req)
        self.assertIs(addr, _addr)

    @unittest.skip
    def test__ensure_process_task(self):
        raise NotImplementedError

    @unittest.skip
    async def test_process(self):
        raise NotImplementedError

    @unittest.skip
    def test_handle_unsupported_request(self):
        raise NotImplementedError

    @unittest.skip
    def test_handle_send_request(self):
        raise NotImplementedError

    @unittest.skip
    def test_handle_recv_request(self):
        raise NotImplementedError

    @unittest.skip
    def test_handle_event_request(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_server.py
DragonHPC/dragon
import asyncio
from datetime import timedelta
import logging
from pathlib import Path
import ssl
from tempfile import TemporaryDirectory
import unittest
from unittest.mock import MagicMock, patch

from dragon.transport.tcp import transport

from test_transport import TestMessages


def setUpModule():
    """Disable logging while running these tests."""
    logging.disable()


def tearDownModule():
    """Re-enable logging after running these tests."""
    logging.disable(logging.NOTSET)


class StreamTransportTestCase(unittest.IsolatedAsyncioTestCase):

    def setUp(self):
        addr = transport.Address.from_netloc('127.0.0.1:7575')
        self.transport = transport.StreamTransport(addr)

    def tearDown(self):
        del self.transport

    @unittest.skip
    async def test_close(self):
        raise NotImplementedError

    async def test_wait_serving(self):
        self.assertIsNone(self.transport._server)
        with patch.object(self.transport, '_server', autospec=asyncio.base_events.Server) as server:
            server.is_serving.side_effect = [False, False, True]
            await self.transport.wait_serving()
        self.assertEqual(server.is_serving.call_count, 3)

    @patch.object(transport.StreamTransport, 'add_connection')
    async def test_accept_connection(self, add_connection):
        # Transport.accept_connection() completes the address handshake and
        # adds the connection to internal state.
        addr = transport.Address.from_netloc('127.0.0.1:8888')
        reader = MagicMock(spec=asyncio.StreamReader)
        writer = MagicMock(spec=asyncio.StreamWriter)

        with patch.object(transport.Address, 'do_handshake', return_value=addr) as do_handshake:
            await self.transport.accept_connection(reader, writer)

        do_handshake.assert_awaited_once_with(reader, writer)
        add_connection.assert_called_once_with(addr, reader, writer)

    @patch.object(transport.StreamTransport, 'add_connection')
    @patch('asyncio.open_connection')
    async def test__open_connection(self, open_connection, add_connection):
        # Transport._open_connection() connects to a transport server and
        # initiates the address handshake. Once established, the connection
        # is added to the transport's internal state.
        addr = transport.Address.from_netloc('127.0.0.1:8888')
        reader = MagicMock(spec=asyncio.StreamReader)
        writer = MagicMock(spec=asyncio.StreamWriter)

        open_connection.return_value = (reader, writer)

        with patch.object(transport.Address, 'do_handshake', return_value=addr) as do_handshake:
            r, w = await self.transport._open_connection(addr)

        self.assertEqual(open_connection.await_args.args, (str(addr.host), int(addr.port)))
        do_handshake.assert_awaited_once_with(reader, writer)
        add_connection.assert_called_once_with(addr, reader, writer)

    @patch.object(transport.StreamTransport, '_do_recv', new_callable=MagicMock)
    @patch('asyncio.create_task')
    def test_add_connection(self, create_task, _do_recv):
        task = MagicMock(spec=asyncio.Task)
        create_task.return_value = task

        addr = transport.Address.from_netloc('127.0.0.1:8888')
        reader = MagicMock(spec=asyncio.StreamReader)
        writer = MagicMock(spec=asyncio.StreamWriter)
        self.transport.add_connection(addr, reader, writer)

        # Verify writer added to _writers list for addr
        self.assertIn(writer, self.transport._writers[addr])
        # Verify recv task started
        _do_recv.assert_called_once_with(addr, reader)
        create_task.assert_called_once()
        # Verify recv task added to set of recv tasks for addr
        self.assertIn(task, self.transport._recv_tasks[addr])

    @patch.object(transport, 'writer_addrs', None)
    @patch.object(transport, 'close_writer')
    @patch.object(transport.StreamTransport, '_open_connection')
    async def test_connect(self, open_connection, close_writer):
        def _open_connection(addr):
            # Simulate connection errors when trying to open a new connection
            if open_connection.await_count < 5:
                raise ConnectionRefusedError
            reader = MagicMock(spec=asyncio.StreamReader)
            writer = MagicMock(spec=asyncio.StreamWriter)
            writer.is_closing.return_value = False
            # XXX Skip the handshake!
            # XXX Do not call add_connection() since that will start recv
            # XXX tasks. Instead, just update the _writers list.
            #self.transport.add_connection(addr, reader, writer)
            self.transport._writers[addr].append(writer)
            return reader, writer
        open_connection.side_effect = _open_connection

        # Verify self connections are prevented
        with self.assertRaises(AssertionError):
            await self.transport.connect(self.transport.addr)

        self.addAsyncCleanup(self.transport.close)

        addr = transport.Address.from_netloc('127.0.0.1:8888')

        # Verify new connection opened
        writer = await self.transport.connect(addr)
        open_connection.assert_awaited_with(addr)
        self.assertIn(writer, self.transport._writers[addr])

        # Verify existing connection is returned
        open_connection.reset_mock()
        writer2 = await self.transport.connect(addr)
        open_connection.assert_not_awaited()
        self.assertIs(writer2, writer)

        writer.is_closing.return_value = True
        open_connection.reset_mock()
        writer3 = await self.transport.connect(addr)
        open_connection.assert_awaited_with(addr)
        # Check that closed writer has been removed
        self.assertNotIn(writer, self.transport._writers[addr])
        self.assertIn(writer3, self.transport._writers[addr])
        # XXX Not sure we can speculate about how many recv tasks may be running
        # XXX at this point

    @unittest.skip
    async def _do_recv(self):
        raise NotImplementedError

    @unittest.skip
    def test__handle_send(self):
        raise NotImplementedError

    @unittest.skip
    def test__ensure_send_task(self):
        raise NotImplementedError

    @unittest.skip
    async def test__do_send(self):
        raise NotImplementedError


class StreamTransportConnectionTestCase(TestMessages, unittest.IsolatedAsyncioTestCase):

    @classmethod
    def setUpClass(cls):
        super().setUpClass()
        cls.addr = transport.Address.from_netloc('127.0.0.1:7575')

    def setUp(self):
        self.transport = transport.StreamTransport(self.addr)

    def tearDown(self):
        del self.transport

    def _assert_transport_empty(self):
        # Verify transport data structures are cleaned up
        self.assertIsNone(self.transport._task)
        self.assertFalse(self.transport._recv_tasks)
        self.assertFalse(self.transport._send_tasks)
        self.assertFalse(self.transport._writers)

    async def _open_client_connection(self):
        opts = self.transport.default_connection_options.new_child(self.transport.connection_options[self.transport.addr])
        return await asyncio.open_connection(
            str(self.transport.addr.host),
            int(self.transport.addr.port),
            **opts,
        )

    async def test_accept_connection(self):
        client_req = self.SendRequest(seqno=42)
        client_addr = transport.Address.from_netloc('127.0.0.1:8686')

        self.transport.start()
        await self.transport.wait_serving()
        try:
            reader, writer = await self._open_client_connection()
            try:
                server_addr = await client_addr.do_handshake(reader, writer)
                self.assertEqual(server_addr, self.transport.addr)

                await transport.write_message(writer, client_req)
            finally:
                await transport.close_writer(writer)

            req, addr = await self.transport.read_request()
            self.assertEqual(req, client_req)
            self.assertEqual(addr, client_addr)
        finally:
            await self.transport.stop()

        self._assert_transport_empty()

    async def _start_server(self, accept_connection):
        return await asyncio.start_server(
            accept_connection,
            str(self.transport.addr.host),
            int(self.transport.addr.port),
            **self.transport.server_options,
        )

    async def test_open_connection(self):
        server_req = self.SendRequest(seqno=42)
        server_addr = transport.Address.from_netloc('127.0.0.1:8686')

        async def accept_connection(reader, writer):
            try:
                client_addr = await server_addr.do_handshake(reader, writer)
                self.assertEqual(client_addr, self.transport.addr)

                await transport.write_message(writer, server_req)
            finally:
                await transport.close_writer(writer)

        server = await self._start_server(accept_connection)
        async with server:
            task = asyncio.create_task(server.serve_forever())
            try:
                await self.transport._open_connection(self.transport.addr)
                req, addr = await self.transport.read_request()
                self.assertEqual(req, server_req)
                self.assertEqual(addr, server_addr)
            finally:
                await self.transport.close()
                task.cancel()
                try:
                    await task
                except asyncio.CancelledError:
                    pass

        self._assert_transport_empty()


try:
    from dragon.transport.x509 import CertificateAuthority, dragon_uri, generate_server_csr, generate_server_self_signed_cert, pem_encode, pem_encode_pkcs8
except ImportError:
    pass
else:
    class StreamTransportSSLConnectionTestCase(StreamTransportConnectionTestCase):

        """StreamTransport tests where transport supports basic SSL using
        self-signed certs and no verification.
        """

        @classmethod
        def setUpClass(cls):
            super().setUpClass()
            cls._uri = dragon_uri(0, cls.addr.host, cls.addr.port)
            # Create temporary directory for key, cert, and CA files
            cls._tempdir = TemporaryDirectory()
            cls._workdir = Path(cls._tempdir.name)
            cls._generate_ssl_files()

        @classmethod
        def tearDownClass(cls):
            cls._tempdir.cleanup()

        def setUp(self):
            super().setUp()
            self.transport.server_options.update(ssl=self._create_server_ssl_context())
            self.transport.default_connection_options.update(ssl=self._create_client_ssl_context())

        @classmethod
        def _generate_ssl_files(cls):
            # Generate self-signed certificate; important to use same IP addr as
            # transport will use.
            key, cert = generate_server_self_signed_cert(cls.addr.host, cls._uri, valid_for=timedelta(minutes=5))
            # Write key file
            cls._keyfile = cls._workdir / 'server-key.pem'
            with open(cls._keyfile, 'wb') as f:
                f.write(pem_encode_pkcs8(key))
            # Write cert file
            cls._certfile = cls._workdir / 'server-cert.pem'
            with open(cls._certfile, 'wb') as f:
                f.write(pem_encode(cert))

        def _create_server_ssl_context(self):
            """Returns SSL Context for servers."""
            ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
            ctx.load_cert_chain(self._certfile, self._keyfile)
            return ctx

        def _create_client_ssl_context(self):
            """Returns SSL Context for client connections."""
            ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
            # Disable server verification since it's a self signed cert
            ctx.check_hostname = False
            ctx.verify_mode = ssl.CERT_NONE
            return ctx


    class StreamTransportMutualTLSConnectionTestCase(StreamTransportSSLConnectionTestCase):
        """StreamTransport tests where transport supports mutual-TLS using certs
        issued by trusted CA.
        """

        @classmethod
        def _generate_ssl_files(cls):
            ca = CertificateAuthority.generate()
            key, csr = generate_server_csr(cls.addr.host, cls._uri)
            cert = ca.issue_server_certificate(csr, valid_for=timedelta(minutes=5))
            # Write CA file
            cls._cafile = cls._workdir / 'ca.pem'
            with open(cls._cafile, 'wb') as f:
                f.write(pem_encode(ca.certificate))
            # Write key file
            cls._keyfile = cls._workdir / 'server-key.pem'
            with open(cls._keyfile, 'wb') as f:
                f.write(pem_encode_pkcs8(key))
            # Write cert file
            cls._certfile = cls._workdir / 'server-cert.pem'
            with open(cls._certfile, 'wb') as f:
                f.write(pem_encode(cert))

        def _create_server_ssl_context(self):
            """Returns SSL Context for servers."""
            ctx = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
            ctx.load_cert_chain(self._certfile, self._keyfile)
            ctx.load_verify_locations(self._cafile)
            # Require client verification (e.g., enable privacy without authentication)
            ctx.verify_mode = ssl.CERT_REQUIRED
            return ctx

        def _create_client_ssl_context(self):
            """Returns SSL Context for client connections."""
            ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH)
            ctx.load_cert_chain(self._certfile, self._keyfile)
            ctx.load_verify_locations(self._cafile)
            assert ctx.verify_mode == ssl.CERT_REQUIRED
            return ctx


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_stream_transport.py
DragonHPC/dragon
import unittest


class StreamsTestCase(unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    async def test_close_writer(self):
        raise NotImplementedError

    @unittest.skip
    async def test_create_streams(self):
        raise NotImplementedError

    @unittest.skip
    async def test_create_pipe_streams(self):
        raise NotImplementedError

    @unittest.skip
    async def test_create_pipe_connections(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_streams.py
DragonHPC/dragon
import unittest


class TaskTestCase(unittest.IsolatedAsyncioTestCase):

    @unittest.skip
    def test_run_forever(self):
        raise NotImplementedError

    @unittest.skip
    async def test_cancel_all_tasks(self):
        raise NotImplementedError

    @unittest.skip
    async def test_task_mixin(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_task.py
DragonHPC/dragon
import asyncio
from functools import partial
import logging
import unittest
from unittest.mock import call, patch
from uuid import uuid4

from dragon.transport.tcp import errno
from dragon.transport.tcp import messages
from dragon.transport.tcp import transport


def setUpModule():
    """Disable logging while running these tests."""
    logging.disable()


def tearDownModule():
    """Re-enable logging after running these tests."""
    logging.disable(logging.NOTSET)


class TestMessages:
    """Mixin that provides messages for testing."""

    @classmethod
    def setUpClass(cls):
        cls.SendRequest = partial(messages.SendRequest,
            seqno=None, timeout=0.5, channel_sd=b'channel desc',
            return_mode=messages.SendReturnMode.WHEN_BUFFERED, sendhid=uuid4(),
            payload=b'payload',
        )

        cls.RecvRequest = partial(messages.RecvRequest,
            seqno=None, timeout=0.5, channel_sd=b'channel desc',
        )


class TransportTestCase(TestMessages, unittest.IsolatedAsyncioTestCase):

    def setUp(self):
        self.transport = transport.Transport(transport.LOOPBACK_ADDRESS_IPv4)
        self.srcaddr = self.transport.addr
        self.destaddr = self.transport.addr

    def tearDown(self):
        del self.transport

    @patch.object(transport.Transport, '_handle_send')
    async def test_request_response(self, handle_send):
        """Request/Response round-trip

        Verifies the `asyncio.Future` returned from `Transport.write_request`
        is being properly tracked and completed when a corresponding `Response`
        is handled.
        """
        # Write request
        req = self.SendRequest()
        fut = self.transport.write_request(req, self.destaddr)

        # Verify that the seuqnce number was set
        self.assertEqual(req.seqno, 1)

        # Check that _handle_send was called properly
        handle_send.assert_called_once_with(req, self.destaddr)

        # Verify that the returned future is being tracked in _responses
        self.assertIn(req.seqno, self.transport._responses)
        self.assertIs(fut, self.transport._responses[req.seqno])

        # Verify that future is not done
        self.assertFalse(fut.done())

        # Handle corresponding response
        resp = messages.SendResponse(req.seqno)
        self.transport._handle_recv(resp, self.destaddr)

        # Verify the future is no longer being tracked in _responses
        self.assertNotIn(req.seqno, self.transport._responses)

        # Verify the future is done
        self.assertTrue(fut.done())

        # Verify future is completed with response
        fut_resp, fut_addr = await fut
        self.assertIs(fut_resp, resp)
        self.assertIs(fut_addr, self.destaddr)

    @patch.object(transport.Transport, '_handle_send')
    async def test_write_request(self, handle_send):
        # Write request
        req = self.SendRequest()
        fut = self.transport.write_request(req, self.destaddr)

        # Verify that the seuqnce number was set
        self.assertEqual(req.seqno, 1)

        # Check that _handle_send was called properly
        handle_send.assert_called_once_with(req, self.destaddr)

        # Verify that the returned future is being tracked in _responses
        self.assertIn(req.seqno, self.transport._responses)
        self.assertIs(fut, self.transport._responses[req.seqno])

        # Verify that future is not done
        self.assertFalse(fut.done())

    @patch.object(transport.Transport, '_handle_send', side_effect=RuntimeError)
    async def test_write_request_handle_send_exception(self, handle_send):
        """Future response is not managed when exception raised in _handle_send()

        Verifies the `Transport._responses` dictionary is properly cleaned
        up (i.e., the corresponding `asyncio.Future` is removed) if
        `Transport._handle_send` raises an exception.
        """
        # Verify RuntimeError is raised when writing request
        self.assertRaises(RuntimeError, self.transport.write_request, self.SendRequest(), self.transport.addr)

        # Verify that _responses is empty
        self.assertFalse(self.transport._responses)

    @patch.object(transport.Transport, '_handle_send')
    async def test_write_response(self, handle_send):
        resp = messages.SendResponse(seqno=42)
        self.transport.write_response(resp, self.srcaddr)
        handle_send.assert_called_once_with(resp, self.srcaddr)

    @patch.object(transport.Transport, '_handle_recv')
    def test_handle_send(self, handle_recv):
        # Transport._handle_send() only supports messages to itself
        self.assertRaises(AssertionError, self.transport._handle_send, None, None)

        for mode in (messages.SendReturnMode.IMMEDIATELY, messages.SendReturnMode.WHEN_BUFFERED):
            handle_recv.reset_mock()
            req = self.SendRequest(seqno=42, return_mode=mode)
            self.transport._handle_send(req, self.transport.addr)
            handle_recv.assert_has_calls([
                call(messages.SendResponse(req.seqno), self.transport.addr),
                call(req, self.transport.addr),
            ])

        handle_recv.reset_mock()
        req = self.SendRequest(seqno=42, return_mode=messages.SendReturnMode.WHEN_RECEIVED)
        self.transport._handle_send(req, self.transport.addr)
        handle_recv.assert_called_once_with(req, self.transport.addr)

    @patch.object(transport.Transport, '_handle_send')
    async def test_handle_recv(self, handle_send):
        """Sends DRAGON_NOT_IMPLEMENTED error when unsupported message type is received

        Verifies that a `DRAGON_NOT_IMPLEMENTED` error response is sent when
        an unsupported message type is received.
        """
        # Note that Hello messages are only exchanged during connection
        # handshake and are not a `Request` or `Response` and therefore
        # are not supported by `Transport`.
        msg = messages.Hello(host=self.transport.addr.host, port=8080)
        self.transport._handle_recv(msg, self.transport.addr)

        # Expected response
        resp = messages.ErrorResponse(0, errno.DRAGON_NOT_IMPLEMENTED, f'Unsupported type of message: {type(msg)}')
        handle_send.assert_called_once_with(resp, self.transport.addr)

    async def test_handle_recv_response(self):
        # Initialize transport to expect a response
        msg = messages.SendResponse(seqno=42)
        fut = self.transport._responses[msg.seqno] = asyncio.Future()

        # Receive response message
        self.transport._handle_recv(msg, self.destaddr)

        # Verify the future is no longer being tracked in _responses
        self.assertNotIn(msg.seqno, self.transport._responses)

        # Verify that future is completed
        self.assertTrue(fut.done())
        resp, addr = await fut
        self.assertIs(resp, msg)
        self.assertIs(addr, self.destaddr)

    async def test_handle_recv_unexpected_response(self):
        msg = messages.SendResponse(seqno=42)
        self.assertFalse(self.transport._responses)
        # _handle_recv() must silently fail since there is not a corresponding
        # future.
        self.transport._handle_recv(msg, self.transport.addr)

    async def test_handle_recv_request(self):
        with patch.object(self.transport._requests, 'put_nowait') as put_nowait:
            msg = self.SendRequest(seqno=42)
            self.transport._handle_recv(msg, self.transport.addr)
            self.assertEqual(self.transport._last_request[self.transport.addr], msg.seqno)
            put_nowait.assert_called_once_with((msg, self.transport.addr))

    async def test_handle_recv_replayed_requests(self):
        with patch.object(self.transport._requests, 'put_nowait') as put_nowait:
            # Send original request
            msg = self.SendRequest(seqno=42)
            self.transport._handle_recv(msg, self.transport.addr)

            # Verify the last request seqno is correct
            self.assertEqual(self.transport._last_request[self.transport.addr], msg.seqno)
            # Verify the request was put on the _requests queue
            put_nowait.assert_called_once_with((msg, self.transport.addr))

            # Reset mock to see evaluate replays
            put_nowait.reset_mock()

            # Replay previous requests
            self.transport._handle_recv(msg, self.transport.addr)
            self.transport._handle_recv(self.SendRequest(seqno=9), self.transport.addr)

            # Verify the last request seqno is unaffected
            self.assertEqual(self.transport._last_request[self.transport.addr], msg.seqno)
            # Verify the replayed requests were NOT put on the _requests queue
            put_nowait.assert_not_called()

    async def test_read_request(self):
        msg = self.SendRequest(seqno=42)
        self.transport._requests.put_nowait((msg, self.transport.addr))

        req, addr = await self.transport.read_request()
        self.assertIs(req, msg)
        self.assertIs(addr, self.transport.addr)

    async def test_read_response(self):
        msg = messages.SendResponse(seqno=42)
        self.transport._responses[msg.seqno] = asyncio.Future()
        #asyncio.get_event_loop().call_soon(fut.set_result, (msg, self.destaddr))
        asyncio.get_event_loop().call_soon(self.transport._handle_recv, msg, self.destaddr)
        resp, addr = await self.transport.read_response(msg.seqno)
        self.assertIs(resp, msg)
        self.assertIs(addr, self.destaddr)
        self.assertNotIn(resp.seqno, self.transport._responses)


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_transport.py
DragonHPC/dragon
import unittest


class UtilsTestCase(unittest.TestCase):

    @unittest.skip
    def test_unget_nowait(self):
        raise NotImplementedError

    @unittest.skip
    def test_seconds_remaining(self):
        raise NotImplementedError

    @unittest.skip
    def test_mem_descr_msg(self):
        pass

    @unittest.skip
    def test_mem_pool_msg(self):
        pass

    @unittest.skip
    def test_create_msg(self):
        pass

    @unittest.skip
    def test_attach_channel(self):
        raise NotImplementedError

    @unittest.skip
    def test_open_handle(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/tcp/test_util.py
DragonHPC/dragon
from collections import ChainMap
import logging
import os
import shlex
import time
import unittest

from dragon.channels import Channel
from dragon.infrastructure import facts as dfacts
from dragon.infrastructure import messages as dmsg
from dragon.infrastructure.connection import Connection
from dragon.launcher.util import next_tag
from dragon.managed_memory import MemoryPool
from dragon.transport import start_transport_agent
from dragon.utils import B64


LOGGER = logging.getLogger('test.transport')
NUM_GW_CHANNELS_PER_NODE=1


def setUpModule():
    logging.basicConfig(
        format='%(asctime)s %(levelname)-9s %(process)d %(name)s:%(filename)s:%(lineno)d %(message)s',
        level=logging.NOTSET,
    )
    logging.disable(logging.NOTSET)


def tearDownModule():
    logging.disable()


class LocalServicesInterface:

    NUM_NODES = 1
    ENV = None

    @classmethod
    def setUpClass(cls):
        if not hasattr(cls, 'ARGS'):
            raise NotImplementedError('Missing args to launch transport agent')

    def setUp(self):
        LOGGER.debug('Setting up')

        self.nodes = {
            i: {
                'host_name': f'node-{i}',
                'host_id': str(i),
                'ip_addrs': [f'127.0.0.1:{i+8000}'],
            } for i in range(self.NUM_NODES)
        }

        # Create channels for each node
        for i, n in self.nodes.items():
            # Create memory pool
            n['mempool'] = MemoryPool(
                int(dfacts.DEFAULT_SINGLE_DEF_SEG_SZ),
                f'{os.getuid()}_{os.getpid()}_{i}' + dfacts.DEFAULT_POOL_SUFFIX,
                dfacts.default_pool_muid_from_index(i),
            )
            # Create agent I/O connection
            n['conn'] = Connection(
                inbound_initializer=Channel(n['mempool'], dfacts.shepherd_cuid_from_index(i)),
                outbound_initializer=Channel(n['mempool'], dfacts.transport_cuid_from_index(i)),
            )
            # Set shepherd's channel (i.e., input for local services)
            n['shep_cd'] = B64.bytes_to_str(n['conn'].inbound_chan.serialize())
            # Create gateway channels
            n['gw_chs'] = [Channel(n['mempool'], dfacts.gw_cuid_from_index(i, NUM_GW_CHANNELS_PER_NODE) + j) for j in range(NUM_GW_CHANNELS_PER_NODE)]

        LOGGER.info('Memory pools and channels created')

        # Create LAChannelsInfo
        la_ch_info = dmsg.LAChannelsInfo(
            tag=next_tag(),
            nodes_desc=self.nodes,
            gs_cd='',
            num_gw_channels=NUM_GW_CHANNELS_PER_NODE,
        )

        LOGGER.info(f'LAChannelsInfo created: {la_ch_info.get_sdict()}')

        # Start agents
        for i, n in self.nodes.items():
            n['agent'] = start_transport_agent(
                node_index=str(i),
                in_ch_sdesc=B64(n['conn'].outbound_chan.serialize()),
                args=self.ARGS,
                env=self.ENV,
                gateway_channels=n['gw_chs'],
            )

        LOGGER.info('Agents started')

        # Send LAChannelsInfo
        for n in self.nodes.values():
            n['conn'].send(la_ch_info.serialize())

        LOGGER.info('Sent all LAChannelsInfo')

        # Wait for each node to reply with pings
        for i, n in self.nodes.items():
            ping_msg = dmsg.parse(n['conn'].recv())
            assert isinstance(ping_msg, dmsg.TAPingSH), f'Did not receive TAPingSH from node {i}'

        LOGGER.info('Received all TAPingSH')
        LOGGER.debug('Setup completed')

    def tearDown(self):
        LOGGER.info('Tearing down')

        # Halt
        halt_msg = dmsg.SHHaltTA(next_tag()).serialize()
        for n in self.nodes.values():
            n['conn'].send(halt_msg)

        LOGGER.info('Sent all SHHaltTA')

        # Wait for each node to reply
        for i, n in self.nodes.items():
            halted_msg = dmsg.parse(n['conn'].recv())
            assert isinstance(halted_msg, dmsg.TAHalted), f'Did noot receive TAHalted from node {i}'

        LOGGER.info('Received all TAHalted')

        for n in self.nodes.values():
            # Stop agent
            n['agent'].terminate()
            n['agent'].wait(timeout=3)
            n['agent'].kill()
            # Destroy gateway channels
            for ch in n['gw_chs']:
                ch.destroy()
            # Close the connection
            n['conn'].close()
            # Destroy I/O channels since they're externally managed
            n['conn'].outbound_chan.destroy()
            n['conn'].inbound_chan.destroy()
            # Destroy memory pool
            n['mempool'].destroy()

        LOGGER.debug('Teardown completed')

    def test_allow_agent_control_loops_to_poll_once(self):
        # TODO Send a GatewayMessage to each agent and verify it arrived
        # Wait for control loop to poll at least once
        time.sleep(6)


class TCPTestCase(LocalServicesInterface, unittest.TestCase):

    NUM_NODES = 3
    ARGS = shlex.split(f'tcp --log-level {logging.getLevelName(LOGGER.getEffectiveLevel())} --no-dragon-logging')


@unittest.skip
class HSTATestCase(LocalServicesInterface, unittest.TestCase):

    ARGS = shlex.split(f'hsta --log-level {logging.getLevelName(LOGGER.getEffectiveLevel())} --no-dragon-logging')


if __name__ == '__main__':
    unittest.main()

test/transport/test_lsif.py
DragonHPC/dragon
import unittest
import time
import os
import random
import pickle
import logging

logging.captureWarnings(True)

import multiprocessing as mp
from dragon.channels import Channel, Message, ChannelSendH, ChannelRecvH, register_gateways_from_env, discard_gateways, \
    MASQUERADE_AS_REMOTE, POLLIN
from dragon.managed_memory import MemoryPool
from dragon.utils import B64
import asyncio
from collections import defaultdict

from dragon.mpbridge.queues import DragonQueue
import dragon.infrastructure.connection as dconn

SUCCESS = 0
FAILURE = 1
TIMEOUT = 2

SILENT = True

def status(*args, **kwargs):
    if not SILENT:
        print(*args, **kwargs)


from dragon.transport.tcp.transport import Transport
from dragon.transport.tcp.transport import LOOPBACK_ADDRESS_IPv4
from dragon.transport.tcp.server import Server
from dragon.transport.tcp.client import Client
from dragon.transport.tcp.task import cancel_all_tasks

GW_ENV_PREFIX = "DRAGON_GW"

def mk_msg(pool, nbytes):
    msg = Message.create_alloc(pool, nbytes)
    mview = msg.bytes_memview()
    for i in range(nbytes):
        mview[i] = i % 256

    return msg

def validate_msg(msg):
    # Here we can do any message validation that needs
    # to be done.
    mview = msg.bytes_memview()

    for i in range(len(mview)):
        if (mview[i] != i%256):
            print(f'The message contents were incorrectly {mview[i]} at location {i}')
            return i

    return len(mview)

async def agent(channels, control):
    # Create a Transport instance to handle Request/Response I/O.
    # By default, Transport._handle_send() implements a loopback
    # strategy by calling Tranport._handle_recv().
    transport = Transport(LOOPBACK_ADDRESS_IPv4)

    # Create a Server instance to process Requests (i.e., remote
    # channel operations).
    server = Server(transport)

    # Create Client instances for each gateway channel to send
    # Requests to the Server and process Responses in order to
    # complete gateway messages.
    clients = {}

    # Using the default transport, the node mapping should
    # ensure every address is the transport's address, for
    # consistency.
    nodes = defaultdict(lambda: transport.addr)

    # Assuming the gateway channels have been created and
    # CHANNELS is a list of the corresponding names...
    for name in channels:
        gw_ser_str = os.environ[name]
        gw_ser = B64.from_str(gw_ser_str).decode()
        clients[name] = Client(gw_ser, transport, nodes)

    # Note that Transport is not a TaskMixin as there is nothing that
    # requires a separate thread of execution, so there is no need to
    # call transport.start() or await transport.wait_started() as
    # is the case with StreamTransport.

    # Start the server
    server.start(name=f'Transport-server')

    status('Transport service started.')

    # Start each client
    for name, client in clients.items():
        client.start(name=f'Transport-client-{name}')

    # At this point, the server and client tasks are running and you
    # can do whatever, but if you need a control loop...
    try:
        while control.value == 1:
            await asyncio.sleep(5.0)
    finally:
        # Shut down clients and server
        status('Transport service is shutting down')
        await asyncio.gather(*[c.stop() for c in clients.values()], return_exceptions=True)
        try:
            await server.stop()
        except:
            pass
        # Cancel all remaining tasks for good measure
        await cancel_all_tasks()

def start_transport(control):
    done = False
    id = 1
    channels = []
    while not done:
        id_str = GW_ENV_PREFIX+str(id)
        if id_str in os.environ:
            channels.append(id_str)
            id += 1
        else:
            done = True

    status('Added gateway channels to environment. Now starting transport service.')

    logging.disable(logging.ERROR)
    asyncio.run(agent(channels, control))

def add_gw_to_env(gw_ch, id):
    ser_gw = gw_ch.serialize()
    encoded_ser_gw = B64(ser_gw)
    os.environ[GW_ENV_PREFIX+str(id)] = str(encoded_ser_gw)

def echo_gw_and_attach():
    gw_channels = []
    id = 1
    id_str = GW_ENV_PREFIX + str(id)
    while id_str in os.environ:
        ser_gw = B64.from_str(os.environ[id_str])
        status("Found", id_str, "in environment")
        gw_ch = Channel.attach(ser_gw.decode())
        status('Attached to gateway channel', id)
        gw_channels.append(gw_ch)
        id += 1
        id_str = GW_ENV_PREFIX + str(id)

def get_msg_from_channel(ch_ser, msg_nbytes, ret_value, timeout=None):
    ch = Channel.attach(ch_ser)
    recvh = ChannelRecvH(ch)
    recvh.open()
    try:
        msg = recvh.recv(timeout=timeout)
        # Here we can do any message validation that needs
        # to be done.
        mview = msg.bytes_memview()
        if (validate_msg(msg) == msg_nbytes):
            ret_value.value = SUCCESS
            status("Received message sent to masquerading remote channel and verified its contents.")
        else:
            ret_value.value = FAILURE
        recvh.close()
        ch.detach()
    except Exception as ex:
        print("Exception getting message from channel.")
        print(ex)
        ret_value.value = FAILURE

def send_msg_to_channel(ch_ser, msg_nbytes, ret_value, timeout=None, sleep_secs=0):
    time.sleep(sleep_secs)
    ch = Channel.attach(ch_ser)
    pool = ch.get_pool()
    sendh = ChannelSendH(ch)
    sendh.open()
    msg = mk_msg(pool, msg_nbytes)
    try:
        sendh.send(msg, timeout=timeout)
        sendh.close()
        ch.detach()
        ret_value.value = SUCCESS
    except Exception as ex:
        print("Exception sending message to channel.")
        print(ex)
        ret_value.value = FAILURE

class SingleNodeTransportBench(unittest.TestCase):
    """Unit tests for the Single Node emulation of a transport service.
    """

    @classmethod
    def setUpClass(cls):
        pool_name = f"dragon_transport_test_{os.getpid()}"
        pool_size = 1073741824  # 1GB Pool
        pool_uid = 1
        cls.mpool = MemoryPool(pool_size, pool_name, pool_uid)
        cls.gw_ch1 = Channel(cls.mpool, c_uid=1, block_size=256)
        add_gw_to_env(cls.gw_ch1, 1)

        cls.transport_running = mp.Value('i', 1, lock=False)
        cls.transport_proc = mp.Process(target=start_transport, args=(cls.transport_running,))
        cls.transport_proc.start()

        register_gateways_from_env()

        cls.opts = dconn.ConnectionOptions(default_pool=cls.mpool)

    @classmethod
    def tearDownClass(cls):
        cls.transport_running.value = 0
        cls.transport_proc.join()
        discard_gateways()
        cls.gw_ch1.destroy()
        cls.mpool.destroy()

    def setUp(self):
        self.user_ch1 = Channel(self.mpool, 3, flags=MASQUERADE_AS_REMOTE)
        self.user_ch2 = Channel(self.mpool, 4, flags=MASQUERADE_AS_REMOTE)

    def tearDown(self):
        self.user_ch1.destroy()
        self.user_ch2.destroy()

    def test_setup_and_teardown(self):
        pass

    def test_encoding_decoding(self):
        x = b'Hello World'
        y = B64(x)
        self.assertEqual(str(y),"SGVsbG8gV29ybGQ=")
        self.assertEqual(y.decode(),x)
        x = b'How are you?'
        y = B64(x)
        self.assertEqual(str(y),"SG93IGFyZSB5b3U/")
        self.assertEqual(y.decode(),x)

    def test_environment_attach(self):
        proc = mp.Process(target=echo_gw_and_attach)
        proc.start()
        proc.join()
        self.assertEqual(proc.exitcode, 0)

    def test_send(self):
        sendh = ChannelSendH(self.user_ch1)
        sendh.open()
        msg = mk_msg(self.mpool, 1)
        sendh.send(msg)
        # Pass the serialized original self.user_ch1
        # to get_msg_from_channel to verify that the
        # transport service forwarded the sent message
        # to the channel. Then that, combined with the
        # completion of the send in this test, will
        # verify the complete path from bootstrapping the
        # gw channels, through channels, the gw, and the
        # transport service and back out the other end.
        ch_ser = self.user_ch1.serialize()
        ret_value = mp.Value('i', FAILURE, lock=False)
        proc = mp.Process(target=get_msg_from_channel, args=(ch_ser, 1, ret_value))
        proc.start()
        proc.join()
        self.assertEqual(proc.exitcode, 0)
        self.assertEqual(ret_value.value, SUCCESS)

    def test_send_serialized(self):
        sendh = ChannelSendH(self.user_ch1)
        sendh.open()
        msg = mk_msg(self.mpool, 512)
        sendh.send(msg)
        # Pass the serialized original self.user_ch1
        # to get_msg_from_channel to verify that the
        # transport service forwarded the sent message
        # to the channel. Then that, combined with the
        # completion of the send in this test, will
        # verify the complete path from bootstrapping the
        # gw channels, through channels, the gw, and the
        # transport service and back out the other end.
        ch_ser = self.user_ch1.serialize()
        ret_value = mp.Value('i', FAILURE, lock=False)
        proc = mp.Process(target=get_msg_from_channel, args=(ch_ser, 512, ret_value))
        proc.start()
        proc.join()
        self.assertEqual(proc.exitcode, 0)
        self.assertEqual(ret_value.value, SUCCESS)

    def test_receive_timeout(self):
        recvh = ChannelRecvH(self.user_ch1)
        recvh.open()
        # No message is sent to the channel. This forces a timeout within
        # the transport agent since here the channel looks like a remote
        # channel.
        try:
            self.assertRaises(TimeoutError, recvh.recv, timeout=2)
        finally:
            recvh.close()

    def test_receive(self):
        recvh = ChannelRecvH(self.user_ch1)
        recvh.open()
        ch_ser = self.user_ch1.serialize()
        ret_value = mp.Value('i', FAILURE, lock=False)
        proc = mp.Process(target=send_msg_to_channel, args=(ch_ser, 1, ret_value))
        proc.start()
        msg = recvh.recv()
        # Here we can do any message validation that needs
        # to be done.
        mview = msg.bytes_memview()
        self.assertEqual(validate_msg(msg), 1, 'The message did not have the correct contents')
        status('Received message from masquerading remote channel via remote receive and verified its contents.')
        recvh.close()
        proc.join()
        self.assertEqual(proc.exitcode, 0)
        self.assertEqual(ret_value.value, SUCCESS)

    def test_poll(self):
        status("Starting test_poll")
        #try:
        recvh = ChannelRecvH(self.user_ch1)
        recvh.open()
        ch_ser = self.user_ch1.serialize()
        ret_value = mp.Value('i', FAILURE, lock=False)
        proc = mp.Process(target=send_msg_to_channel, args=(ch_ser, 1, ret_value, None, 2))
        proc.start()
        rc = self.user_ch1.poll(event_mask=POLLIN)
        self.assertEqual(rc, True)
        msg = recvh.recv()
        # Here we can do any message validation that needs
        # to be done.
        mview = msg.bytes_memview()
        self.assertEqual(validate_msg(msg), 1, 'The message did not have the correct contents')
        status('Polled for and received message from masquerading remote channel via remote receive and verified its contents.')
        recvh.close()
        proc.join()
        self.assertEqual(proc.exitcode, 0)
        self.assertEqual(ret_value.value, SUCCESS)
        #except Exception as ex:
        #    self.assertTrue(False, "There was an exception in the test code.")

    @classmethod
    def _queue_proc(self, q, obj1, obj2):
        msg = q.get()
        assert obj1 == msg
        q.put(obj2)

        q.close()

    def test_queue_with_external_channel(self):
        # In the case of an externally managed channel, we need to
        # provide a memory pool which will be used as the buffer_pool
        # for the Many2ManyWriting adapter.
        # This is because we don't call the Dragon start method and
        # dparms.this_process.default_pd that we use in case of a
        # remote channel is not set in Queue.
        q = DragonQueue(m_uid=self.mpool, _ext_channel=self.user_ch1)

        y = [1, 2, 3]
        q.put(y)
        self.assertEqual(q.qsize(), 1, "We could not get the proper size from a queue on a remote channel.")
        x = q.get()
        self.assertEqual(x, y, "The Queue implementation did not pass a list correctly through it.")
        y = ['a', 'b', 'c']
        q.put(y)
        x = q.get()
        self.assertEqual(x, y, "The Queue implementation did not pass a list correctly through it.")

        obj1 = {"hyena": "hyena"}
        obj2 = bytes(10000)
        q.put(obj1)
        proc = mp.Process(target=self._queue_proc, args=(q, obj1, obj2,))
        proc.start()
        proc.join()
        self.assertEqual(proc.exitcode, 0)
        msg = q.get()
        self.assertEqual(msg, obj2)

        q.close()

    def test_queue_ordering_with_external_channel(self):
        # In the case of an externally managed channel, we need to
        # provide a memory pool which will be used as the buffer_pool
        # for the Many2ManyWriting adapter.
        # This is because we don't call the Dragon start method and
        # dparms.this_process.default_pd that we use in case of a
        # remote channel is not set in Queue.
        q = DragonQueue(m_uid=self.mpool, _ext_channel=self.user_ch1)

        # This test assumes ordering on queue messages and tests that
        y = [1, 2, 3]
        q.put(y)
        x = q.get()
        self.assertEqual(x, y, "The Queue implementation did not pass a list correctly through it.")
        yy = ['a', 'b', 'c']
        q.put(yy)
        q.put(y)
        x = q.get()
        self.assertEqual(x, yy, "The Queue implementation did not pass a list correctly through it.")
        x = q.get()
        self.assertEqual(x, y, "The Queue implementation did not pass a list correctly through it.")

        q.close()

    @staticmethod
    def _connection_proc(some_msg, connin, connout):
        if connin:
            msg = connin.recv()
            assert some_msg == msg
            connin.ghost_close()
        if connout:
            connout.send(some_msg)
            connout.ghost_close()

    def test_connection_with_external_channels(self):
        msg = random.randbytes(1024)

        connout = dconn.Connection(outbound_initializer=self.user_ch1, options=self.opts)
        connin = dconn.Connection(inbound_initializer=self.user_ch2, options=self.opts)
        pconnout = dconn.Connection(outbound_initializer=self.user_ch2, options=self.opts)
        pconnin = dconn.Connection(inbound_initializer=self.user_ch1, options=self.opts)
        proc = mp.Process(target=self._connection_proc, args=(msg, pconnin, pconnout, ))
        proc.start()

        connout.send(msg)
        msg_back = connin.recv()
        self.assertEqual(msg, msg_back)

        proc.join()
        self.assertEqual(proc.exitcode, 0)
        connout.ghost_close()
        connin.ghost_close()
        pconnin.ghost_close()
        pconnout.ghost_close()

    @staticmethod
    def _connection_proc_bytes(some_bytes, connin, connout):
        if connout:
            connout.send_bytes(some_bytes)
            connout.ghost_close()
        if connin:
            rec_bytes = connin.recv_bytes()
            assert some_bytes == rec_bytes
            connin.ghost_close()

    def test_connection_with_external_channels_test_proc_send_bytes(self):
        connin = dconn.Connection(inbound_initializer=self.user_ch2, options=self.opts)
        pconnout = dconn.Connection(outbound_initializer=self.user_ch2, options=self.opts)

        for lg_size in range(25):
            some_bytes = random.randbytes(2 ** lg_size)
            proc = mp.Process(target=self._connection_proc_bytes, args=(some_bytes, None, pconnout, ))
            proc.start()

            rec_bytes = connin.recv_bytes()
            self.assertEqual(some_bytes, rec_bytes)
            proc.join()
            self.assertEqual(proc.exitcode, 0)

        connin.ghost_close()
        pconnout.ghost_close()

    def test_connection_with_external_channels_test_proc_recv_bytes(self):
        connout = dconn.Connection(outbound_initializer=self.user_ch1, options=self.opts)
        pconnin = dconn.Connection(inbound_initializer=self.user_ch1, options=self.opts)

        # The following for loop fails and we have added the above lines to at least test one case of
        # having the process receiving bytes
        for lg_size in range(25):
            some_bytes = random.randbytes(2 ** lg_size)
            proc = mp.Process(target=self._connection_proc_bytes, args=(some_bytes, pconnin, None, ))
            proc.start()
            connout.send_bytes(some_bytes)
            proc.join()
            self.assertEqual(proc.exitcode, 0)

        connout.ghost_close()
        pconnin.ghost_close()

    def pass_an_obj(self, obj):
        writer = dconn.Connection(outbound_initializer=self.user_ch1, options=self.opts)
        reader = dconn.Connection(inbound_initializer=self.user_ch1, options=self.opts)

        proc = mp.Process(target=self._connection_proc, args=(obj, None, writer, ))
        proc.start()

        rec_obj = reader.recv()
        proc.join()
        self.assertEqual(proc.exitcode, 0)
        self.assertEqual(obj, rec_obj)

        writer.ghost_close()
        reader.ghost_close()

    def test_connection_with_external_channels_test_objs(self):
        self.pass_an_obj('hyena')
        self.pass_an_obj({'hyena': 'hyena'})

        self.pass_an_obj(bytes(100000))

        self.pass_an_obj(bytes(200000))

        self.pass_an_obj(bytearray(2 ** 20))

        for lg_size in range(25):
            self.pass_an_obj(bytearray(2 ** lg_size))

        for lg_size in range(25):
            self.pass_an_obj(random.randbytes(2 ** lg_size))

    def test_connection_with_external_channels_test_recv_bytes_from_send(self):
        # required for multiprocessing.Connection unit tests...
        writer = dconn.Connection(outbound_initializer=self.user_ch1, options=self.opts)
        reader = dconn.Connection(inbound_initializer=self.user_ch1, options=self.opts)

        my_obj = {'hyenas': 'are awesome'}
        writer.send(my_obj)
        # would like this to raise, but....
        ser_obj = reader.recv_bytes()
        rec_obj = pickle.loads(ser_obj)
        self.assertEqual(my_obj, rec_obj)

        writer.ghost_close()
        reader.ghost_close()

    def test_connection_with_external_channels_test_recv_from_send_bytes(self):
        # again, weird usage but multiprocessing.Connection unit tests need
        # this functionality

        writer = dconn.Connection(outbound_initializer=self.user_ch1, options=self.opts)
        reader = dconn.Connection(inbound_initializer=self.user_ch1, options=self.opts)

        my_obj = {'hyenas': 'are awesome'}
        ser_obj = pickle.dumps(my_obj)
        writer.send_bytes(ser_obj)
        rec_obj = reader.recv()
        self.assertEqual(my_obj, rec_obj)

        writer.ghost_close()
        reader.ghost_close()

    @staticmethod
    def _connection_proc_poll(msg, conn):
        rc = conn.poll(timeout=20.0)
        assert rc == True
        msg_recv = conn.recv()
        assert msg == msg_recv

    def test_connection_with_external_channels_test_poll(self):
        msg = random.randbytes(1024)
        writer = dconn.Connection(outbound_initializer=self.user_ch1, options=self.opts)
        reader = dconn.Connection(inbound_initializer=self.user_ch1, options=self.opts)

        proc = mp.Process(target=self._connection_proc_poll, args=(msg, reader, ))
        proc.start()
        time.sleep(3)
        writer.send(msg)
        proc.join()
        self.assertEqual(proc.exitcode, 0)

        writer.ghost_close()
        reader.ghost_close()


if __name__ == "__main__":
    mp.set_start_method("spawn")
    unittest.main()

test/transport/test_tcp_transport.py
DragonHPC/dragon
import unittest

from dragon.transport.x509 import CertificateAuthority, dragon_uri, generate_server_csr, generate_server_self_signed_cert, pem_encode, pem_encode_pkcs8


class X509TestCase(unittest.TestCase):

    @unittest.skip
    def test_generate_private_key(self):
        raise NotImplementedError

    @unittest.skip
    def test_pem_encode_pkcs8(self):
        raise NotImplementedError

    @unittest.skip
    def test_certificate_builder(self):
        raise NotImplementedError

    @unittest.skip
    def test_add_ca_key_usage_extensions(self):
        raise NotImplementedError

    @unittest.skip
    def test_add_server_key_usage_extensions(self):
        raise NotImplementedError

    @unittest.skip
    def test_sign(self):
        raise NotImplementedError

    @unittest.skip
    def test_pem_encode(self):
        raise NotImplementedError

    @unittest.skip
    def test_get_common_name(self):
        raise NotImplementedError

    @unittest.skip
    def test_check_subject_alternative_name(self):
        raise NotImplementedError

    @unittest.skip
    def test_server_csr_builder(self):
        raise NotImplementedError

    @unittest.skip
    def test_generate_server_csr(self):
        raise NotImplementedError

    @unittest.skip
    def test_generate_server_self_signed_cert(self):
        raise NotImplementedError

    @unittest.skip
    def test_dragon_uri(self):
        raise NotImplementedError


class CATestCase(unittest.TestCase):

    @unittest.skip
    def test_generate(self):
        raise NotImplementedError

    @unittest.skip
    def test_issue_server_certificate(self):
        raise NotImplementedError


if __name__ == '__main__':
    unittest.main()

test/transport/test_x509.py
DragonHPC/dragon
#!/usr/bin/env python3

import unittest
import random
import string
import pickle
import math
import multiprocessing as mp
from dragon.managed_memory import MemoryPool, MemoryAlloc, MemoryAllocations, MemoryPoolAttr, \
    DragonPoolError, DragonPoolAttachFail, DragonPoolCreateFail, DragonMemoryError, AllocType

from dragon.dtypes import DragonError

# limits are for tests
MIN_POOL_SIZE = 32768
MAX_POOL_SIZE = 2 ** 31
DEFAULT_UID = 123456

# NOTE: The following function is dependent on the serialization of memory descriptors.
# It relies on the following format:
#    bytes offset     description
#    0-7              length of pool serialized descriptor
#    8-15             m_uid of pool
#    16-23            hostid of pool
#    ...              No dependencies beginning at byte offset 24
def mk_remote_mem_ser(local_mem_ser):
    mem_ser_array = bytearray(local_mem_ser)
    mem_ser_array[8] = 99 # change m_uid
    mem_ser_array[16] = 99 # change hostid
    mem_ser_remote = bytes(mem_ser_array)
    return mem_ser_remote

# NOTE: The following function is dependent on the serialization of memory pool descriptors.
# It relies on the following format:
#    bytes offset     description
#    0-7              m_uid of pool
#    8-15             hostid of pool
#    ...              No dependencies beginning at byte offset 16
def mk_remote_pool_ser(local_pool_ser):
    pool_ser_array = bytearray(local_pool_ser)
    pool_ser_array[0] = 99 # change m_uid
    pool_ser_array[8] = 99 # change hostid
    pool_ser_remote = bytes(pool_ser_array)
    return pool_ser_remote


class MemPoolCreateTest(unittest.TestCase):
    seed = 0

    @classmethod
    def setUpClass(cls):
        cls.seed = random.getrandbits(100)
        random.seed(a=cls.seed)
        # print(f'Our random seed is {cls.seed}')

    def test_create(self):
        size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        mpool = MemoryPool(size, "mpool_test", DEFAULT_UID, None)
        self.assertIsInstance(mpool, MemoryPool,
                              msg=f'Create failed to produce a memory pool of size {size}')
        mpool.destroy()

    def test_destroy(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "mpool_test", DEFAULT_UID, None)
        pool_ser = mpool.serialize()
        mpool.destroy()
        with self.assertRaises(DragonPoolAttachFail):
            _ = MemoryPool.attach(pool_ser)

    @unittest.skip("Skipping for now until bug in ref counting is found")
    def test_attach(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "mpool_test", DEFAULT_UID, None)
        pool_ser = mpool.serialize()
        mpool2 = MemoryPool.attach(pool_ser)
        mpool2.detach()
        mpool.destroy()

    @unittest.skip("Skipping for now until bug in ref counting is found")
    def test_attach_ref_count(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "mpool_test", DEFAULT_UID, None)
        pool_ser = mpool.serialize()
        mpool.detach()

        mpools = []
        for i in range(10):
            mpools.append(MemoryPool.attach(pool_ser))

        for i in range(9):
            mpools[i].detach()

        # If any of the other detaches had removed our knowledge of the pool, this should fail
        # Success means we have not detached
        mpools[9].destroy()

    def test_serialized_info_no_attach(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "mpool_test", DEFAULT_UID, None)
        pool_ser = mpool.serialize()

        (uid, fname) = MemoryPool.serialized_uid_fname(pool_ser)
        try:
            self.assertEqual(uid, DEFAULT_UID)
            self.assertEqual(fname, "/_dragon_mpool_test_manifest")
        finally:
            mpool.destroy()

    def test_destroy_not_created(self):
        # This doesn't really make sense to do, but it does validate destroy errors correctly
        mpool = MemoryPool.__new__(MemoryPool)
        with self.assertRaises(DragonPoolError,
                               msg=f'Destroying an uncreated pool did not raise error'):
            mpool.destroy()

    def test_negative_size(self):
        with self.assertRaises(OverflowError,
                               msg=f'Creating a negative size pool did not raise an OverflowError'):
            _ = MemoryPool(-12, "mpool_test", DEFAULT_UID, None)

    def test_negative_uid(self):
        with self.assertRaises(OverflowError):
            _ = MemoryPool(MIN_POOL_SIZE, "mpool_test", -1, None)

    def test_zero_uid(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "mpool_test", 0, None)
        self.assertIsInstance(mpool, MemoryPool)
        mpool.destroy()

    def test_create_prealloc_blocks(self):
        prealloc_blocks = [8,16,11]
        mpool = MemoryPool(MAX_POOL_SIZE, "mpool_test", DEFAULT_UID, prealloc_blocks)
        self.assertIsInstance(mpool, MemoryPool,
                              msg=f'Failed to create memory pool of size f{MAX_POOL_SIZE} with pre-allocated blocks')
        mpool.destroy()

    def test_create_large_mpool(self):
        mpool = MemoryPool(MAX_POOL_SIZE, "mpool_test", DEFAULT_UID, None)
        self.assertIsInstance(mpool, MemoryPool,
                              msg=f'Failed to create memory pool of size f{MAX_POOL_SIZE}')
        mpool.destroy()

    def test_create_small_mpool(self):
        with self.assertRaises(DragonPoolCreateFail, msg=f'Creating an undersized pool did not raise a RuntimeError'):
            self.mpool = MemoryPool(100, "mpool_test", DEFAULT_UID, None)

    def test_create_zero_mpool(self):
        with self.assertRaises(DragonPoolCreateFail,
                               msg=f'Failed to raise RuntimeError on creation of 0 size memory pool'):
            _ = MemoryPool(0, "mpool_test", DEFAULT_UID, None)

    def test_create_empty_name_mpool(self):
        with self.assertRaises(DragonPoolCreateFail):
            _ = MemoryPool(MIN_POOL_SIZE, "", DEFAULT_UID, None)

    def test_create_large_uid_mpool(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "mpool_test", 2 ** 63 + 1234, None)
        self.assertIsInstance(mpool, MemoryPool)
        mpool.destroy()

    def test_create_numerical_name_mpool(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "12345", DEFAULT_UID, None)
        self.assertIsInstance(mpool, MemoryPool)
        mpool.destroy()

    def test_create_int_name_mpool(self):
        with self.assertRaises(TypeError):
            _ = MemoryPool(MIN_POOL_SIZE, 12345, DEFAULT_UID, None)

    def test_create_string_bytes(self):
        with self.assertRaises(TypeError):
            _ = MemoryPool("32768", "mpool_test", DEFAULT_UID, None)

    def test_create_string_uid(self):
        with self.assertRaises(TypeError):
            _ = MemoryPool(MIN_POOL_SIZE, "mpool_test", "1", None)

    def test_create_two(self):
        size1 = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        size2 = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        mpool1 = MemoryPool(size1, "mpool_test1", DEFAULT_UID, None)
        mpool2 = MemoryPool(size2, "mpool_test2", DEFAULT_UID + 1, None)
        self.assertIsInstance(mpool1, MemoryPool)
        self.assertIsInstance(mpool2, MemoryPool)
        mpool1.destroy()
        mpool2.destroy()

    def test_create_two_mpool_same_name(self):
        mpool1 = MemoryPool(34567, "mpool_test", DEFAULT_UID, None)
        with self.assertRaises(DragonPoolCreateFail):
            _ = MemoryPool(45678, "mpool_test", DEFAULT_UID + 1, None)
        mpool1.destroy()

    def test_create_two_mpool_same_uid(self):
        mpool1 = MemoryPool(34567, "mpool_test", DEFAULT_UID, None)
        with self.assertRaises(DragonPoolCreateFail):
            _ = MemoryPool(45678, "mpool_test", DEFAULT_UID, None)
        mpool1.destroy()

    def test_create_two_mpool_same_name_uid(self):
        mpool1 = MemoryPool(34567, "mpool_test", DEFAULT_UID, None)
        with self.assertRaises(DragonPoolCreateFail):
            _ = MemoryPool(34567, "mpool_test", DEFAULT_UID, None)
        mpool1.destroy()

    def test_create_destroy_many(self):
        for i in range(random.randint(25, 150)):
            size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
            mpool = MemoryPool(size, "mpool_test", DEFAULT_UID, None)
            self.assertIsInstance(mpool, MemoryPool)
            mpool.destroy()

    def test_create_max_name(self):
        size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        # Expect shm_open() to limit names to _POSIX_PATH_MAX, which is
        # typically 256 bytes. Recall that names are prefixed with `"/" +
        # DRAGON_MEMORY_DEFAULT_FILE_PREFIX` (e.g., `"_dragon_"`) and suffixed
        # with `"_manifest."`, see _alloc_pool_shm(). This implies the longest
        # possible name given to `MemoryPool.create()` is 236 bytes, after being
        # UTF-8 encoded and accounting for the terminating `\0` byte.
        name = "a" * 236
        mpool = MemoryPool(size, name, 1, None)
        self.assertIsInstance(mpool, MemoryPool)
        mpool.destroy()

    #@unittest.skip("Limit is 236, bugfix PE-38813 incoming")
    def test_create_too_long_name(self):
        size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        # Based on the comment above, you might think any name over 238 bytes
        # will be too long; however, maximum paths can be finicky. Instead,
        # this function creates a name that is guaranteed to fill the maximum
        # size of the buffer passed to shm_open(), which is
        # DRAGON_MEMORY_DEFAULT_MAX_FILE_NAM_LENGTH bytes.
        name = "a" * 4096
        with self.assertRaises(DragonPoolCreateFail):
            _ = MemoryPool(size, name, 1, None)

    def test_create_random_name(self):
        name = ''.join(random.choices(string.ascii_letters + string.digits,
                                      k=random.randint(1, 236)))
        size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        mpool = MemoryPool(size, name, 1, None)
        self.assertIsInstance(mpool, MemoryPool)
        mpool.destroy()

    def test_create_random(self):
        name = ''.join(random.choices(string.ascii_letters + string.digits,
                                      k=random.randint(1, 236)))
        size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        uid = random.randint(0, 2 ** 64)
        mpool = MemoryPool(size, name, uid, None)
        self.assertIsInstance(mpool, MemoryPool)
        mpool.destroy()

    def test_create_size_float(self):
        with self.assertRaises(TypeError):
            _ = MemoryPool(34567.987, "mpool_test", DEFAULT_UID, None)

    def test_create_uid_float(self):
        with self.assertRaises(TypeError):
            _ = MemoryPool(MIN_POOL_SIZE, "mpool_test", 1.8, None)

    def test_create_too_few_args(self):
        with self.assertRaises(TypeError):
            _ = MemoryPool(MIN_POOL_SIZE)

    def test_create_too_many_args(self):
        with self.assertRaises(TypeError):
            _ = MemoryPool(MIN_POOL_SIZE, "mpool_test", DEFAULT_UID, None, None)


class MemoryPoolAllocTest(unittest.TestCase):

    def setUp(self):
        self.size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        self.mpool = MemoryPool(self.size, "mpool_test", DEFAULT_UID, None)

    def tearDown(self):
        self.mpool.destroy()

    def test_alloc_free(self):
        mem = self.mpool.alloc(512)
        self.assertEqual(1, self.mpool.get_allocations().num_allocs)
        mem.free()
        self.assertEqual(0, self.mpool.get_allocations().num_allocs)

    def test_clone(self):
        mem = self.mpool.alloc(512)
        mem.get_memview()[100] = 75
        num_allocs = self.mpool.get_allocations().num_allocs
        self.assertEqual(1, num_allocs)
        mem2 = mem.clone(100)
        self.assertEqual(mem2.get_memview()[0], mem.get_memview()[100])
        self.assertEqual(mem.size, mem2.size + 100)
        mem.free()

    def test_alloc_random(self):
        num_allocations = random.randint(25, 100)
        max_size = max(MIN_POOL_SIZE, math.floor(self.size / num_allocations))
        num_curr_allocs = 0
        curr_allocs = []
        for i in range(num_allocations):
            alloc_size = random.randint(256, max_size)
            mem = self.mpool.alloc(alloc_size)
            curr_allocs.append(mem)
            num_curr_allocs += 1
            self.assertEqual(num_curr_allocs, self.mpool.get_allocations().num_allocs)
            num_frees = math.floor(random.random() * num_curr_allocs)
            for j in range(num_frees):
                index = random.randint(0, num_curr_allocs - 1)
                curr_allocs[index].free()
                del curr_allocs[index]
                num_curr_allocs -= 1
                self.assertEqual(num_curr_allocs, self.mpool.get_allocations().num_allocs)
        for alloc in curr_allocs:
            alloc.free()
        self.assertEqual(self.mpool.get_allocations().num_allocs, 0)

    def test_get_id(self):
        mem = self.mpool.alloc(512)
        allocs = self.mpool.get_allocations()
        self.assertIsInstance(allocs, MemoryAllocations)
        self.assertEqual(allocs.alloc_id(0), 1)
        mem.free()

    def test_alloc_string(self):
        with self.assertRaises(TypeError):
            _ = self.mpool.alloc("512")

    def test_alloc_negative(self):
        with self.assertRaises(Exception):
            _ = self.mpool.alloc(-1)

    def test_alloc_float(self):
        with self.assertRaises(TypeError):
            _ = self.mpool.alloc(512.789)

    def test_alloc_zero(self):
        with self.assertRaises(RuntimeError):
            _ = self.mpool.alloc(0)

    def test_alloc_all(self):
        mem = self.mpool.alloc(self.size)
        self.assertEqual(1, self.mpool.get_allocations().num_allocs)
        mem.free()

    def test_alloc_small(self):
        mem = self.mpool.alloc(1)
        self.assertEqual(1, self.mpool.get_allocations().num_allocs)
        mem.free()

    def test_alloc_all_incremental(self):
        # Loop over 1000 times to stress test
        for i in range(0, 1000):
            size = self.size
            allocs = []

            # Random size out of remaining pool
            alloc_size = random.randint(1, size)
            # Set to >= power of 2 to prevent overflow
            alloc_size = 1 << (alloc_size - 1).bit_length()

            # While we have a valid block left, allocate and repeat
            while size >= alloc_size:
                # Reduce remaining size
                size -= alloc_size
                # Allocate
                mem = self.mpool.alloc(alloc_size)
                allocs.append(mem)

                if size == 0:
                    break

                alloc_size = random.randint(1, size)
                alloc_size = 1 << (alloc_size - 1).bit_length()

            # Free allocations to repeat
            for alloc in allocs:
                alloc.free()

    def test_return_codes(self):
        x = DragonError.SUCCESS
        self.assertEqual(x.name, 'SUCCESS')
        self.assertEqual(x.value, 0)

        y = DragonError.INVALID_ARGUMENT
        self.assertEqual(y.name, 'INVALID_ARGUMENT')



class MemoryPoolAllocNoSetupTests(unittest.TestCase):

    def test_alloc_after_destroy(self):
        with self.assertRaises(DragonPoolError):
            self.mpool = MemoryPool(random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE),
                                           "mpool_test", 1, None)
            self.mpool.destroy()
            _ = self.mpool.alloc(512)

    def test_alloc_more_than_size(self):
        mpool = MemoryPool(MIN_POOL_SIZE, "mpool_test", 1, None)
        mem1 = mpool.alloc(10000)  # Powers of 2, allocates 16k
        mem2 = mpool.alloc(10000)  # Powers of 2, allocates 16k
        # Out of memory, expect failure
        with self.assertRaises(DragonMemoryError):
            _ = mpool.alloc(5000)

        mem1.free()
        mem2.free()
        mpool.destroy()


class MemoryPoolAttachTests(unittest.TestCase):

    def setUp(self):
        self.mpool_size = random.randint(MIN_POOL_SIZE, MAX_POOL_SIZE)
        self.mpool = MemoryPool(self.mpool_size, "mpool_test", 1, None)

    def tearDown(self):
        self.mpool.destroy()

    def test_attach(self):
        mem1 = self.mpool.alloc(512)
        mem_ser = mem1.serialize()
        mem2 = MemoryAlloc.attach(mem_ser)
        memview1 = mem1.get_memview()
        memview2 = mem2.get_memview()
        memview2[0:5] = b'Hello'
        self.assertEqual(memview1[0:5], memview2[0:5])
        self.assertEqual(memview1[0:5], b'Hello')
        memview1[0:11] = b'Hello world'
        self.assertEqual(memview1[0:11], memview2[0:11])
        self.assertEqual(memview1[0:11], b'Hello world')
        mem1.free()

    def test_attach_to_number(self):
        with self.assertRaises(TypeError):
            _ = MemoryAlloc.attach(1)

    def test_attach_to_string(self):
        with self.assertRaises(TypeError):
            _ = MemoryAlloc.attach("mpool_test")

    def test_attach_to_mpool(self):
        with self.assertRaises(TypeError):
            pool = MemoryPool(MIN_POOL_SIZE, "temp", 2, None)
            _ = MemoryAlloc.attach(pool)
        pool.destroy()

    def test_attach_allocation(self):
        with self.assertRaises(TypeError):
            mem1 = self.mpool.alloc(512)
            _ = MemoryAlloc.attach(mem1)
        mem1.free()

    def test_attach_bytes(self):
        with self.assertRaises(TypeError):
            mem1 = self.mpool.alloc(512)
            bytes_obj = pickle.dumps(mem1)
            _ = MemoryAlloc.attach(bytes_obj)
        mem1.free()

    def test_attach_many(self):
        mem1 = self.mpool.alloc(512)
        mem_ser = mem1.serialize()
        mem2 = MemoryAlloc.attach(mem_ser)
        mem_ser2 = mem2.serialize()
        mem3 = MemoryAlloc.attach(mem_ser2)
        memview1 = mem1.get_memview()
        memview2 = mem2.get_memview()
        memview3 = mem3.get_memview()
        memview2[0:5] = b'Hello'
        self.assertEqual(memview1[0:5], memview2[0:5])
        self.assertEqual(memview1[0:5], memview3[0:5])
        self.assertEqual(memview1[0:5], b'Hello')
        memview1[0:11] = b'Hello world'
        self.assertEqual(memview1[0:11], memview2[0:11])
        self.assertEqual(memview1[0:11], memview3[0:11])
        self.assertEqual(memview1[0:11], b'Hello world')
        mem1.free()

    def test_attach_detach(self):
        mem1 = self.mpool.alloc(512)
        mem_ser = mem1.serialize()
        mem2 = MemoryAlloc.attach(mem_ser)
        memview1 = mem1.get_memview()
        memview2 = mem2.get_memview()
        memview2[0:5] = b'Hello'
        self.assertEqual(memview1[0:5], memview2[0:5])
        self.assertEqual(memview1[0:5], b'Hello')

        # Assert that after detaching, we cannot get a reference to the underlying memory
        with self.assertRaises(DragonMemoryError):
            mem2.detach()
            _ = mem2.get_memview()

        mem1.free()

    def test_remote(self):
        mem1 = self.mpool.alloc(512)
        mem_ser = mem1.serialize()
        mem_ser_remote = mk_remote_mem_ser(mem_ser)
        mem2 = MemoryAlloc.attach(mem_ser_remote)
        mem_ser2 = mem2.serialize()
        mem3 = MemoryAlloc.attach(mem_ser2)
        mem4 =  mem3.clone(100)
        mem4_ser = mem4.serialize()
        mem5 = MemoryAlloc.attach(mem4_ser)

        rmt_pool_ser = mk_remote_pool_ser(self.mpool.serialize())
        rmt_pool = MemoryPool.attach(rmt_pool_ser)
        rmt_pool.detach()

        rmt_pool = MemoryPool.attach(rmt_pool_ser)
        with self.assertRaises(DragonPoolError):
            rmt_pool.destroy()

        with self.assertRaises(DragonPoolError):
            mem6 = rmt_pool.alloc(100)

        with self.assertRaises(DragonPoolError):
            mem6 = rmt_pool.alloc_blocking(100)

        with self.assertRaises(DragonPoolError):
            allocs = rmt_pool.get_allocations()

        with self.assertRaises(DragonPoolError):
            rmt_pool.allocation_exists(AllocType.DATA,20)

        rmt_pool.detach()

        with self.assertRaises(DragonMemoryError):
            mem5.get_memview()

        with self.assertRaises(DragonMemoryError):
            mem5.free()

        mem1.free()

if __name__ == '__main__':
    mp.set_start_method('spawn', force=True)
    unittest.main(verbosity=2)

test/utils/test_basic_mempool.py
DragonHPC/dragon
#!/usr/bin/env python3

import time
import unittest
import mmap
import multiprocessing as mp
from multiprocessing import shared_memory as shm
from multiprocessing import Array
from parameterized import parameterized
from parameterized import parameterized_class
from dragon.locks import DragonLock, Type
from enum import Enum


class Mem(Enum):
    MMAP = 0
    SHM = 1


# Driver class that contains all shared tests for FIFO vs Greedy and SharedMem vs MMAP
@parameterized_class([
    {"lock_type": Type.FIFO, "mem_type": Mem.MMAP},
    {"lock_type": Type.FIFO, "mem_type": Mem.SHM},
    {"lock_type": Type.FIFOLITE, "mem_type": Mem.MMAP},
    {"lock_type": Type.FIFOLITE, "mem_type": Mem.SHM},
    {"lock_type": Type.GREEDY, "mem_type": Mem.MMAP},
    {"lock_type": Type.GREEDY, "mem_type": Mem.SHM},
])
class LockTest(unittest.TestCase):

    def setUp(self) -> None:
        self.memsize = DragonLock.size(self.lock_type)
        self.memsize = int(4096.0 * ((self.memsize / 4096.0) + 1.0))
        if self.mem_type == Mem.MMAP:
            self.memobj = mmap.mmap(-1, self.memsize)
        elif self.mem_type == Mem.SHM:
            self.shmemobj = shm.SharedMemory(create=True, size=self.memsize)
            self.memobj = self.shmemobj.buf
        else:
            raise RuntimeError("Invalid memory type")

        self.hdl = None

    def tearDown(self) -> None:
        if self.hdl:
            # If locked, nothing happens
            # If unlocked, guarantees we can do unlock and avoid undefined behavior
            self.hdl.try_lock()
            self.hdl.unlock()
            self.hdl.destroy()
            del self.hdl

        if self.mem_type == Mem.MMAP:
            self.memobj.close()
            del self.memobj

        elif self.mem_type == Mem.SHM:
            self.shmemobj.close()
            self.shmemobj.unlink()
            del self.shmemobj

    # Count atomically, make sure the locks work as expected
    @staticmethod
    def atomic_count(lockmem, iters, tdata) -> None:
        hdl = DragonLock.attach(lockmem)

        for i in range(iters):
            hdl.lock()
            tdata[0] += 1
            hdl.unlock()

        hdl.detach()

    # Used to benchmark lock/unlock performance across multiple processes
    def bench_ops(self, lockmem, warmup, iters, tdata) -> None:
        hdl = DragonLock.attach(lockmem)

        for i in range(warmup + iters):
            hdl.lock()

            if i == warmup:
                start_time = time.monotonic_ns()

            hdl.unlock()

        end_time = time.monotonic_ns()
        total = 1.0 * (end_time - start_time) * 1e-9
        ops = (1.0 * iters) / total

        # Multiprocessing Arrays are thread safe according to docs
        tdata[0] += total
        tdata[1] += ops

        hdl.detach()

    def test_init_single(self) -> None:
        self.hdl = DragonLock.init(self.lock_type, self.memobj)
        self.hdl.lock()
        self.hdl.unlock()

    @parameterized.expand([
        ["2Thread", 2],
        ["4Thread", 4]
    ])
    def test_lock_count(self, name: str, maxprocs: int) -> None:
        self.hdl = DragonLock.init(self.lock_type, self.memobj)
        self.hdl.lock()

        iters = 1000
        procs = []

        tdata = Array('i', range(1))

        for i in range(maxprocs):
            proc = mp.Process(target=self.atomic_count,
                              args=(self.memobj, iters, tdata))
            proc.start()
            procs.append(proc)

        self.hdl.unlock()

        for proc in procs:
            proc.join()

        self.assertEqual(maxprocs * iters, tdata[0])

    @parameterized.expand([
        ["1Thread", 1],
        ["2Thread", 2],
        ["4Thread", 4]
    ])
    def test_multi_lock(self, name: str, maxprocs: int) -> None:
        self.hdl = DragonLock.init(self.lock_type, self.memobj)
        self.hdl.lock()

        warmup = 1000
        iters = 10000
        tdata = Array('d', range(2))
        procs = []

        for i in range(maxprocs):
            proc = mp.Process(target=self.bench_ops,
                              args=(self.memobj, warmup, iters, tdata))
            proc.start()
            procs.append(proc)

        self.hdl.unlock()

        for proc in procs:
            proc.join()

        avg_time = tdata[0] / maxprocs
        avg_ops = tdata[1] / maxprocs
        # print(f"Average time/ops: {avg_time} ({avg_ops} ops/s)")

    # TODO: This currently fails on FIFO locks
    def test_try_lock(self):
        self.hdl = DragonLock.init(self.lock_type, self.memobj)
        self.hdl.lock()
        expected = 0
        res = self.hdl.try_lock()
        self.assertEqual(expected, res)
        self.hdl.unlock()
        expected = 1
        res = self.hdl.try_lock()
        self.assertEqual(expected, res)
        self.hdl.unlock()

    def worker_try_lock(self, lockmem, procid):
        hdl = DragonLock.attach(lockmem)
        res = hdl.try_lock()
        hdl.detach()
        if res != 0:
            raise RuntimeError(f"Process {procid} got lock when it should not have")

    @parameterized.expand([
        ["1Thread", 1],
        ["2Thread", 2],
        ["4Thread", 4]
    ])
    def test_multi_try_lock(self, name: str, maxprocs: int) -> None:
        self.hdl = DragonLock.init(self.lock_type, self.memobj)
        self.hdl.lock()

        procs = []
        for i in range(maxprocs):
            proc = mp.Process(target=self.worker_try_lock, args=(self.memobj, i))
            proc.start()
            procs.append(proc)

        time.sleep(1)

        for proc in procs:
            proc.join()

        self.hdl.unlock()

    @unittest.skip("Skipping until double unlock behavior is well defined")
    def test_fail_double_unlock(self):
        self.hdl = DragonLock.init(self.lock_type, self.memobj)
        self.hdl.lock()
        self.hdl.unlock()
        # Should throw an exception on attempting to double unlock
        self.assertRaises(self.hdl.unlock())
        # Delete the handle object so tearDown() doesn't cause runtime errors
        del self.hdl


if __name__ == '__main__':
    unittest.main()

test/utils/test_locks.py
DragonHPC/dragon
#!/usr/bin/env python3
import os
import sys
import time
import unittest
import random
import string
import subprocess
import logging
import queue
from contextlib import redirect_stderr
from unittest.mock import MagicMock, patch
from io import StringIO

import dragon.infrastructure.messages as dmsg
import dragon.infrastructure.parameters as dparm
import dragon.infrastructure.facts as dfacts
import dragon.dlogging.util as dlog
import dragon.launcher.util as dlutil
import dragon.infrastructure.util as dutil
import dragon.utils as du

from dragon.dlogging.logger import DragonLogger, DragonLoggingError
from dragon.channels import ChannelEmpty
from dragon.utils import B64


def generate_msg(length):
    return ''.join(random.choice(string.ascii_letters + string.digits + ' ') for i in range(length))


def run_logger(info_msg, debug_msg):

    # set up my stdin for listening to my channel descriptor
    logger_sdesc = B64.from_str(os.environ.get(dfacts.DRAGON_LOGGER_SDESC))

    # Log stuff
    level, full_fname = dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST,
                                              logger_sdesc=logger_sdesc)

    if level <= logging.DEBUG:
        l_stdout = dutil.NewlineStreamWrapper(sys.stdout, read_intent=False)
        l_stdout.send(full_fname)

    log = logging.getLogger(dlog.DragonLoggingServices.TEST)
    log.info(info_msg)
    log.debug(debug_msg)


def run_logging_subprocesses(nproc, level, sdesc):

    # Update environment
    the_env = os.environ
    the_env['DRAGON_LOGGER_SDESC'] = du.B64.bytes_to_str(sdesc)

    info_msgs = []
    debug_msgs = []
    log_procs = []
    log_stdins = []
    fnames = []
    for i in range(nproc):
        info_msgs.append(generate_msg(random.randrange(20, 100)))
        debug_msgs.append(generate_msg(random.randrange(20, 100)))
        run_string = f"from utils.test_logging import run_logger;  run_logger('{info_msgs[-1]}', '{debug_msgs[-1]}')"
        log_procs.append(subprocess.Popen(['python3', '-c', run_string], env=the_env, stdout=subprocess.PIPE,
                                          stdin=subprocess.PIPE))
        log_stdins.append(dutil.NewlineStreamWrapper(log_procs[-1].stdout, write_intent=False))

    if level <= logging.DEBUG:
        fnames = [log_stdin.recv().rstrip() for log_stdin in log_stdins]
    else:
        fnames = None

    for p in log_procs:
        p.stdin.close()
        p.stdout.close()
        p.wait()

    return info_msgs, debug_msgs, fnames


def close_handlers(log):

    for handler in log.handlers:
        log.removeHandler(handler)
        handler.close()


class LoggingTest(unittest.TestCase):

    def setUp(self):
        self.mpool = dlog._get_logging_mpool(0)
        self.logger = None

    def tearDown(self):
        if self.logger is None:
            self.mpool.destroy()
        else:
            self.logger.destroy()  # This also nukes the underlying pool, change that behavior?

    def test_create(self):
        self.logger = DragonLogger(self.mpool)

    def test_put(self):
        self.logger = DragonLogger(self.mpool)
        self.logger.put("A")

    def test_put_get(self):
        self.logger = DragonLogger(self.mpool)
        self.logger.put('aaaa')
        a = self.logger.get()
        self.assertEqual(a, 'aaaa')

    def test_put_get_many(self):
        self.logger = DragonLogger(self.mpool)
        abcs = "abcdefghijklmnopqrstuvwxyz"
        for i in range(0, len(abcs)):
            self.logger.put(abcs[i])

        for i in range(0, len(abcs)):
            msg = self.logger.get()
            self.assertEqual(msg, abcs[i])

    def test_get_timeout(self):
        # What else would be good test cases for timeout besides "it times out"?
        self.logger = DragonLogger(self.mpool)
        with self.assertRaises(ChannelEmpty):
            self.logger.get(logging.WARNING, timeout=1)

    def test_put_get_priority(self):
        self.logger = DragonLogger(self.mpool)
        self.logger.put("debug", logging.DEBUG)
        self.logger.put("info", logging.INFO)
        self.logger.put("warning", logging.WARNING)
        self.logger.put("error", logging.ERROR)

        # Only get messages at warning or above
        msg_debug = self.logger.get(logging.WARNING)
        msg_info = self.logger.get(logging.WARNING)
        msg_warning = self.logger.get(logging.WARNING)
        msg_error = self.logger.get(logging.WARNING)

        self.assertEqual(msg_debug, None)
        self.assertEqual(msg_info, None)
        self.assertEqual(msg_warning, "warning")
        self.assertEqual(msg_error, "error")


class TestLogHandler(unittest.TestCase):

    def setUp(self) -> None:
        self.info_log_message = self.get_random_string(72)
        self.debug_log_message = self.get_random_string(72)

    def get_random_string(self, length: int) -> str:
        # choose from all lowercase letter
        return ''.join(random.choice(string.ascii_lowercase) for i in range(length))

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "INFO"})
    @patch('dragon.dlogging.util.DragonLogger', autospec=True)
    def test_dragon_loghandler_mock(self, mock_dragon_logger) -> None:

        instance_mock = MagicMock()
        mock_dragon_logger.attach.return_value = instance_mock
        serialized_log_descr = B64.from_str(self.get_random_string(1024))

        level, _ = dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST,
                                         logger_sdesc=serialized_log_descr)
        py_logger = logging.getLogger(dlog.DragonLoggingServices.TEST)
        py_logger.info(self.info_log_message)
        py_logger.debug(self.debug_log_message)

        mock_dragon_logger.attach.assert_called_with(serialized_log_descr.decode(), mpool=None)
        self.assertEqual(instance_mock.put.call_count, 1)

        serial_msg, level = instance_mock.mock_calls[-1].args
        msg = dmsg.parse(serial_msg)
        self.assertEqual(msg.msg, self.info_log_message)
        self.assertEqual(level, logging.INFO)

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "INFO"})
    def test_dragon_loghandler_to_dragon_logging_layer(self) -> None:
        from logging import INFO

        dragon_logger = dlog.setup_dragon_logging(0)
        dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST,
                              logger_sdesc=B64(dragon_logger.serialize()))
        py_logger = logging.getLogger(dlog.DragonLoggingServices.TEST)
        py_logger.info(self.info_log_message)
        msg = dmsg.parse(dragon_logger.get(INFO))
        self.assertEqual(msg.msg, self.info_log_message)

        dragon_logger.destroy()  # This also nukes the underlying pool


class TestLoggingSubprocesses(unittest.TestCase):

    def setUp(self):

        self.dragon_logger = None

    def create_logger(self):

        # Create the logger children will attach to
        self.dragon_logger = dlog.setup_dragon_logging(0)

    def destroy_logger(self):

        self.dragon_logger.destroy()
        self.dragon_logger = None

    def get_messages(self, nmsg, level):

        total = 0
        iter_max = 500
        niter = 0  # set a break point
        msgs = []
        while (total != nmsg or niter > iter_max):
            try:
                msgs.append(dmsg.parse(self.dragon_logger.get(level)))
                total += 1
            except (DragonLoggingError, ChannelEmpty):
                pass
            time.sleep(0.01)
            niter += 1

        # Make sure we didn't time out
        self.assertLessEqual(niter, iter_max)
        return msgs

    def run_subprocesses(self, nproc, level):

        # Create logger
        self.create_logger()

        # Run logging subprocesses
        info_msgs, debug_msgs, fnames = run_logging_subprocesses(nproc, level, self.dragon_logger.serialize())

        if level is logging.DEBUG:
            self.nmsgs = 2 * nproc
        else:
            self.nmsgs = nproc

        log_msgs = self.get_messages(self.nmsgs, level)

        # Destroy my logger
        self.destroy_logger()

        return info_msgs, debug_msgs, log_msgs

    def check_msgs(self, info_msgs, debug_msgs, log_msgs, level):

        # Get the actual log message out of the object
        log_msgs = [msg.msg for msg in log_msgs]
        self.assertEqual(len(log_msgs), self.nmsgs)

        if level is logging.DEBUG:
            input_msgs = debug_msgs + info_msgs
        else:
            input_msgs = info_msgs

        for msg in input_msgs:
            self.assertTrue(msg in log_msgs)
            self.assertEqual(log_msgs.count(msg), 1)
            log_msgs.remove(msg)

        self.assertEqual(len(log_msgs), 0)

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "INFO"})
    def test_one_to_one_logging_info(self):
        nproc = 1
        level = logging.INFO
        info_msgs, debug_msgs, log_msgs = self.run_subprocesses(nproc, level)
        self.check_msgs(info_msgs, debug_msgs, log_msgs, level)

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "DEBUG"})
    def test_one_to_one_logging_debug(self):
        nproc = 1
        level = logging.DEBUG
        info_msgs, debug_msgs, log_msgs = self.run_subprocesses(nproc, level)
        self.check_msgs(info_msgs, debug_msgs, log_msgs, level)

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "INFO"})
    def test_many_to_one_logging_info(self):
        nproc = 2
        level = logging.INFO
        info_msgs, debug_msgs, log_msgs = self.run_subprocesses(nproc, level)
        self.check_msgs(info_msgs, debug_msgs, log_msgs, level)

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "DEBUG"})
    def test_many_to_one_logging_debug(self):
        nproc = 2
        level = logging.DEBUG
        info_msgs, debug_msgs, log_msgs = self.run_subprocesses(nproc, level)
        self.check_msgs(info_msgs, debug_msgs, log_msgs, level)


class TestLoggingInfrastructure(unittest.TestCase):

    def setUp(self):
        # Create a dragon logger
        self.dragon_logger = dlog.setup_dragon_logging(0)
        self.destroyed_dragon = False
        self.log_levels = [logging.CRITICAL, logging.ERROR, logging.WARNING, logging.INFO,
                           logging.DEBUG, logging.NOTSET]
        self.dummy_msgs = [dmsg.GSIsUp, dmsg.GSTeardown, dmsg.SHTeardown]
        self._tag = 0

    def tearDown(self):
        # Destroy dragon logger
        if not self.destroyed_dragon:
            self.dragon_logger.destroy()
            self.destroyed_dragon = True

        # Make sure the root logger has closed any handlers since the atexit()
        # function isn't going to get called by logging
        for handler in logging.root.handlers:
            logging.root.removeHandler(handler)
            handler.close()
        logging.shutdown()

    def get_tag(self):
        tag = self._tag
        self._tag += 1
        return tag

    def write_logs(self, log, n_entries):

        messages = [(random.choice(self.log_levels), generate_msg(random.randrange(20, 100)))
                    for _ in range(n_entries)]
        for msg in messages:
            log.log(*msg)

        return messages

    def flush_logs(self, level, test_queue):

        for _ in range(self.dragon_logger.num_logs()):
            msg = self.dragon_logger.get(level)
            if msg is not None:
                test_queue.send(dmsg.parse(msg).serialize())

    def send_dummy_msgs(self, n_msgs, test_queue):

        for _ in range(n_msgs):
            msg_class = random.choice(self.dummy_msgs)
            test_queue.send(msg_class(tag=self.get_tag()).serialize())

    def check_log_entries(self, input_msgs, level, log_queue):

        # Grab messages from the logging queue. This will return a list of records
        logged_msgs = []
        while not log_queue.empty():
            logged_msgs.append(log_queue.get())
        logged_msgs = [(msg.level, msg.msg) for msg in logged_msgs]

        # First remove input_msgs that are < level from our input messages since
        # those wouldn't have been logged
        actually_logged = [(msg_level, msg) for msg_level, msg in input_msgs if msg_level >= level]

        self.assertEqual(len(logged_msgs), len(actually_logged))
        for msg in actually_logged:
            self.assertTrue(msg in logged_msgs)
            self.assertEqual(logged_msgs.count(msg), 1)
            logged_msgs.remove(msg)
        self.assertEqual(len(logged_msgs), 0)

    def clear_test_queue(self, test_queue, log_queue):
        # Get log entries by clearing the queue. Decorator should
        # pull off log messages
        unwrapped_get = dlutil.get_with_timeout.__wrapped__

        def get_with_timeout(test_queue, timeout=0.01):
            return unwrapped_get(test_queue, timeout=timeout)

        wrapped_get = dlutil.logging_queue_monitor(get_with_timeout, log_test_queue=log_queue)
        while not test_queue.empty():
            try:
                wrapped_get(test_queue, timeout=0.01)
            except TimeoutError:
                pass

    def test_logging_FE_queue_simple(self):
        # Set up basic logging for a given service
        level, _ = dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST,
                                         logger_sdesc=B64(self.dragon_logger.serialize()))
        log = logging.getLogger(dlog.DragonLoggingServices.TEST)
        test_queue = dlutil.SRQueue()
        log_queue = queue.SimpleQueue()

        # write a slew of messages to the queue
        group_1 = self.write_logs(log, 5)
        self.flush_logs(level, test_queue)
        self.send_dummy_msgs(5, test_queue)
        group_2 = self.write_logs(log, 4)
        self.flush_logs(level, test_queue)
        self.send_dummy_msgs(3, test_queue)

        # Check the logging queue got the messages correctly
        self.clear_test_queue(test_queue, log_queue)
        self.check_log_entries(group_1 + group_2, level, log_queue)

    def test_log_to_wrong_parent_log(self):
        level, _ = dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST,
                                         logger_sdesc=B64(self.dragon_logger.serialize()))

        fmter = logging.Formatter(dlog.default_services_fmt)
        not_a_service = dlog.DragonLoggingServices.TEST
        log = logging.getLogger(not_a_service)
        test_queue = dlutil.SRQueue()
        log_queue = queue.SimpleQueue()

        self.write_logs(log, 5)
        self.flush_logs(level, test_queue)
        self.clear_test_queue(test_queue, log_queue)

        logged_msgs = []
        while not log_queue.empty():
            logged_msgs.append(log_queue.get())
        logged_msgs = [fmter.format(logging.makeLogRecord(msg.get_sdict())) for msg in logged_msgs]

        # Make sure formatted entries have service and name info:
        for msg in logged_msgs:
            self.assertIn(dlog.DragonLoggingServices.TEST, msg)
            self.assertIn(not_a_service, msg)

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "DEBUG"})
    def test_debug_level_logging(self):
        level, _ = dlog.setup_FE_logging(basedir=os.getcwd())
        log = logging.getLogger(dlog.DragonLoggingServices.LA_FE)
        log.info('FE logger up')

        nproc = 1
        info_msgs, debug_msgs, fnames = run_logging_subprocesses(nproc, level, self.dragon_logger.serialize())

        written_msgs = info_msgs + debug_msgs
        log.info("fe done")

        # Check file exists
        for fname in fnames:
            self.assertTrue(os.path.exists(fname))

        # Open file, parse logs and make sure all the messages are in there
        for file in fnames:
            with open(file, 'r') as f:
                log_msgs = f.readlines()
            os.remove(file)

        self.assertEqual(len(written_msgs), len(log_msgs))

        for msg in written_msgs:
            match = [m for m in log_msgs if msg in m]
            self.assertEqual(len(match), 1)
            log_msgs.remove(match[-1])

        self.assertEqual(len(log_msgs), 0)

        # Make sure all handlers associated with FE logging
        # are closed beofre starting next tests
        dlog.close_FE_logging()

    @patch.dict(os.environ, {"DRAGON_LOG_DEVICE_STDERR": "DEBUG", "DRAGON_LA_LOG_DIR": "/not/a/real/dir"})
    def test_undefined_dragon_la_log_dir(self):
        level, fn = dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST,
                                          logger_sdesc=B64(self.dragon_logger.serialize()))
        log = logging.getLogger(dlog.DragonLoggingServices.TEST)
        info_msg = [generate_msg(random.randrange(20, 100))]
        debug_msg = [generate_msg(random.randrange(20, 100))]
        log.info(info_msg[-1])
        log.debug(debug_msg[-1])
        written_msgs = info_msg + debug_msg

        # Open file, parse logs and make sure all the messages are in there
        with open(fn, 'r') as f:
            log_msgs = f.readlines()
        os.remove(fn)

        self.assertEqual(len(written_msgs), len(log_msgs))
        for msg in written_msgs:
            match = [m for m in log_msgs if msg in m]
            self.assertEqual(len(match), 1)
            log_msgs.remove(match[-1])
        self.assertEqual(len(log_msgs), 0)

    def test_logging_to_terminal(self):
        with redirect_stderr(StringIO()) as f:
            level, fname = dlog.setup_FE_logging(basedir=os.getcwd(), add_console=True)
            log = logging.getLogger(dlog.DragonLoggingServices.LA_FE)
            msgs = ['FE logger up']
            log.info(msgs[-1])
            msgs.append('test warning')
            log.warning(msgs[-1])

        # get data off of stderr and check for accuracy
        stderr_msgs = f.getvalue().rstrip().split('\n')
        for msg in msgs:
            match = [m for m in stderr_msgs if msg in m]
            self.assertEqual(len(match), 1)
            stderr_msgs.remove(match[-1])
        self.assertEqual(len(stderr_msgs), 0)

        os.remove(fname)
        dlog.close_FE_logging()

    def test_write_to_destroyed_dragon_logger(self):
        # Set up logging for BE services
        level, fn = dlog.setup_BE_logging(service=dlog.DragonLoggingServices.TEST,
                                          logger_sdesc=B64(self.dragon_logger.serialize()))
        log = logging.getLogger(dlog.DragonLoggingServices.TEST)

        # Destroy dragon logger
        self.dragon_logger.destroy()
        self.destroyed_dragon = True
        self.assertRaises(DragonLoggingError, log.info, generate_msg(random.randrange(20, 100)))


if __name__ == '__main__':
    unittest.main(verbosity=2)

test/utils/test_logging.py
DragonHPC/dragon
#!/usr/bin/env python3

import time
import unittest
import multiprocessing as mp
import time
from dragon.managed_memory import MemoryPool, MemoryAlloc, DragonPoolError, DragonPoolAllocationNotAvailable
import os

SMALL_TIMEOUT_VAL = 0.2
BIGGER_TIMEOUT_VAL = 1.2

def worker_attach(q=None, pool_ser=None, mem_ser=None):
    if pool_ser is not None and q is not None:
        mpool = MemoryPool.attach(pool_ser)
        mem = mpool.alloc(512)
        memview = mem.get_memview()
        memview[0:5] = b"Hello"
        q.put(mem.serialize())
        mpool.detach()
    elif mem_ser is not None:
        mem = MemoryAlloc.attach(mem_ser)
        memview = mem.get_memview()
        memview[0:1] = b"Y"
    else:
        raise RuntimeError("Invalid arguments to worker_attach")

def worker_blocking_alloc(q, pool_ser, pid):
    mpool = MemoryPool.attach(pool_ser)
    mem = mpool.alloc_blocking(3000*pid+100)
    memview = mem.get_memview()
    memview[0:5] = b"Hello"
    q.put(mem.serialize())
    mpool.detach()

def worker_blocking_alloc_timeout(pool_ser, TIMEOUT_VAL):
    mpool = MemoryPool.attach(pool_ser)
    try:
        mem = mpool.alloc_blocking(3100,timeout=TIMEOUT_VAL)
        rc = 1
    except TimeoutError as ex:
        rc = 0
    except DragonPoolError:
        rc = 2
    finally:
        mpool.detach()

    exit(rc)

def worker_attach_via_pickle(q, mpool):
    mem = mpool.alloc(512)
    memview = mem.get_memview()
    memview[:5] = b"Howdy"
    q.put(mem.serialize())
    mpool.detach()

class MemPoolTest(unittest.TestCase):
    mpool = None

    @classmethod
    def setUpClass(cls):
        # Create a 32kb memory pool ahead of time for all tests
        cls.mpool = MemoryPool(32768, "mpool_test", 1, None)

    @classmethod
    def tearDownClass(cls):
        cls.mpool.destroy()

    def test_exhaust_pool(self):
        allocations = []
        try:
            while True:
                mem = self.mpool.alloc(1)
                allocations.append(mem)
        except DragonPoolAllocationNotAvailable as ex:
            self.assertTrue("DRAGON_DYNHEAP_REQUESTED_SIZE_NOT_AVAILABLE" in str(ex))

        for mem in allocations:
            mem.free()

    def test_memory_serialize_attach(self):
        mem1 = self.mpool.alloc(512)
        mem_ser = mem1.serialize()

        mem2 = MemoryAlloc.attach(mem_ser)

        memview1 = mem1.get_memview()
        memview2 = mem2.get_memview()

        memview1[0:5] = b"Hello"
        self.assertEqual(memview1[0:5], memview2[0:5])

        memview2[0:1] = b"Y"
        self.assertEqual(memview1[0:5], memview2[0:5])
        mem1.free()

    def test_check_allocations(self):
        mem1 = self.mpool.alloc(512)
        mem2 = self.mpool.alloc(512)

        allocs = self.mpool.get_allocations()
        self.assertEqual(2, allocs.num_allocs)

        for i in range(allocs.num_allocs):
            self.assertEqual(0, allocs.alloc_type(i).value)
            self.assertTrue(self.mpool.allocation_exists(allocs.alloc_type(i), allocs.alloc_id(i)))

        mem1.free()
        mem2.free()

    def test_multiprocess_pool(self):
        pool_ser = self.mpool.serialize()

        q = mp.Queue()
        proc = mp.Process(target=worker_attach, args=(q, pool_ser, None))
        proc.start()
        proc.join()

        mem_ser = q.get()
        mem = MemoryAlloc.attach(mem_ser)
        memview = mem.get_memview()
        self.assertEqual(b"Hello", memview[0:5])
        mem.free()

    def test_multiprocess_pool_via_pickle(self):
        q = mp.Queue()
        proc = mp.Process(target=worker_attach_via_pickle, args=(q, self.mpool))
        proc.start()
        proc.join()

        mem_ser = q.get()
        mem = MemoryAlloc.attach(mem_ser)
        memview = mem.get_memview()
        self.assertEqual(b"Howdy", memview[:5])
        mem.free()

    def test_multiprocess_alloc(self):
        msg = b"Hello"
        mem = self.mpool.alloc(512)
        memview = mem.get_memview()
        memview[0:5] = msg
        mem_ser = mem.serialize()

        proc = mp.Process(target=worker_attach, args=(None, None, mem_ser))
        proc.start()
        proc.join()
        self.assertEqual(b"Yello", memview[0:5])
        mem.free()

    def test_multiprocess_blocking_alloc(self):
        pool_ser = self.mpool.serialize()
        # Getting a big allocation so subsequent allocs will block.
        mem = self.mpool.alloc(32000)
        proc_list = []
        q = mp.Queue()
        for k in range(10):
            proc = mp.Process(target=worker_blocking_alloc, args=(q, pool_ser, k))
            proc_list.append(proc)
            proc.start()

        time.sleep(1)
        mem.free()

        for k in range(10):
            mem_ser = q.get()
            mem = MemoryAlloc.attach(mem_ser)
            memview = mem.get_memview()
            self.assertEqual(b"Hello", memview[0:5])
            mem.free()

        for k in range(10):
            proc_list[k].join()

    def test_small_blocking_timeout(self):
        pool_ser = self.mpool.serialize()
        # Getting a big allocation so subsequent allocs will block.
        mem = self.mpool.alloc(32000)

        proc = mp.Process(target=worker_blocking_alloc_timeout, args=(pool_ser,SMALL_TIMEOUT_VAL))
        start = time.monotonic()
        proc.start()
        proc.join()
        delta = time.monotonic() - start

        rc = proc.exitcode

        mem.free()

        self.assertEqual(rc, 0, "Blocking Alloc process did not timeout.")

        self.assertGreater(delta, SMALL_TIMEOUT_VAL, 'time lapsed should be bigger than timeout.')
        self.assertLessEqual(delta, 1+SMALL_TIMEOUT_VAL, "blocking allocation timeout took too long.")

    def test_bigger_blocking_timeout(self):
        pool_ser = self.mpool.serialize()
        # Getting a big allocation so subsequent allocs will block.
        mem = self.mpool.alloc(32000)

        proc = mp.Process(target=worker_blocking_alloc_timeout, args=(pool_ser,BIGGER_TIMEOUT_VAL))
        start = time.monotonic()
        proc.start()
        proc.join()
        delta = time.monotonic() - start

        rc = proc.exitcode

        mem.free()

        self.assertEqual(rc, 0, "Blocking Alloc process did not timeout.")

        self.assertGreater(delta, BIGGER_TIMEOUT_VAL, 'time lapsed should be bigger than timeout.')
        self.assertLessEqual(delta, 1+BIGGER_TIMEOUT_VAL, "blocking allocation timeout took too long.")

    def test_zero_blocking_timeout(self):
        pool_ser = self.mpool.serialize()
        # Getting a big allocation so subsequent allocs will block.
        mem = self.mpool.alloc(32000)

        proc = mp.Process(target=worker_blocking_alloc_timeout, args=(pool_ser,0))
        start = time.monotonic()
        proc.start()
        proc.join()
        delta = time.monotonic() - start

        rc = proc.exitcode

        mem.free()

        self.assertEqual(rc, 2, "Blocking Alloc process timed out with zero timeout.")

        self.assertGreater(delta, 0, 'There should have been a near immediate return.')
        self.assertLessEqual(delta, 1, "blocking allocation timeout took too long.")


    def test_big_allocation(self):
        mem_bytes = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES')
        big_size = int(mem_bytes * 0.55)
        big_pool = MemoryPool(big_size, "big_pool_test", 2)
        mem = big_pool.alloc(big_size)
        memview = mem.get_memview()
        memview[big_size-1] = 42
        self.assertEqual(memview[big_size-1], 42)
        mem.free()
        big_pool.destroy()

if __name__ == '__main__':
    mp.set_start_method('spawn', force=True)
    unittest.main(verbosity=2)

test/utils/test_mempool.py
DragonHPC/dragon
#/usr/bin/env python3

import unittest
import array
from dragon.pheap import PriorityHeap


class PHeapTest(unittest.TestCase):

    def setUp(self):
        self.cap = 15
        self.nvals = 1
        self.base = 3
        self.size = PriorityHeap.size(self.cap, self.nvals)
        # @MCB TODO: Parameterize this to test with multiprocessing
        self.mem = bytearray(self.size)

    def tearDown(self):
        del self.mem

    def test_init(self):
        """
        Make sure init works without throwing an exception
        """
        hdl = PriorityHeap.create(self.base, self.cap, self.nvals, self.mem)
        hdl.destroy()

    def test_simple_insert(self):
        hdl = PriorityHeap.create(self.base, self.cap, self.nvals, self.mem)

        arr = array.array('Q', [0] * self.cap)
        for i in range(self.cap):
            for j in range(self.nvals):
                arr[j] = i * self.nvals + j

            hdl.insert(arr)

        for i in range(self.cap):
            v = array.array('Q', [0] * self.nvals)
            p = int()
            hdl.extract(v, p)

            # Since nothing is inserted as urgent, everything should have equal value and priority
            for j in range(self.nvals):
                self.assertEqual(i * self.nvals + j, v[j])

        hdl.destroy()

    def test_urgent_insert(self):
        hdl = PriorityHeap.create(self.base, self.cap, self.nvals, self.mem)

        expected = [5, 0, 2, 3, 4, 1, 6, 7, 8, 9, 10, 11, 12, 13, 14]
        arr = array.array('Q', [0] * self.cap)
        for i in range(self.cap):
            for j in range(self.nvals):
                arr[j] = i * self.nvals + j

            if i != 5:
                hdl.insert(arr)
            else:
                hdl.insert(arr, urgent=True)

        for i in range(self.cap):
            v = array.array('Q', [0] * self.nvals)
            p = int()
            hdl.extract(v, p)

            if i == 0:
                for j in range(self.nvals):
                    self.assertEqual(5 * self.nvals + j, v[j])
            elif i <= 5:
                for j in range(self.nvals):
                    self.assertEqual((i-1) * self.nvals + j, v[j])
            else:
                for j in range(self.nvals):
                    self.assertEqual(i * self.nvals + j , v[j])

        hdl.destroy()


if __name__ == '__main__':
    unittest.main()

test/utils/test_pyheap.py
DragonHPC/dragon