<commit_before>from celery import task
from celery.result import AsyncResult

from .models import TeamSubmission

import logging

logger = logging.getLogger(__name__)

@task()
def create_shellai_tag(instance):
    """Tags the repo's HEAD as "ShellAI" to provide a default tag for
    the arena to use"""
    team_name = instance.team.name

    if instance.repository.task_id is not None:
        # Wait for the repo to be created
        AsyncResult(instance.repository.task_id).wait()
        msg = "Waiting for {}'s repository to be created..."
        logger.info(msg.format(team_name))

    logger.info("{}'s repository is ready".format(team_name))

    # Create a submission for the HEAD commit
    TeamSubmission.objects.create(team=instance.team,
                                  commit=instance.repository.repo['HEAD'].id,
                                  name="ShellAI",
                                  submitter=None)
    logger.info("Tagged {}'s repo".format(team_name))
<commit_msg>Handle attempts to tag empty shell repos<commit_after>from celery import task
from celery.result import AsyncResult

from .models import TeamSubmission

import logging

logger = logging.getLogger(__name__)

@task()
def create_shellai_tag(instance):
    """Tags the repo's HEAD as "ShellAI" to provide a default tag for
    the arena to use"""
    team_name = instance.team.name

    if instance.repository.task_id is not None:
        # Wait for the repo to be created
        AsyncResult(instance.repository.task_id).wait()
        msg = "Waiting for {}'s repository to be created..."
        logger.info(msg.format(team_name))

    logger.info("{}'s repository is ready".format(team_name))

    try:
        commit = instance.repository.repo['HEAD']
    except KeyError:
        # Log an error if we can't get a commit
        msg = "Unable to tag {}'s repo. Bad ref 'HEAD'. Is the repo empty?"
        logger.error(msg.format(team_name))
    else:
        # Create a submission for the HEAD commit
        TeamSubmission.objects.create(teamclient=instance,
                                      commit=commit.id,
                                      name="ShellAI",
                                      submitter=None)
        logger.info("Tagged {}'s repo".format(team_name))
<|endoftext|>


<commit_before>"""

"""

from framework.routing import Rule, json_renderer
from website.routes import OsfWebRenderer

from . import views

settings_routes = {
    'rules': [],
    'prefix': '/api/v1',
}

widget_routes = {
    'rules': [
        Rule([
            '/project/<pid>/osffiles/widget/',
            '/project/<pid>/node/<nid>/osffiles/widget/',
        ], 'get', views.osffiles_widget, json_renderer),
    ],
    'prefix': '/api/v1',
}

web_routes = {

    'rules': [

        Rule([
            '/project/<pid>/osffiles/',
            '/project/<pid>/node/<nid>/osffiles/',
        ], 'get', views.get_osffiles, OsfWebRenderer('../addons/osffiles/templates/osffiles_tree.mako')),

        Rule([
            '/project/<pid>/osffiles/<fid>/',
            '/project/<pid>/node/<nid>/osffiles/<fid>/',
        ], 'get', views.view_file, OsfWebRenderer('../addons/osffiles/templates/osffiles_file.mako')),

    ]

}

api_routes = {

    'rules': [

        Rule(
            [
                '/project/<pid>/osffiles/hgrid/',
                '/project/<pid>/node/<nid>/osffiles/hgrid/',
                '/project/<pid>/osffiles/hgrid/<path:path>/',
                '/project/<pid>/node/<nid>/osffiles/hgrid/<path:path>/',
            ],
            'get',
            views.get_osffiles,
            json_renderer,
        ),

        Rule([
            '/project/<pid>/file_paths/',
            '/project/<pid>/node/<nid>/file_paths/',
        ], 'get', views.list_file_paths, json_renderer),

        # Download file
        Rule([
            '/project/<pid>/osffiles/<fid>/',
            '/project/<pid>/node/<nid>/osffiles/<fid>/',
        ], 'get', views.download_file, json_renderer),

        # Download file by version
        Rule([
            '/project/<pid>/osffiles/<fid>/version/<vid>/',
            '/project/<pid>/node/<nid>/osffiles/<fid>/version/<vid>/',
        ], 'get', views.download_file_by_version, json_renderer),

        Rule(
            [
                '/project/<pid>/osffiles/',
                '/project/<pid>/node/<nid>/osffiles/',
            ],
            'post',
            views.upload_file_public,
            json_renderer,
        ),
        Rule(
            [
                '/project/<pid>/osffiles/<fid>/',
                '/project/<pid>/node/<nid>/osffiles/<fid>/',
            ],
            'delete',
            views.delete_file,
            json_renderer,
        ),

    ],

    'prefix': '/api/v1',

}
<commit_msg>Include old URLs for backwards compatibility<commit_after>"""

"""

from framework.routing import Rule, json_renderer
from website.routes import OsfWebRenderer

from . import views

settings_routes = {
    'rules': [],
    'prefix': '/api/v1',
}

widget_routes = {
    'rules': [
        Rule([
            '/project/<pid>/osffiles/widget/',
            '/project/<pid>/node/<nid>/osffiles/widget/',
        ], 'get', views.osffiles_widget, json_renderer),
    ],
    'prefix': '/api/v1',
}

web_routes = {

    'rules': [

        Rule([
            '/project/<pid>/osffiles/',
            '/project/<pid>/node/<nid>/osffiles/',
        ], 'get', views.get_osffiles, OsfWebRenderer('../addons/osffiles/templates/osffiles_tree.mako')),

        Rule([
            '/project/<pid>/osffiles/<fid>/',
            '/project/<pid>/node/<nid>/osffiles/<fid>/',
        ], 'get', views.view_file, OsfWebRenderer('../addons/osffiles/templates/osffiles_file.mako')),

    ]

}

api_routes = {

    'rules': [

        Rule(
            [
                '/project/<pid>/osffiles/hgrid/',
                '/project/<pid>/node/<nid>/osffiles/hgrid/',
                '/project/<pid>/osffiles/hgrid/<path:path>/',
                '/project/<pid>/node/<nid>/osffiles/hgrid/<path:path>/',
            ],
            'get',
            views.get_osffiles,
            json_renderer,
        ),

        Rule([
            '/project/<pid>/file_paths/',
            '/project/<pid>/node/<nid>/file_paths/',
        ], 'get', views.list_file_paths, json_renderer),

        # Download file
        Rule([
            '/project/<pid>/osffiles/<fid>/',
            '/project/<pid>/node/<nid>/osffiles/<fid>/',
            # Note: Added these old URLs for backwards compatibility with
            # hard-coded links.
            '/project/<pid>/files/download/<fid>/',
            '/project/<pid>/node/<nid>/files/download/<fid>/',
        ], 'get', views.download_file, json_renderer),

        # Download file by version
        Rule([
            '/project/<pid>/osffiles/<fid>/version/<vid>/',
            '/project/<pid>/node/<nid>/osffiles/<fid>/version/<vid>/',
        ], 'get', views.download_file_by_version, json_renderer),

        Rule(
            [
                '/project/<pid>/osffiles/',
                '/project/<pid>/node/<nid>/osffiles/',
            ],
            'post',
            views.upload_file_public,
            json_renderer,
        ),
        Rule(
            [
                '/project/<pid>/osffiles/<fid>/',
                '/project/<pid>/node/<nid>/osffiles/<fid>/',
            ],
            'delete',
            views.delete_file,
            json_renderer,
        ),

    ],

    'prefix': '/api/v1',

}
<|endoftext|>


<commit_before>from __future__ import unicode_literals

import re
from collections import Counter

from rtmbot.core import Plugin

class PluginRepeat(Plugin):
    def __init__(self,slack_client=None, plugin_config=None):
        # because of the way plugins are called we must explicitly pass the
        # arguments to the super
        super(PluginRepeat, self).__init__(
            slack_client=slack_client,
            plugin_config=plugin_config
        )
        with open("words.txt") as f:
            self.words = f.read().split()

    def process_message(self, data):
        print(data)
        if (data['channel'].startswith("C") and data['type']=='message'):
            word_counter = self.check_for_flag_words(data['text'])
            if word_counter:
                resp = self.slack_client.api_call('im.open', user =data['user'])
                if resp['ok']:
                    self.outputs.append([
                        resp['channel']['id'],
                        'I counted you saying the following words this many '
                        'times: \n{}'.format(word_counter)
                    ])

    def check_for_flag_words(self, message):
        # type: (str) -> Dict
        """

        :rtype: Dict
        """
        cnt = Counter()
        delims = '!"#$%&()*+,./:;<=>?@[\\]^_`{|}~\t\n\x0b\x0c\r '
        pattern = r"[{}]".format(delims)
        message_array = re.split(pattern, message.lower())
        for word in message_array:
            word_v2 = word.replace(u"\u2019", "")
            word_v3 = word_v2.replace("'", "")
            if word_v3 in self.words:
                cnt[word_v3] += 1
        return dict(cnt)<commit_msg>Add load_word_list function and chain replace calls<commit_after>from __future__ import unicode_literals

import re
from collections import Counter

from rtmbot.core import Plugin

class PluginRepeat(Plugin):
    def __init__(self,slack_client=None, plugin_config=None):
        # because of the way plugins are called we must explicitly pass the
        # arguments to the super
        super(PluginRepeat, self).__init__(
            slack_client=slack_client,
            plugin_config=plugin_config
        )
        self.words = self.load_word_list("words.txt")

    def load_word_list(self, filepath):
        with open(filepath) as f:
            words = f.read().split()
            return words

    def process_message(self, data):
        print(data)
        if (data['channel'].startswith("C") and data['type']=='message'):
            word_counter = self.check_for_flag_words(data['text'])
            if word_counter:
                resp = self.slack_client.api_call('im.open', user =data['user'])
                if resp['ok']:
                    self.outputs.append([
                        resp['channel']['id'],
                        'I counted you saying the following words this many '
                        'times: \n{}'.format(word_counter)
                    ])

    def check_for_flag_words(self, message):
        # type: (str) -> Dict
        """

        :rtype: Dict
        """
        cnt = Counter()
        delims = '!"#$%&()*+,./:;<=>?@[\\]^_`{|}~\t\n\x0b\x0c\r '
        pattern = r"[{}]".format(delims)
        message_array = re.split(pattern, message.lower())
        for word in message_array:
            word_v2 = word.replace(u"\u2019", "").replace("'", "")
            if word_v2 in self.words:
                cnt[word_v2] += 1
        return dict(cnt)<|endoftext|>


<commit_before>import re
from abc import ABCMeta
from inspect import isfunction

from six import with_metaclass

from neupy.utils import AttributeKeyDict

__all__ = ("SharedDocsMeta", "SharedDocs", "SharedDocsException",
           "SharedDocsABCMeta", 'shared_docs')

def merge_dicts(left_dict, right_dict):
    """
    Merge two dictionaries in one.

    Parameters
    ----------
    left_dict : dict
    right_dict : dict

    Returns
    -------
    dict
    """
    return dict(left_dict, **right_dict)

def find_numpy_doc_indent(docs):
    """
    Find indent for Numpy styled documentation and return
    number of shifts inside of it.

    Parameters
    ----------
    docs : str

    Returns
    -------
    int or None
        Returns number of indentations in documentation. If
        it doesn't identify indentation function output will
        be ``None`` value.
    """

    indent_detector = re.compile(r"(?P<indent>\ *)(?P<dashes>-{3,})")
    indent_info = indent_detector.findall(docs)

    if indent_info:
        indent, _ = indent_info[0]
        return len(indent)

def iter_parameters(docs):
    """
    Find parameters defined in the documentation.

    Parameters
    ----------
    docs : str

    Yields
    ------
    tuple
        Yields tuple that contain 3 values. There are: parameter
        name, parameter type and parameter description.
    """

    n_indents = find_numpy_doc_indent(docs)
    doc_indent = ' ' * n_indents if n_indents else ''
    parser = re.compile(
        r"(?P<name>\w+?)\s*\:\s*(?P<type>[^\n]+)"
        r"((?P<description>(\n{indent}\ +[^\n]+)|(\n))*)"
        "".format(indent=doc_indent)
    )

    for name, type_, desc, _, _, _ in parser.findall(docs):
        yield (name, type_, desc)

def iter_methods(docs):
    """
    Find methods defined in the documentation.

    Parameters
    ----------
    docs : str

    Yields
    ------
    tuple
        Yields tuple that contain 3 values. There are: method
        name, method parameters and method description.
    """

    n_indents = find_numpy_doc_indent(docs)
    doc_indent = ' ' * n_indents if n_indents else ''
    parser = re.compile(
        r"(?P<name>\w+?)(\((.+?)?\))"
        r"((?P<description>\n{indent}\ +[^\n]+)*)"
        "".format(indent=doc_indent)
    )

    for name, func_params, _, desc, _ in parser.findall(docs):
        yield (name, func_params, desc)

def parse_warns(docs):
    """
    Find warning defined in the documentation.

    Parameters
    ----------
    docs : str

    Returns
    -------
    str or None
        Returns warnings from documentation or ``None`` if
        function didn't find it.
    """

    parser = re.compile(r"Warns\s+-+\s+(?P<warntext>(.+\n)+)")
    doc_warns = parser.findall(docs)

    if not doc_warns:
        return None

    doc_warns, _ = doc_warns[0]
    return doc_warns

def parse_variables_from_docs(instance, parent_instances):
    """
    Parse documentation with NumPy style and returns all
    extracted information.

    Parameters
    ----------
    instances : list
        List of objects that has documentations.

    Returns
    -------
    dict
        Variables parsed from the documentations.
    """
    variables = {}

    if not parent_instances:
        return variables

    for instance in parent_instances:
        parent_docs = instance.__doc__

        if parent_docs is None:
            continue

        parent_variables = AttributeKeyDict()

        for name, type_, desc in iter_parameters(parent_docs):
            parameter = "{} : {}{}".format(name, type_, desc.rstrip())
            parent_variables[name] = parameter

        for name, func_params, desc in iter_methods(parent_docs):
            parent_variables[name] = ''.join([name, func_params, desc])

        doc_warns = parse_warns(parent_docs)
        if doc_warns is not None:
            parent_variables['Warns'] = doc_warns

        parent_name = instance.__name__
        variables[parent_name] = parent_variables

    return variables

def format_docs(instance, parent_instances):
    """
    Format instance's documentation.

    Parameters
    ----------
    instance : object
        Any object that has documentation.
    parent_instances : list
        List of object that has documentations. Function will
        extract all information from theirs documentations and
        it will use them to format main instance documentation.

    Returns
    -------
    str
        Formated documentation.

    Raises
    ------
    SharedDocsException
        If function cannot format documentation properly.
    """
    try:
        instance_docs = instance.__doc__
        variables = parse_variables_from_docs(instance, parent_instances)
        instance_new_docs = instance_docs.format(**variables)
        # If we have multiple spaces between words, we need
        # to trim them. For instance:
        # change "hello   world" to "hello world"
        instance_new_docs = re.sub(
            pattern=r'([\S]+)(\ {2,})([\S]+)',
            repl=r'\1 \3',
            string=instance_new_docs
        )
        return instance_new_docs

    except Exception as exception:
        exception_classname = exception.__class__.__name__
        raise SharedDocsException(
            "Can't format documentation for class `{}`. "
            "Catched `{}` exception with message: {}".format(
                instance.__name__,
                exception_classname,
                exception
            )
        )

class SharedDocsException(Exception):
    """
    Exception that help identify problems related to shared
    documentation.
    """

def inherit_docs_for_methods(class_, attrs):
    """
    Class methods inherit documentation from the parent
    classes in case if methods doesn't have it.

    Parameters
    ----------
    class_ : object
    attrs : dict
        Class attributes.
    """
    for attrname, attrvalue in attrs.items():
        if not isfunction(attrvalue) or attrvalue.__doc__ is not None:
            continue

        for parent_class in class_.__mro__:
            if not hasattr(parent_class, attrname):
                continue

            parent_attrvalue = getattr(parent_class, attrname)
            if parent_attrvalue.__doc__ is not None:
                attrvalue.__doc__ = parent_attrvalue.__doc__
                break

class SharedDocsMeta(type):
    """
    Meta-class for shared documentation. This class conatains
    main functionality that help inherit parameters and methods
    descriptions from parent classes. This class automaticaly
    format class documentation using basic python format syntax
    for objects.

    Attributes
    ----------
    inherit_method_docs : bool
        ``True`` means that methods that doesn't have
        documentation will be inherited from the parent
        methods. ``False`` will disable this option for
        the specified class. Defaults to ``True``.
    """
    def __new__(cls, clsname, bases, attrs):
        new_class = super(SharedDocsMeta, cls).__new__(cls, clsname,
                                                       bases, attrs)

        if attrs.get('inherit_method_docs', True):
            inherit_docs_for_methods(new_class, attrs)

        if new_class.__doc__ is None:
            return new_class

        class_docs = new_class.__doc__
        n_indents = find_numpy_doc_indent(class_docs)

        if n_indents is not None:
            new_class.__doc__ = format_docs(new_class, new_class.__mro__)

        return new_class

class SharedDocsABCMeta(SharedDocsMeta, ABCMeta):
    """
    Meta-class that combine ``SharedDocsMeta`` and ``ABCMeta``
    meta-classes.
    """

class SharedDocs(with_metaclass(SharedDocsMeta)):
    """
    Main class that provide with shared documentation
    functionality.
    """

def shared_docs(parent_function):
    """
    Decorator shares documentation between functions.

    Parameters
    ----------
    parent_function : object
        Any object that has documentation.
    """
    def decorator(function):
        function.__doc__ = format_docs(function, [parent_function])
        return function
    return decorator
<commit_msg>Create a separate func that checks docs existance for the object.<commit_after>import re
from abc import ABCMeta
from inspect import isfunction

from six import with_metaclass

from neupy.utils import AttributeKeyDict

__all__ = ("SharedDocsMeta", "SharedDocs", "SharedDocsException",
           "SharedDocsABCMeta", 'shared_docs')

def merge_dicts(left_dict, right_dict):
    """
    Merge two dictionaries in one.

    Parameters
    ----------
    left_dict : dict
    right_dict : dict

    Returns
    -------
    dict
    """
    return dict(left_dict, **right_dict)

def find_numpy_doc_indent(docs):
    """
    Find indent for Numpy styled documentation and return
    number of shifts inside of it.

    Parameters
    ----------
    docs : str

    Returns
    -------
    int or None
        Returns number of indentations in documentation. If
        it doesn't identify indentation function output will
        be ``None`` value.
    """

    indent_detector = re.compile(r"(?P<indent>\ *)(?P<dashes>-{3,})")
    indent_info = indent_detector.findall(docs)

    if indent_info:
        indent, _ = indent_info[0]
        return len(indent)

def iter_parameters(docs):
    """
    Find parameters defined in the documentation.

    Parameters
    ----------
    docs : str

    Yields
    ------
    tuple
        Yields tuple that contain 3 values. There are: parameter
        name, parameter type and parameter description.
    """

    n_indents = find_numpy_doc_indent(docs)
    doc_indent = ' ' * n_indents if n_indents else ''
    parser = re.compile(
        r"(?P<name>\w+?)\s*\:\s*(?P<type>[^\n]+)"
        r"((?P<description>(\n{indent}\ +[^\n]+)|(\n))*)"
        "".format(indent=doc_indent)
    )

    for name, type_, desc, _, _, _ in parser.findall(docs):
        yield (name, type_, desc)

def iter_methods(docs):
    """
    Find methods defined in the documentation.

    Parameters
    ----------
    docs : str

    Yields
    ------
    tuple
        Yields tuple that contain 3 values. There are: method
        name, method parameters and method description.
    """

    n_indents = find_numpy_doc_indent(docs)
    doc_indent = ' ' * n_indents if n_indents else ''
    parser = re.compile(
        r"(?P<name>\w+?)(\((.+?)?\))"
        r"((?P<description>\n{indent}\ +[^\n]+)*)"
        "".format(indent=doc_indent)
    )

    for name, func_params, _, desc, _ in parser.findall(docs):
        yield (name, func_params, desc)

def parse_warns(docs):
    """
    Find warning defined in the documentation.

    Parameters
    ----------
    docs : str

    Returns
    -------
    str or None
        Returns warnings from documentation or ``None`` if
        function didn't find it.
    """

    parser = re.compile(r"Warns\s+-+\s+(?P<warntext>(.+\n)+)")
    doc_warns = parser.findall(docs)

    if not doc_warns:
        return None

    doc_warns, _ = doc_warns[0]
    return doc_warns

def parse_variables_from_docs(instance, parent_instances):
    """
    Parse documentation with NumPy style and returns all
    extracted information.

    Parameters
    ----------
    instances : list
        List of objects that has documentations.

    Returns
    -------
    dict
        Variables parsed from the documentations.
    """
    variables = {}

    if not parent_instances:
        return variables

    for instance in parent_instances:
        parent_docs = instance.__doc__

        if parent_docs is None:
            continue

        parent_variables = AttributeKeyDict()

        for name, type_, desc in iter_parameters(parent_docs):
            parameter = "{} : {}{}".format(name, type_, desc.rstrip())
            parent_variables[name] = parameter

        for name, func_params, desc in iter_methods(parent_docs):
            parent_variables[name] = ''.join([name, func_params, desc])

        doc_warns = parse_warns(parent_docs)
        if doc_warns is not None:
            parent_variables['Warns'] = doc_warns

        parent_name = instance.__name__
        variables[parent_name] = parent_variables

    return variables

def format_docs(instance, parent_instances):
    """
    Format instance's documentation.

    Parameters
    ----------
    instance : object
        Any object that has documentation.
    parent_instances : list
        List of object that has documentations. Function will
        extract all information from theirs documentations and
        it will use them to format main instance documentation.

    Returns
    -------
    str
        Formated documentation.

    Raises
    ------
    SharedDocsException
        If function cannot format documentation properly.
    """
    try:
        instance_docs = instance.__doc__
        variables = parse_variables_from_docs(instance, parent_instances)
        instance_new_docs = instance_docs.format(**variables)
        # If we have multiple spaces between words, we need
        # to trim them. For instance:
        # change "hello   world" to "hello world"
        instance_new_docs = re.sub(
            pattern=r'([\S]+)(\ {2,})([\S]+)',
            repl=r'\1 \3',
            string=instance_new_docs
        )
        return instance_new_docs

    except Exception as exception:
        exception_classname = exception.__class__.__name__
        raise SharedDocsException(
            "Can't format documentation for class `{}`. "
            "Catched `{}` exception with message: {}".format(
                instance.__name__,
                exception_classname,
                exception
            )
        )

class SharedDocsException(Exception):
    """
    Exception that help identify problems related to shared
    documentation.
    """

def has_docs(value):
    """ Checks whether object has documentation.

    Parameters
    ----------
    value : object

    Returns
    -------
    bool
        Function returns ``True`` if object has a documentation
        and ``False`` otherwise.
    """
    return value.__doc__ is not None

def inherit_docs_for_methods(class_, attrs):
    """
    Class methods inherit documentation from the parent
    classes in case if methods doesn't have it.

    Parameters
    ----------
    class_ : object
    attrs : dict
        Class attributes.
    """
    for attrname, attrvalue in attrs.items():
        if not isfunction(attrvalue) or has_docs(attrvalue):
            continue

        for parent_class in class_.__mro__:
            if not hasattr(parent_class, attrname):
                continue

            parent_attrvalue = getattr(parent_class, attrname)
            if has_docs(parent_attrvalue):
                attrvalue.__doc__ = parent_attrvalue.__doc__
                break

class SharedDocsMeta(type):
    """
    Meta-class for shared documentation. This class conatains
    main functionality that help inherit parameters and methods
    descriptions from parent classes. This class automaticaly
    format class documentation using basic python format syntax
    for objects.

    Attributes
    ----------
    inherit_method_docs : bool
        ``True`` means that methods that doesn't have
        documentation will be inherited from the parent
        methods. ``False`` will disable this option for
        the specified class. Defaults to ``True``.
    """
    def __new__(cls, clsname, bases, attrs):
        new_class = super(SharedDocsMeta, cls).__new__(cls, clsname,
                                                       bases, attrs)

        if attrs.get('inherit_method_docs', True):
            inherit_docs_for_methods(new_class, attrs)

        if new_class.__doc__ is None:
            return new_class

        class_docs = new_class.__doc__
        n_indents = find_numpy_doc_indent(class_docs)

        if n_indents is not None:
            new_class.__doc__ = format_docs(new_class, new_class.__mro__)

        return new_class

class SharedDocsABCMeta(SharedDocsMeta, ABCMeta):
    """
    Meta-class that combine ``SharedDocsMeta`` and ``ABCMeta``
    meta-classes.
    """

class SharedDocs(with_metaclass(SharedDocsMeta)):
    """
    Main class that provide with shared documentation
    functionality.
    """

def shared_docs(parent_function):
    """
    Decorator shares documentation between functions.

    Parameters
    ----------
    parent_function : object
        Any object that has documentation.
    """
    def decorator(function):
        function.__doc__ = format_docs(function, [parent_function])
        return function
    return decorator
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
MPI support classes.
"""

import inspect
import os
import re
import subprocess
import sys

from mpi4py import MPI
import msgpack
import psutil
import shortuuid
import twiggy
import zmq

from tools.logging import format_name, setup_logger, set_excepthook

def getargnames(f):
    """
    Get names of a callable's arguments.

    Parameters
    ----------
    f : callable
        Function to examine.

    Results
    -------
    args : list of str
        Argument names.

    Notes
    -----
    For instance methods, the `self` argument is omitted.
    """

    spec = inspect.getargspec(f)
    if inspect.ismethod(f):
        return spec.args[1:]
    else:
        return spec.args

def args_to_dict(f, *args, **kwargs):
    """
    Combine sequential and named arguments in single dictionary.

    Parameters
    ----------
    f : callable
        Function to which the arguments will be passed.
    args : tuple
        Sequential arguments.
    kwargs : dict
        Named arguments.

    Returns
    -------
    d : dict
        Maps argument names to values.
    """

    spec = inspect.getargspec(f)
    d = {}

    arg_names = getargnames(f)
    assert len(arg_names) <= args
    for arg, val in zip(arg_names, args):
        d[arg] = val
    for arg, val in kwargs.iteritems():
        if arg in d:
            raise ValueError('\'%s\' already specified in positional args' % arg)
        d[arg] = val
    return d

class PollerChecker(object):
    """
    Wrapper class to facilitate creation and use of ZMQ pollers.

    Parameters
    ----------
    sock : zmq.Socket
        ZeroMQ socket to poll.
    direction : int
        Polling direction.
    """

    def __init__(self, sock, direction=zmq.POLLIN):
        self._sock = sock
        self._poller = zmq.Poller()
        self._poller.register(sock, direction)

    def check(self, timeout=None):
        """
        Check for I/O.

        Parameters
        ----------
        timeout : float, int
            Timeout in milliseconds. If None, no timeout is assumed.

        Returns
        -------
        status : bool
            True if transmitted messages are available.
        """

        socks = dict(self._poller.poll(timeout))
        if self._sock in socks:
            return True
        else:
            return False

class Worker(object):
    """
    MPI worker class.

    This class repeatedly executes a work method.

    Parameters
    ----------
    data_tag : int
        MPI tag to identify data messages transmitted between worker nodes.
    ctrl_tag : int
        MPI tag to identify control messages transmitted to worker nodes.
    """

    def __init__(self, data_tag=0, ctrl_tag=1):
        rank = MPI.COMM_WORLD.Get_rank()
        self.logger = twiggy.log.name(format_name('worker %s' % rank))
        set_excepthook(self.logger, True)

        # Tags used to distinguish MPI messages:
        self._data_tag = data_tag
        self._ctrl_tag = ctrl_tag

        # Maximum number of execution steps:
        self.steps = float('inf')

    def do_work(self):
        """
        Work method.

        This method is repeatedly executed by the Worker instance after the
        instance receives a 'start' control message and until it receives a 'stop'
        control message. It should be overridden by child classes.
        """

        self.logger.info('executing do_work')

    def run(self):
        """
        Main body of worker process.
        """

        rank = MPI.COMM_WORLD.Get_rank()
        self.logger.info('running body of worker %s' % rank)

        # Start listening for control messages:
        r_ctrl = []
        r_ctrl.append(MPI.COMM_WORLD.irecv(source=0, tag=self._ctrl_tag))

        running = False
        req = MPI.Request()
        steps = 0
        while True:

            # Handle control messages:
            flag, msg = req.testall(r_ctrl)
            if flag:

                # Deserialize:
                msg = msgpack.loads(msg[0])

                # Start listening for data messages:
                if msg[0] == 'start':
                    self.logger.info('started')
                    running = True

                # Stop listening for data messages:
                elif msg[0] == 'stop':
                    self.logger.info('stopped')
                    running = False

                # Set maximum number of execution steps:
                elif msg[0] == 'steps':
                    if msg[1] == 'inf':
                        self.steps = float('inf')
                    else:
                        self.steps = int(msg[1])
                    self.logger.info('steps set to %s' % self.steps)

                # Quit:
                elif msg[0] == 'quit':
                    self.logger.info('acknowledging quit')

                    # Block until acknowledgment message received:
                    MPI.COMM_WORLD.send(str(rank), dest=0, tag=self._ctrl_tag)
                    self.logger.info('quit')
                    return

                r_ctrl = []
                r_ctrl.append(MPI.COMM_WORLD.irecv(source=0, tag=self._ctrl_tag))

            # Execute work method:
            if running:
                self.do_work()
                steps += 1
                self.logger.info('execution step: %s' % steps)

            # Send acknowledgment back to master if a finite number of steps was
            # specified and then completed:
            if steps > self.steps:
                self.logger.info('acknowledging step')

                # Block until acknowledgment message received:
                MPI.COMM_WORLD.send(str(rank), dest=0, tag=self._ctrl_tag)
                self.logger.info('step')
                return

class Manager(object):
    """
    Self-launching MPI worker manager.

    This class may be used to construct an MPI application consisting of

    - a launcher process with methods for specifying Worker class instances as
      the bodies of MPI nodes and for starting and stopping the application's
      execution;
    - worker processes that perform some processing task; and
    - a master process that relays control messages from the launcher process to the
      worker processes.

    The application should NOT be started via mpiexec.

    Parameters
    ----------
    mpiexec : str
        Name of MPI launcher executable.
    mpiargs : tuple
        Additional arguments to pass to MPI launcher.
    data_tag : int
        MPI tag to identify data messages transmitted between worker nodes.
    ctrl_tag : int
        MPI tag to identify control messages transmitted to worker nodes.

    Notes
    -----
    This class does not require MPI-2 dynamic processing management.

    See Also
    --------
    Worker
    """

    def __init__(self, mpiexec='mpiexec', mpiargs=(), data_tag=0, ctrl_tag=1):
        assert data_tag != ctrl_tag and data_tag != MPI.ANY_TAG and \
                           ctrl_tag != MPI.ANY_TAG

        # MPI launch info:
        self._mpiexec = mpiexec
        self._mpiargs = tuple(mpiargs)

        # Make logger name reflect process identity:
        if self._is_launcher():
            self.logger = twiggy.log.name(format_name('launcher/manager'))
        elif self._is_master():
            self.logger = twiggy.log.name(format_name('master/manager'))
        else:
            self.logger = twiggy.log.name(format_name('worker %s/manager' % MPI.COMM_WORLD.Get_rank()))
        set_excepthook(self.logger, True)

        # Tags used to distinguish MPI messages:
        self._data_tag = data_tag
        self._ctrl_tag = ctrl_tag

        # Worker classes to instantiate:
        self._targets = {}

        # Arguments to pass to worker class constructors:
        self._kwargs = {}

        # Reserve node 0 for use as master:
        self._rank = 1

    @property
    def num_workers(self):
        """
        Number of workers known to manager.
        """

        return len(self._targets)

    def add(self, target, *args, **kwargs):
        """
        Add a worker to an MPI application.

        Parameters
        ----------
        target : Worker
            Worker class to instantiate and run.
        args : sequence
            Sequential arguments to pass to target class constructor.
        kwargs : dict
            Named arguments to pass to target class constructor.

        Returns
        -------
        rank : int
            MPI rank assigned to class.
        """

        # The contents of this method need to run on the launcher (so that it
        # knows how many MPI processes to start), the workers (so that they
        # know how to instantiate their respective classes), and the master (so
        # that it will be aware of all of the MPI processes):
        self.logger.info('adding class %s' % target.__name__)
        assert issubclass(target, Worker)

        rank = self._rank
        self._targets[rank] = target
        self._kwargs[rank] = args_to_dict(target.__init__, *args, **kwargs)
        self._rank += 1
        return rank

    def _in_mpi(self):
        """
        Return True if the current process is an MPI process.
        """

        # Get name of parent process:
        parent_name = psutil.Process(os.getppid()).name()

        # All MPI processes must be children of the launcher process:
        return bool(re.search(os.path.basename(self._mpiexec), parent_name))

    def _is_launcher(self):
        """
        Return True if the current process is the launching process.
        """

        return not self._in_mpi()

    def _is_master(self):
        """
        Return True if the current process is the master MPI process.
        """

        # We need to check if MPI is active because Get_rank() will return zero
        # in the launcher too:
        return self._in_mpi() and (MPI.COMM_WORLD.Get_rank() == 0)

    def _run_launcher(self):
        """
        Asynchronously launch MPI application, connect master and launcher with ZeroMQ.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping launch')
            return
        else:
            self.logger.info('launching application (%s)' % self._rank)

        # Create random IPC interface name:
        env = os.environ.copy()
        env['MASTER_IPC_INT'] = 'ipc://'+shortuuid.uuid()

        # Set up communication socket to master process:
        ctx = zmq.Context()
        self._sock = ctx.socket(zmq.ROUTER)
        self._sock.bind(env['MASTER_IPC_INT'])
        self._pc = PollerChecker(self._sock, zmq.POLLIN)

        # Pass the IPC interface name to the launched process via an
        # environmental variable:
        python_path = sys.executable
        script_name = os.path.basename(sys.argv[0])
        self._mpiexec_proc = subprocess.Popen((self._mpiexec,)+self._mpiargs+\
                                              ('-np', str(self._rank),
                                               python_path, script_name),
                                              stdout=sys.stdout,
                                              stderr=sys.stderr,
                                              stdin=sys.stdin,
                                              env=env)
        self.logger.info('application launched')

        # Synchronize connection:
        while True:
            self._sock.send_multipart(['master', ''])
            if self._pc.check(10):
                self._sock.recv_multipart()
                break
        self.logger.info('launcher synchronized')

    def _run_master(self):
        """
        Body of master MPI process.

        The only function of the master process is to relay messages from the
        launcher to the other nodes.
        """

        if not self._is_master():
            self.logger.info('not in master - skipping _run_master')
            return
        else:
            self.logger.info('running body of master')

        ctx = zmq.Context()
        self._sock = ctx.socket(zmq.DEALER)
        self._sock.setsockopt(zmq.IDENTITY, 'master')
        self._sock.connect(os.environ['MASTER_IPC_INT'])
        self._pc = PollerChecker(self._sock, zmq.POLLIN)

        # Synchronize connection:
        while True:
            if self._pc.check(10):
                self._sock.recv()
                self._sock.send('')
                break
        self.logger.info('master synchronized')

        # Relay messages from launcher to workers until a quit or step message is received:
        # XXX currently only broadcasts messages to workers; could be extended
        # to permit directed transmission to specific workers:
        size = MPI.COMM_WORLD.Get_size()
        while True:

            # Check for messages from launcher:
            if self._pc.check(10):
                msg = self._sock.recv_multipart()

                # Pass any messages on to all of the workers:
                self.logger.info('sending message to workers: '+str(msg))
                for i in xrange(1, size):
                    MPI.COMM_WORLD.isend(msgpack.dumps(msg),
                                         dest=i, tag=self._ctrl_tag)

                # Wait for all workers to acknowledge quit or step message:
                if msg[0] == 'quit' or msg[0] == 'step':
                    req = MPI.Request()
                    r_quit = []
                    self.logger.info('waiting for quit/step acknowledgments')
                    for i in xrange(1, size):
                        r_quit.append(MPI.COMM_WORLD.irecv(source=MPI.ANY_SOURCE,
                                                           tag=self._ctrl_tag))
                    while not req.testall(r_quit):
                        pass

                    self.logger.info('finished running master')
                    break

    def run(self):
        """
        Run MPI application.

        Notes
        -----
        This method must be called after all callable code that is to be run on
        the MPI nodes has been added to the manager. This method will return
        immediately on the launcher but not on the master or worker processes.
        """

        if self._is_launcher():
            self._run_launcher()
        elif self._is_master():
            self._run_master()
        else:

            # Instantiate each target using the specified parameters and
            # execute the target's run() method:
            rank = MPI.COMM_WORLD.Get_rank()
            t = self._targets[rank](**self._kwargs[rank])
            t.run()
            self.logger.info('finished running %s' % rank)

    def start(self):
        """
        Tell the workers to start processing data.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping start')
            return
        self.logger.info('sending start message')
        self._sock.send_multipart(['master', 'start'])

    def stop(self):
        """
        Tell the workers to stop processing data.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping stop')
            return
        self.logger.info('sending stop message')
        self._sock.send_multipart(['master', 'stop'])

    def steps(self, n=float('inf')):
        """
        Tell the workers to run no more than the specified number of steps.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping steps')
            return
        self.logger.info('sending steps message (%s)' % n)
        self._sock.send_multipart(['master', 'steps', str(n)])

    def quit(self):
        """
        Tell the workers to quit.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping quit')
            return
        self.logger.info('sending quit message')
        self._sock.send_multipart(['master', 'quit'])

    def kill(self):
        """
        Kill MPI launcher.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping kill')
            return
        self.logger.info('killing launcher')
        self._mpiexec_proc.kill()

    def wait(self):
        """
        Wait for MPI launcher to exit.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping wait')
            return
        self.logger.info('waiting for launcher to exit')
        self._mpiexec_proc.wait()

    def terminate(self):
        """
        Terminate MPI launcher.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping')
            return
        self.logger.info('terminating launcher')
        self._mpiexec_proc.terminate()

if __name__ == '__main__':
    import time

    setup_logger(stdout=sys.stdout)

    # Define a class whose constructor takes arguments so as to test
    # instantiation of the class by the manager:
    class MyWorker(Worker):
        def __init__(self, x, y, z=None):
            super(MyWorker, self).__init__()
            comm = MPI.COMM_WORLD
            size = comm.Get_size()
            rank = comm.Get_rank()
            name = MPI.Get_processor_name()
            self.logger.info('I am process %d of %d on %s.' % (rank, size, name))
            self.logger.info('init args: %s, %s, %s' % (x, y, z))

    man = Manager()
    man.add(target=MyWorker, x=1, y=2, z=3)
    man.add(MyWorker, 3, 4, 5)
    man.run()
    man.start()
    time.sleep(1)
    man.stop()
    man.quit()
<commit_msg>Rename execution step counter attrib.<commit_after>#!/usr/bin/env python

"""
MPI support classes.
"""

import inspect
import os
import re
import subprocess
import sys

from mpi4py import MPI
import msgpack
import psutil
import shortuuid
import twiggy
import zmq

from tools.logging import format_name, setup_logger, set_excepthook

def getargnames(f):
    """
    Get names of a callable's arguments.

    Parameters
    ----------
    f : callable
        Function to examine.

    Results
    -------
    args : list of str
        Argument names.

    Notes
    -----
    For instance methods, the `self` argument is omitted.
    """

    spec = inspect.getargspec(f)
    if inspect.ismethod(f):
        return spec.args[1:]
    else:
        return spec.args

def args_to_dict(f, *args, **kwargs):
    """
    Combine sequential and named arguments in single dictionary.

    Parameters
    ----------
    f : callable
        Function to which the arguments will be passed.
    args : tuple
        Sequential arguments.
    kwargs : dict
        Named arguments.

    Returns
    -------
    d : dict
        Maps argument names to values.
    """

    spec = inspect.getargspec(f)
    d = {}

    arg_names = getargnames(f)
    assert len(arg_names) <= args
    for arg, val in zip(arg_names, args):
        d[arg] = val
    for arg, val in kwargs.iteritems():
        if arg in d:
            raise ValueError('\'%s\' already specified in positional args' % arg)
        d[arg] = val
    return d

class PollerChecker(object):
    """
    Wrapper class to facilitate creation and use of ZMQ pollers.

    Parameters
    ----------
    sock : zmq.Socket
        ZeroMQ socket to poll.
    direction : int
        Polling direction.
    """

    def __init__(self, sock, direction=zmq.POLLIN):
        self._sock = sock
        self._poller = zmq.Poller()
        self._poller.register(sock, direction)

    def check(self, timeout=None):
        """
        Check for I/O.

        Parameters
        ----------
        timeout : float, int
            Timeout in milliseconds. If None, no timeout is assumed.

        Returns
        -------
        status : bool
            True if transmitted messages are available.
        """

        socks = dict(self._poller.poll(timeout))
        if self._sock in socks:
            return True
        else:
            return False

class Worker(object):
    """
    MPI worker class.

    This class repeatedly executes a work method.

    Parameters
    ----------
    data_tag : int
        MPI tag to identify data messages transmitted between worker nodes.
    ctrl_tag : int
        MPI tag to identify control messages transmitted to worker nodes.
    """

    def __init__(self, data_tag=0, ctrl_tag=1):
        rank = MPI.COMM_WORLD.Get_rank()
        self.logger = twiggy.log.name(format_name('worker %s' % rank))
        set_excepthook(self.logger, True)

        # Tags used to distinguish MPI messages:
        self._data_tag = data_tag
        self._ctrl_tag = ctrl_tag

        # Maximum number of execution steps:
        self.max_steps = float('inf')

    def do_work(self):
        """
        Work method.

        This method is repeatedly executed by the Worker instance after the
        instance receives a 'start' control message and until it receives a 'stop'
        control message. It should be overridden by child classes.
        """

        self.logger.info('executing do_work')

    def run(self):
        """
        Main body of worker process.
        """

        rank = MPI.COMM_WORLD.Get_rank()
        self.logger.info('running body of worker %s' % rank)

        # Start listening for control messages:
        r_ctrl = []
        r_ctrl.append(MPI.COMM_WORLD.irecv(source=0, tag=self._ctrl_tag))

        running = False
        req = MPI.Request()
        steps = 0
        while True:

            # Handle control messages:
            flag, msg = req.testall(r_ctrl)
            if flag:

                # Deserialize:
                msg = msgpack.loads(msg[0])

                # Start listening for data messages:
                if msg[0] == 'start':
                    self.logger.info('started')
                    running = True

                # Stop listening for data messages:
                elif msg[0] == 'stop':
                    self.logger.info('stopped')
                    running = False

                # Set maximum number of execution steps:
                elif msg[0] == 'steps':
                    if msg[1] == 'inf':
                        self.max_steps = float('inf')
                    else:
                        self.max_steps = int(msg[1])
                    self.logger.info('maximum steps set to %s' % self.max_steps)

                # Quit:
                elif msg[0] == 'quit':
                    self.logger.info('acknowledging quit')

                    # Block until acknowledgment message received:
                    MPI.COMM_WORLD.send(str(rank), dest=0, tag=self._ctrl_tag)
                    self.logger.info('quit')
                    return

                r_ctrl = []
                r_ctrl.append(MPI.COMM_WORLD.irecv(source=0, tag=self._ctrl_tag))

            # Execute work method:
            if running:
                self.do_work()
                steps += 1
                self.logger.info('execution step: %s' % steps)

            # Send acknowledgment back to master if a finite number of steps was
            # specified and then completed:
            if steps > self.max_steps:
                self.logger.info('acknowledging steps')

                # Block until acknowledgment message received:
                MPI.COMM_WORLD.send(str(rank), dest=0, tag=self._ctrl_tag)
                self.logger.info('steps')
                return

class Manager(object):
    """
    Self-launching MPI worker manager.

    This class may be used to construct an MPI application consisting of

    - a launcher process with methods for specifying Worker class instances as
      the bodies of MPI nodes and for starting and stopping the application's
      execution;
    - worker processes that perform some processing task; and
    - a master process that relays control messages from the launcher process to the
      worker processes.

    The application should NOT be started via mpiexec.

    Parameters
    ----------
    mpiexec : str
        Name of MPI launcher executable.
    mpiargs : tuple
        Additional arguments to pass to MPI launcher.
    data_tag : int
        MPI tag to identify data messages transmitted between worker nodes.
    ctrl_tag : int
        MPI tag to identify control messages transmitted to worker nodes.

    Notes
    -----
    This class does not require MPI-2 dynamic processing management.

    See Also
    --------
    Worker
    """

    def __init__(self, mpiexec='mpiexec', mpiargs=(), data_tag=0, ctrl_tag=1):
        assert data_tag != ctrl_tag and data_tag != MPI.ANY_TAG and \
                           ctrl_tag != MPI.ANY_TAG

        # MPI launch info:
        self._mpiexec = mpiexec
        self._mpiargs = tuple(mpiargs)

        # Make logger name reflect process identity:
        if self._is_launcher():
            self.logger = twiggy.log.name(format_name('launcher/manager'))
        elif self._is_master():
            self.logger = twiggy.log.name(format_name('master/manager'))
        else:
            self.logger = twiggy.log.name(format_name('worker %s/manager' % MPI.COMM_WORLD.Get_rank()))
        set_excepthook(self.logger, True)

        # Tags used to distinguish MPI messages:
        self._data_tag = data_tag
        self._ctrl_tag = ctrl_tag

        # Worker classes to instantiate:
        self._targets = {}

        # Arguments to pass to worker class constructors:
        self._kwargs = {}

        # Reserve node 0 for use as master:
        self._rank = 1

    @property
    def num_workers(self):
        """
        Number of workers known to manager.
        """

        return len(self._targets)

    def add(self, target, *args, **kwargs):
        """
        Add a worker to an MPI application.

        Parameters
        ----------
        target : Worker
            Worker class to instantiate and run.
        args : sequence
            Sequential arguments to pass to target class constructor.
        kwargs : dict
            Named arguments to pass to target class constructor.

        Returns
        -------
        rank : int
            MPI rank assigned to class.
        """

        # The contents of this method need to run on the launcher (so that it
        # knows how many MPI processes to start), the workers (so that they
        # know how to instantiate their respective classes), and the master (so
        # that it will be aware of all of the MPI processes):
        self.logger.info('adding class %s' % target.__name__)
        assert issubclass(target, Worker)

        rank = self._rank
        self._targets[rank] = target
        self._kwargs[rank] = args_to_dict(target.__init__, *args, **kwargs)
        self._rank += 1
        return rank

    def _in_mpi(self):
        """
        Return True if the current process is an MPI process.
        """

        # Get name of parent process:
        parent_name = psutil.Process(os.getppid()).name()

        # All MPI processes must be children of the launcher process:
        return bool(re.search(os.path.basename(self._mpiexec), parent_name))

    def _is_launcher(self):
        """
        Return True if the current process is the launching process.
        """

        return not self._in_mpi()

    def _is_master(self):
        """
        Return True if the current process is the master MPI process.
        """

        # We need to check if MPI is active because Get_rank() will return zero
        # in the launcher too:
        return self._in_mpi() and (MPI.COMM_WORLD.Get_rank() == 0)

    def _run_launcher(self):
        """
        Asynchronously launch MPI application, connect master and launcher with ZeroMQ.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping launch')
            return
        else:
            self.logger.info('launching application (%s)' % self._rank)

        # Create random IPC interface name:
        env = os.environ.copy()
        env['MASTER_IPC_INT'] = 'ipc://'+shortuuid.uuid()

        # Set up communication socket to master process:
        ctx = zmq.Context()
        self._sock = ctx.socket(zmq.ROUTER)
        self._sock.bind(env['MASTER_IPC_INT'])
        self._pc = PollerChecker(self._sock, zmq.POLLIN)

        # Pass the IPC interface name to the launched process via an
        # environmental variable:
        python_path = sys.executable
        script_name = os.path.basename(sys.argv[0])
        self._mpiexec_proc = subprocess.Popen((self._mpiexec,)+self._mpiargs+\
                                              ('-np', str(self._rank),
                                               python_path, script_name),
                                              stdout=sys.stdout,
                                              stderr=sys.stderr,
                                              stdin=sys.stdin,
                                              env=env)
        self.logger.info('application launched')

        # Synchronize connection:
        while True:
            self._sock.send_multipart(['master', ''])
            if self._pc.check(10):
                self._sock.recv_multipart()
                break
        self.logger.info('launcher synchronized')

    def _run_master(self):
        """
        Body of master MPI process.

        The only function of the master process is to relay messages from the
        launcher to the other nodes.
        """

        if not self._is_master():
            self.logger.info('not in master - skipping _run_master')
            return
        else:
            self.logger.info('running body of master')

        ctx = zmq.Context()
        self._sock = ctx.socket(zmq.DEALER)
        self._sock.setsockopt(zmq.IDENTITY, 'master')
        self._sock.connect(os.environ['MASTER_IPC_INT'])
        self._pc = PollerChecker(self._sock, zmq.POLLIN)

        # Synchronize connection:
        while True:
            if self._pc.check(10):
                self._sock.recv()
                self._sock.send('')
                break
        self.logger.info('master synchronized')

        # Relay messages from launcher to workers until a quit or step message is received:
        # XXX currently only broadcasts messages to workers; could be extended
        # to permit directed transmission to specific workers:
        size = MPI.COMM_WORLD.Get_size()
        while True:

            # Check for messages from launcher:
            if self._pc.check(10):
                msg = self._sock.recv_multipart()

                # Pass any messages on to all of the workers:
                self.logger.info('sending message to workers: '+str(msg))
                for i in xrange(1, size):
                    MPI.COMM_WORLD.isend(msgpack.dumps(msg),
                                         dest=i, tag=self._ctrl_tag)

                # Wait for all workers to acknowledge quit or step message:
                if msg[0] == 'quit' or msg[0] == 'step':
                    req = MPI.Request()
                    r_quit = []
                    self.logger.info('waiting for quit/step acknowledgments')
                    for i in xrange(1, size):
                        r_quit.append(MPI.COMM_WORLD.irecv(source=MPI.ANY_SOURCE,
                                                           tag=self._ctrl_tag))
                    while not req.testall(r_quit):
                        pass

                    self.logger.info('finished running master')
                    break

    def run(self):
        """
        Run MPI application.

        Notes
        -----
        This method must be called after all callable code that is to be run on
        the MPI nodes has been added to the manager. This method will return
        immediately on the launcher but not on the master or worker processes.
        """

        if self._is_launcher():
            self._run_launcher()
        elif self._is_master():
            self._run_master()
        else:

            # Instantiate each target using the specified parameters and
            # execute the target's run() method:
            rank = MPI.COMM_WORLD.Get_rank()
            t = self._targets[rank](**self._kwargs[rank])
            t.run()
            self.logger.info('finished running %s' % rank)

    def start(self):
        """
        Tell the workers to start processing data.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping start')
            return
        self.logger.info('sending start message')
        self._sock.send_multipart(['master', 'start'])

    def stop(self):
        """
        Tell the workers to stop processing data.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping stop')
            return
        self.logger.info('sending stop message')
        self._sock.send_multipart(['master', 'stop'])

    def steps(self, n=float('inf')):
        """
        Tell the workers to run no more than the specified number of steps.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping steps')
            return
        self.logger.info('sending steps message (%s)' % n)
        self._sock.send_multipart(['master', 'steps', str(n)])

    def quit(self):
        """
        Tell the workers to quit.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping quit')
            return
        self.logger.info('sending quit message')
        self._sock.send_multipart(['master', 'quit'])

    def kill(self):
        """
        Kill MPI launcher.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping kill')
            return
        self.logger.info('killing launcher')
        self._mpiexec_proc.kill()

    def wait(self):
        """
        Wait for MPI launcher to exit.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping wait')
            return
        self.logger.info('waiting for launcher to exit')
        self._mpiexec_proc.wait()

    def terminate(self):
        """
        Terminate MPI launcher.
        """

        if not self._is_launcher():
            self.logger.info('not in launcher - skipping')
            return
        self.logger.info('terminating launcher')
        self._mpiexec_proc.terminate()

if __name__ == '__main__':
    import time

    setup_logger(stdout=sys.stdout)

    # Define a class whose constructor takes arguments so as to test
    # instantiation of the class by the manager:
    class MyWorker(Worker):
        def __init__(self, x, y, z=None):
            super(MyWorker, self).__init__()
            comm = MPI.COMM_WORLD
            size = comm.Get_size()
            rank = comm.Get_rank()
            name = MPI.Get_processor_name()
            self.logger.info('I am process %d of %d on %s.' % (rank, size, name))
            self.logger.info('init args: %s, %s, %s' % (x, y, z))

    man = Manager()
    man.add(target=MyWorker, x=1, y=2, z=3)
    man.add(MyWorker, 3, 4, 5)
    man.run()
    man.start()
    time.sleep(1)
    man.stop()
    man.quit()
<|endoftext|>


<commit_before>import numpy as np
import scipy as sp
import numpy.random as np_rnd

from ngboost.distns import Normal

from ngboost.scores import MLE, CRPS
from ngboost.learners import default_tree_learner, default_linear_learner
from ngboost.distns.normal import Normal

class NGBoost(object):

    def __init__(self, Dist=Normal, Score=MLE(),
                 Base=default_tree_learner, natural_gradient=False,
                 n_estimators=500, learning_rate=0.01, minibatch_frac=1.0,
                 verbose=True, tol=1e-4):
        self.Dist = Dist
        self.Score = Score
        self.Base = Base
        self.natural_gradient = natural_gradient
        self.n_estimators = n_estimators
        self.learning_rate = learning_rate
        self.minibatch_frac = minibatch_frac
        self.verbose = verbose
        self.init_params = None
        self.base_models = []
        self.scalings = []
        self.tol = tol

    def pred_param(self, X, max_iter=None):
        m, n = X.shape
        params = np.ones((m, self.Dist.n_params)) * self.init_params
        for i, (models, s) in enumerate(zip(self.base_models, self.scalings)):
            if max_iter and i == max_iter:
                break
            resids = np.array([model.predict(X) for model in models]).T
            params -= self.learning_rate * resids * s
        return params

    def sample(self, X, Y, params):
        sample_size = int(self.minibatch_frac * len(Y))
        idxs = np_rnd.choice(np.arange(len(Y)), sample_size, replace=False)
        return idxs, X[idxs,:], Y[idxs], params[idxs, :]

    def fit_base(self, X, grads):
        models = [self.Base().fit(X, g) for g in grads.T]
        fitted = np.array([m.predict(X) for m in models]).T
        self.base_models.append(models)
        return fitted

    def line_search(self, resids, start, Y, scale_init=1):
        S = self.Score
        D_init = self.Dist(start.T)
        loss_init = S.loss(D_init, Y)
        scale = scale_init
        while True:
            scaled_resids = resids * scale
            D = self.Dist((start - scaled_resids).T)
            loss = S.loss(D, Y)
            norm = np.mean(np.linalg.norm(scaled_resids, axis=1))
            if not np.isnan(loss) and (loss < loss_init or norm < self.tol) and\
               np.linalg.norm(scaled_resids, axis=1).mean() < 5.0:
                break
            scale = scale * 0.5
        self.scalings.append(scale)
        return scale

    def fit(self, X, Y, X_val = None, Y_val = None):

        loss_list = []
        val_loss_list = []
        self.fit_init_params_to_marginal(Y)

        params = self.pred_param(X)
        if X_val is not None and Y_val is not None:
            val_params = self.pred_param(X_val)

        for itr in range(self.n_estimators):
            _, X_batch, Y_batch, P_batch = self.sample(X, Y, params)

            D = self.Dist(P_batch.T)
            S = self.Score

            loss_list += [S.loss(D, Y_batch)]
            loss = loss_list[-1]
            grads = S.natural_grad(D, Y_batch)

            proj_grad = self.fit_base(X_batch, grads)
            scale = self.line_search(proj_grad, P_batch, Y_batch)

            params -= self.learning_rate * scale * np.array([m.predict(X) for m in self.base_models[-1]]).T

            val_loss = 0
            if X_val is not None and Y_val is not None:
                val_params -= self.learning_rate * scale * np.array([m.predict(X_val) for m in self.base_models[-1]]).T
                val_loss = S.loss(self.Dist(val_params.T), Y_val).mean()
                val_loss_list += [val_loss]
                if np.mean(np.array(val_loss_list[-5:])) > \
                   np.mean(np.array(val_loss_list[-10:-5])):
                    if self.verbose:
                        print(f"== Quitting at iteration / VAL {itr} (val_loss={val_loss:.4f})")
                    break

            if self.verbose:
                grad_norm = np.linalg.norm(grads, axis=1).mean() * scale
                print(f"[iter {itr}] loss={loss:.4f} val_loss={val_loss:.4f} scale={scale:.4f} "
                      f"norm={grad_norm:.4f}")

            if np.linalg.norm(proj_grad, axis=1).mean() < self.tol:
                if self.verbose:
                    print(f"== Quitting at iteration / GRAD {itr}")
                break

        return loss_list, val_loss_list

    def fit_init_params_to_marginal(self, Y, iters=1000):
        try:
            E = Y['Event']
            T = Y['Time'].reshape((-1, 1))[E == 1]
        except:
            T = Y
        self.init_params = self.Dist.fit(T)
        return

    def pred_dist(self, X, max_iter=None):
        params = np.asarray(self.pred_param(X, max_iter))
        dist = self.Dist(params.T)
        return dist

    def predict(self, X):
        dist = self.pred_dist(X)
        return list(dist.loc.flatten())

    def staged_predict(self, X, max_iter=None):
        predictions = []
        m, n = X.shape
        params = np.ones((m, self.Dist.n_params)) * self.init_params
        for i, (models, s) in enumerate(zip(self.base_models, self.scalings)):
            if max_iter and i == max_iter:
                break
            resids = np.array([model.predict(X) for model in models]).T
            params -= self.learning_rate * resids * s
            dists = self.Dist(np.asarray(params).T)
            predictions.append(dists.loc.flatten())
        return predictions

    def staged_pred_dist(self, X, max_iter=None):
        predictions = []
        m, n = X.shape
        params = np.ones((m, self.Dist.n_params)) * self.init_params
        for i, (models, s) in enumerate(zip(self.base_models, self.scalings)):
            if max_iter and i == max_iter:
                break
            resids = np.array([model.predict(X) for model in models]).T
            params -= self.learning_rate * resids * s
            dists = self.Dist(np.asarray(params).T)
            predictions.append(dists)
        return predictions
<commit_msg>Add verbose_eval option.<commit_after>import numpy as np
import scipy as sp
import numpy.random as np_rnd

from ngboost.distns import Normal

from ngboost.scores import MLE, CRPS
from ngboost.learners import default_tree_learner, default_linear_learner
from ngboost.distns.normal import Normal

class NGBoost(object):

    def __init__(self, Dist=Normal, Score=MLE(),
                 Base=default_tree_learner, natural_gradient=False,
                 n_estimators=500, learning_rate=0.01, minibatch_frac=1.0,
                 verbose=True, verbose_eval=1, tol=1e-4):
        self.Dist = Dist
        self.Score = Score
        self.Base = Base
        self.natural_gradient = natural_gradient
        self.n_estimators = n_estimators
        self.learning_rate = learning_rate
        self.minibatch_frac = minibatch_frac
        self.verbose = verbose
        self.verbose_eval = verbose_eval
        self.init_params = None
        self.base_models = []
        self.scalings = []
        self.tol = tol

    def pred_param(self, X, max_iter=None):
        m, n = X.shape
        params = np.ones((m, self.Dist.n_params)) * self.init_params
        for i, (models, s) in enumerate(zip(self.base_models, self.scalings)):
            if max_iter and i == max_iter:
                break
            resids = np.array([model.predict(X) for model in models]).T
            params -= self.learning_rate * resids * s
        return params

    def sample(self, X, Y, params):
        sample_size = int(self.minibatch_frac * len(Y))
        idxs = np_rnd.choice(np.arange(len(Y)), sample_size, replace=False)
        return idxs, X[idxs,:], Y[idxs], params[idxs, :]

    def fit_base(self, X, grads):
        models = [self.Base().fit(X, g) for g in grads.T]
        fitted = np.array([m.predict(X) for m in models]).T
        self.base_models.append(models)
        return fitted

    def line_search(self, resids, start, Y, scale_init=1):
        S = self.Score
        D_init = self.Dist(start.T)
        loss_init = S.loss(D_init, Y)
        scale = scale_init
        while True:
            scaled_resids = resids * scale
            D = self.Dist((start - scaled_resids).T)
            loss = S.loss(D, Y)
            norm = np.mean(np.linalg.norm(scaled_resids, axis=1))
            if not np.isnan(loss) and (loss < loss_init or norm < self.tol) and\
               np.linalg.norm(scaled_resids, axis=1).mean() < 5.0:
                break
            scale = scale * 0.5
        self.scalings.append(scale)
        return scale

    def fit(self, X, Y, X_val = None, Y_val = None):

        loss_list = []
        val_loss_list = []
        self.fit_init_params_to_marginal(Y)

        params = self.pred_param(X)
        if X_val is not None and Y_val is not None:
            val_params = self.pred_param(X_val)

        for itr in range(self.n_estimators):
            _, X_batch, Y_batch, P_batch = self.sample(X, Y, params)

            D = self.Dist(P_batch.T)
            S = self.Score

            loss_list += [S.loss(D, Y_batch)]
            loss = loss_list[-1]
            grads = S.natural_grad(D, Y_batch)

            proj_grad = self.fit_base(X_batch, grads)
            scale = self.line_search(proj_grad, P_batch, Y_batch)

            params -= self.learning_rate * scale * np.array([m.predict(X) for m in self.base_models[-1]]).T

            val_loss = 0
            if X_val is not None and Y_val is not None:
                val_params -= self.learning_rate * scale * np.array([m.predict(X_val) for m in self.base_models[-1]]).T
                val_loss = S.loss(self.Dist(val_params.T), Y_val).mean()
                val_loss_list += [val_loss]
                if np.mean(np.array(val_loss_list[-5:])) > \
                   np.mean(np.array(val_loss_list[-10:-5])):
                    if self.verbose:
                        print(f"== Quitting at iteration / VAL {itr} (val_loss={val_loss:.4f})")
                    break

            if self.verbose and int(self.verbose_eval) > 0 and itr % int(self.verbose_eval) == 0:
                grad_norm = np.linalg.norm(grads, axis=1).mean() * scale
                print(f"[iter {itr}] loss={loss:.4f} val_loss={val_loss:.4f} scale={scale:.4f} "
                      f"norm={grad_norm:.4f}")

            if np.linalg.norm(proj_grad, axis=1).mean() < self.tol:
                if self.verbose:
                    print(f"== Quitting at iteration / GRAD {itr}")
                break

        return loss_list, val_loss_list

    def fit_init_params_to_marginal(self, Y, iters=1000):
        try:
            E = Y['Event']
            T = Y['Time'].reshape((-1, 1))[E == 1]
        except:
            T = Y
        self.init_params = self.Dist.fit(T)
        return

    def pred_dist(self, X, max_iter=None):
        params = np.asarray(self.pred_param(X, max_iter))
        dist = self.Dist(params.T)
        return dist

    def predict(self, X):
        dist = self.pred_dist(X)
        return list(dist.loc.flatten())

    def staged_predict(self, X, max_iter=None):
        predictions = []
        m, n = X.shape
        params = np.ones((m, self.Dist.n_params)) * self.init_params
        for i, (models, s) in enumerate(zip(self.base_models, self.scalings)):
            if max_iter and i == max_iter:
                break
            resids = np.array([model.predict(X) for model in models]).T
            params -= self.learning_rate * resids * s
            dists = self.Dist(np.asarray(params).T)
            predictions.append(dists.loc.flatten())
        return predictions

    def staged_pred_dist(self, X, max_iter=None):
        predictions = []
        m, n = X.shape
        params = np.ones((m, self.Dist.n_params)) * self.init_params
        for i, (models, s) in enumerate(zip(self.base_models, self.scalings)):
            if max_iter and i == max_iter:
                break
            resids = np.array([model.predict(X) for model in models]).T
            params -= self.learning_rate * resids * s
            dists = self.Dist(np.asarray(params).T)
            predictions.append(dists)
        return predictions
<|endoftext|>