<commit_before>"""Search module with complex query parsing for Zinnia"""
from django.utils import six

from pyparsing import Word
from pyparsing import alphas
from pyparsing import WordEnd
from pyparsing import Combine
from pyparsing import opAssoc
from pyparsing import Optional
from pyparsing import OneOrMore
from pyparsing import StringEnd
from pyparsing import printables
from pyparsing import quotedString
from pyparsing import removeQuotes
from pyparsing import ParseResults
from pyparsing import CaselessLiteral
from pyparsing import operatorPrecedence

from django.db.models import Q

from zinnia.models.entry import Entry
from zinnia.models.author import Author
from zinnia.settings import STOP_WORDS

def createQ(token):
    """
    Creates the Q() object.
    """
    meta = getattr(token, 'meta', None)
    query = getattr(token, 'query', '')
    wildcards = None

    if isinstance(query, six.string_types):  # Unicode -> Quoted string
        search = query
    else:  # List -> No quoted string (possible wildcards)
        if len(query) == 1:
            search = query[0]
        elif len(query) == 3:
            wildcards = 'BOTH'
            search = query[1]
        elif len(query) == 2:
            if query[0] == '*':
                wildcards = 'START'
                search = query[1]
            else:
                wildcards = 'END'
                search = query[0]

    # Ignore connective words (of, a, an...) and STOP_WORDS
    if (len(search) < 3 and not search.isdigit()) or search in STOP_WORDS:
        return Q()

    if not meta:
        return Q(content__icontains=search) | \
            Q(excerpt__icontains=search) | \
            Q(title__icontains=search)

    if meta == 'category':
        if wildcards == 'BOTH':
            return Q(categories__title__icontains=search) | \
                Q(categories__slug__icontains=search)
        elif wildcards == 'START':
            return Q(categories__title__iendswith=search) | \
                Q(categories__slug__iendswith=search)
        elif wildcards == 'END':
            return Q(categories__title__istartswith=search) | \
                Q(categories__slug__istartswith=search)
        else:
            return Q(categories__title__iexact=search) | \
                Q(categories__slug__iexact=search)
    elif meta == 'author':
        if wildcards == 'BOTH':
            return Q(**{'authors__%s__icontains' % Author.USERNAME_FIELD:
                        search})
        elif wildcards == 'START':
            return Q(**{'authors__%s__iendswith' % Author.USERNAME_FIELD:
                        search})
        elif wildcards == 'END':
            return Q(**{'authors__%s__istartswith' % Author.USERNAME_FIELD:
                        search})
        else:
            return Q(**{'authors__%s__iexact' % Author.USERNAME_FIELD:
                        search})
    elif meta == 'tag':  # TODO: tags ignore wildcards
        return Q(tags__icontains=search)

def unionQ(token):
    """
    Appends all the Q() objects.
    """
    query = Q()
    operation = 'and'
    negation = False

    for t in token:
        if type(t) is ParseResults:  # See tokens recursively
            query &= unionQ(t)
        else:
            if t in ('or', 'and'):  # Set the new op and go to next token
                operation = t
            elif t == '-':  # Next tokens needs to be negated
                negation = True
            else:  # Append to query the token
                if negation:
                    t = ~t
                if operation == 'or':
                    query |= t
                else:
                    query &= t
    return query

NO_BRTS = printables.replace('(', '').replace(')', '')
SINGLE = Word(NO_BRTS.replace('*', ''))
WILDCARDS = Optional('*') + SINGLE + Optional('*') + WordEnd(wordChars=NO_BRTS)
QUOTED = quotedString.setParseAction(removeQuotes)

OPER_AND = CaselessLiteral('and')
OPER_OR = CaselessLiteral('or')
OPER_NOT = '-'

TERM = Combine(Optional(Word(alphas).setResultsName('meta') + ':') +
               (QUOTED.setResultsName('query') |
                WILDCARDS.setResultsName('query')))
TERM.setParseAction(createQ)

EXPRESSION = operatorPrecedence(TERM, [
    (OPER_NOT, 1, opAssoc.RIGHT),
    (OPER_OR, 2, opAssoc.LEFT),
    (Optional(OPER_AND, default='and'), 2, opAssoc.LEFT)])
EXPRESSION.setParseAction(unionQ)

QUERY = OneOrMore(EXPRESSION) + StringEnd()
QUERY.setParseAction(unionQ)

def advanced_search(pattern):
    """
    Parse the grammar of a pattern and build a queryset with it.
    """
    query_parsed = QUERY.parseString(pattern)
    return Entry.published.filter(query_parsed[0]).distinct()
<commit_msg>Handle subtitle and caption in advanced search<commit_after>"""Search module with complex query parsing for Zinnia"""
from django.utils import six

from pyparsing import Word
from pyparsing import alphas
from pyparsing import WordEnd
from pyparsing import Combine
from pyparsing import opAssoc
from pyparsing import Optional
from pyparsing import OneOrMore
from pyparsing import StringEnd
from pyparsing import printables
from pyparsing import quotedString
from pyparsing import removeQuotes
from pyparsing import ParseResults
from pyparsing import CaselessLiteral
from pyparsing import operatorPrecedence

from django.db.models import Q

from zinnia.models.entry import Entry
from zinnia.models.author import Author
from zinnia.settings import STOP_WORDS

def createQ(token):
    """
    Creates the Q() object.
    """
    meta = getattr(token, 'meta', None)
    query = getattr(token, 'query', '')
    wildcards = None

    if isinstance(query, six.string_types):  # Unicode -> Quoted string
        search = query
    else:  # List -> No quoted string (possible wildcards)
        if len(query) == 1:
            search = query[0]
        elif len(query) == 3:
            wildcards = 'BOTH'
            search = query[1]
        elif len(query) == 2:
            if query[0] == '*':
                wildcards = 'START'
                search = query[1]
            else:
                wildcards = 'END'
                search = query[0]

    # Ignore connective words (of, a, an...) and STOP_WORDS
    if (len(search) < 3 and not search.isdigit()) or search in STOP_WORDS:
        return Q()

    if not meta:
        return (Q(title__icontains=search) |
                Q(subtitle__icontains=search) |
                Q(content__icontains=search) |
                Q(excerpt__icontains=search) |
                Q(caption__icontains=search))

    if meta == 'category':
        if wildcards == 'BOTH':
            return (Q(categories__title__icontains=search) |
                    Q(categories__slug__icontains=search))
        elif wildcards == 'START':
            return (Q(categories__title__iendswith=search) |
                    Q(categories__slug__iendswith=search))
        elif wildcards == 'END':
            return (Q(categories__title__istartswith=search) |
                    Q(categories__slug__istartswith=search))
        else:
            return (Q(categories__title__iexact=search) |
                    Q(categories__slug__iexact=search))
    elif meta == 'author':
        if wildcards == 'BOTH':
            return Q(**{'authors__%s__icontains' % Author.USERNAME_FIELD:
                        search})
        elif wildcards == 'START':
            return Q(**{'authors__%s__iendswith' % Author.USERNAME_FIELD:
                        search})
        elif wildcards == 'END':
            return Q(**{'authors__%s__istartswith' % Author.USERNAME_FIELD:
                        search})
        else:
            return Q(**{'authors__%s__iexact' % Author.USERNAME_FIELD:
                        search})
    elif meta == 'tag':  # TODO: tags ignore wildcards
        return Q(tags__icontains=search)

def unionQ(token):
    """
    Appends all the Q() objects.
    """
    query = Q()
    operation = 'and'
    negation = False

    for t in token:
        if type(t) is ParseResults:  # See tokens recursively
            query &= unionQ(t)
        else:
            if t in ('or', 'and'):  # Set the new op and go to next token
                operation = t
            elif t == '-':  # Next tokens needs to be negated
                negation = True
            else:  # Append to query the token
                if negation:
                    t = ~t
                if operation == 'or':
                    query |= t
                else:
                    query &= t
    return query

NO_BRTS = printables.replace('(', '').replace(')', '')
SINGLE = Word(NO_BRTS.replace('*', ''))
WILDCARDS = Optional('*') + SINGLE + Optional('*') + WordEnd(wordChars=NO_BRTS)
QUOTED = quotedString.setParseAction(removeQuotes)

OPER_AND = CaselessLiteral('and')
OPER_OR = CaselessLiteral('or')
OPER_NOT = '-'

TERM = Combine(Optional(Word(alphas).setResultsName('meta') + ':') +
               (QUOTED.setResultsName('query') |
                WILDCARDS.setResultsName('query')))
TERM.setParseAction(createQ)

EXPRESSION = operatorPrecedence(TERM, [
    (OPER_NOT, 1, opAssoc.RIGHT),
    (OPER_OR, 2, opAssoc.LEFT),
    (Optional(OPER_AND, default='and'), 2, opAssoc.LEFT)])
EXPRESSION.setParseAction(unionQ)

QUERY = OneOrMore(EXPRESSION) + StringEnd()
QUERY.setParseAction(unionQ)

def advanced_search(pattern):
    """
    Parse the grammar of a pattern and build a queryset with it.
    """
    query_parsed = QUERY.parseString(pattern)
    return Entry.published.filter(query_parsed[0]).distinct()
<|endoftext|>


<commit_before>import struct

from twisted.internet.protocol import Factory

from zmq3 import Zmq3Protocol, FLAG_MORE

class Zmq1Protocol(Zmq3Protocol):
    header_size = 2
    send_handshake = False
    # http://rfc.zeromq.org/spec:13
    def parseHeader(self, data):
        if not data.startswith(struct.pack('BB', 0x01, 0x00)):
            print "Incorrect server version"
        #print "Server version ok"
        self.proto_state = 2 # no minor header

        # go on if there is any data left
        self.parseFrameData(data, self.header_size)

    def buildGreeting(self):
        # anonymous
        data = struct.pack('BB', 0x01, 0x00)
        return data

    def buildFrame(self, data, more=0):
        data_len = len(data)+1
        if more:
            flags = FLAG_MORE
        else:
            flags = 0

        # headers
        if data_len < 255:
            msg = struct.pack('BB', data_len, flags)
        else:
            msg = struct.pack('BQB', 0xFF, data_len, flags)
        # data
        msg += data
        return msg

    def parseFrameDataChunk(self, data, offset):
        # If offset is too big we're finished parsing
        data_len = len(data)
        if data_len <= offset:
            self._data = data[offset:]
            return
        # Flags and size
        if self.next_part == 0:
            self.size = struct.unpack_from('B', data, offset)[0]-1
            offset += 1
            if self.size == 254:
                self.size = struct.unpack_from('>Q', data, offset)[0]-1
                offset += 8
            flags = struct.unpack_from('B', data, offset)[0]
            offset += 1
            self.more = flags & FLAG_MORE         # only for messages
            self.next_part += 1
            return offset
        # Data
        else:
            if data_len >= offset + self.size:
                self.frameReceived(data[offset:offset+self.size], False, self.more)
                self.next_part = 0
                return offset+self.size
            else:
                # save for later
                self._data = data[offset:]

    def parseFrameData(self, data, offset=0):
        offset = self.parseFrameDataChunk(data, offset)
        while not offset == None:
            offset = self.parseFrameDataChunk(data, offset)

class Zmq1Factory(Factory):
    def __init__(self, type):
        self.type = type
    def buildProtocol(self, addr):
        return Zmq1Protocol(self.type)

<commit_msg>protect from incomplete data<commit_after>import struct

from twisted.internet.protocol import Factory

from zmq3 import Zmq3Protocol, FLAG_MORE

class Zmq1Protocol(Zmq3Protocol):
    header_size = 2
    send_handshake = False
    # http://rfc.zeromq.org/spec:13
    def parseHeader(self, data):
        if not data.startswith(struct.pack('BB', 0x01, 0x00)):
            print "Incorrect server version"
        #print "Server version ok"
        self.proto_state = 2 # no minor header

        # go on if there is any data left
        self.parseFrameData(data, self.header_size)

    def buildGreeting(self):
        # anonymous
        data = struct.pack('BB', 0x01, 0x00)
        return data

    def buildFrame(self, data, more=0):
        data_len = len(data)+1
        if more:
            flags = FLAG_MORE
        else:
            flags = 0

        # headers
        if data_len < 255:
            msg = struct.pack('BB', data_len, flags)
        else:
            msg = struct.pack('BQB', 0xFF, data_len, flags)
        # data
        msg += data
        return msg

    def parseFrameDataChunk(self, data, offset):
        # If offset is too big we're finished parsing
        data_len = len(data)
        if data_len <= offset:
            self._data = data[offset:]
            return
        # Flags and size
        if self.next_part == 0:
            # if not enough data save for later
            if data_len < offset + 3:
                self._data = data[offset:]
                return
            self.size = struct.unpack_from('B', data, offset)[0]-1
            offset += 1
            if self.size == 254:
                if data_len < offset + 9:
                    self._data = data[offset-1:]
                    return
                self.size = struct.unpack_from('>Q', data, offset)[0]-1
                offset += 8
            flags = struct.unpack_from('B', data, offset)[0]
            offset += 1
            self.more = flags & FLAG_MORE         # only for messages
            self.next_part += 1
            return offset
        # Data
        else:
            if data_len >= offset + self.size:
                self.frameReceived(data[offset:offset+self.size], False, self.more)
                self.next_part = 0
                return offset+self.size
            else:
                # save for later
                self._data = data[offset:]

    def parseFrameData(self, data, offset=0):
        offset = self.parseFrameDataChunk(data, offset)
        while not offset == None:
            offset = self.parseFrameDataChunk(data, offset)

class Zmq1Factory(Factory):
    def __init__(self, type):
        self.type = type
    def buildProtocol(self, addr):
        return Zmq1Protocol(self.type)

<|endoftext|>


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

"""Scenario management module.
Reads scenario-related arguments from the command line and uses them to
setup scenarios_to_run(). For each scenario, this generator will yield a
tokenized list of arguments that define that scenario (similar to sys.argv,
but based on a line from a scenario definition file, followed by any options
specified on the command line). The arguments can be retrieved one-by-one
using parse_arg() or (once attached to a model as model.args) by using
get_arg(). These functions use the same syntax as the argparse module.
A queueing system (based on lock files in a queue directory) is used to
ensure that scenarios_to_run() will always return the next unsolved
scenario from the scenario list file, even if the file is edited while this
script is running. This makes it possible to amend the scenario list while
long solver jobs are running. Multiple solvers can also use scenarios_to_run()
in separate processes to select the next job to do.
"""

from __future__ import print_function
import sys, os, time
import argparse, shlex, socket, io, glob
from collections import OrderedDict

from utilities import _ArgumentParser

# load the main model-solver module
import solve

# retrieve base options and command-line arguments
option_file_args = solve.get_option_file_args()
cmd_line_args = sys.argv[1:]

# Parse scenario-manager-related command-line arguments.
# Other command-line arguments will be passed through to solve.py via scenario_cmd_line_args
parser = _ArgumentParser(
    allow_abbrev=False, description='Solve one or more SWITCH scenarios.'
)
parser.add_argument('--scenario', '--scenarios', nargs='+', dest='scenarios', default=[])
#parser.add_argument('--scenarios', nargs='+', default=[])
parser.add_argument("--scenario-list", default="scenarios.txt")
parser.add_argument("--scenario-queue", default="scenario_queue")
parser.add_argument("--job-id", default=None)

#import pdb; pdb.set_trace()
scenario_manager_args = parser.parse_known_args(args=option_file_args + cmd_line_args)[0]
scenario_option_file_args = parser.parse_known_args(args=option_file_args)[1]
scenario_cmd_line_args = parser.parse_known_args(args=cmd_line_args)[1]

requested_scenarios = scenario_manager_args.scenarios
scenario_list_file = scenario_manager_args.scenario_list
scenario_queue_dir = scenario_manager_args.scenario_queue
job_id = scenario_manager_args.job_id

# Note: we could use code like below to collect arguments ad-hoc,
# but then it is difficult to separate the scenario-related arguments from the
# pass-through arguments for solve.py.
# requested_scenarios = (
#     parse_arg("--scenario", nargs='+', default=[])
#     + parse_arg("--scenarios", nargs='+', default=[])
# )
# scenario_list_file = parse_arg("--scenario-list", default="scenarios.txt")
# scenario_queue_dir = parse_arg("--scenario-queue", default="scenario_queue")
#
# job_id = parse_arg("--job-id", default=None)
# scenario_cmd_line_args = sys.argv[1:] # note: this includes scenario-related args, which is bad

# note: we make a best effort to get a unique, persistent job_id for each job.
# this is used to clear the queue of running jobs if a job is stopped and
# restarted. (would be better if other jobs could do this when this job dies
# but it's hard to see how they can detect when this job fails.)
if job_id is None:
    job_id = os.environ.get('JOB_ID') # could be set by user
if job_id is None:
    job_id = os.environ.get('JOBID') # could be set by user
if job_id is None:
    job_id = os.environ.get('SLURM_JOBID')
if job_id is None:
    job_id = os.environ.get('OMPI_MCA_ess_base_jobid')
if job_id is None:
    # construct one
    job_id = socket.gethostname() + '_' + str(os.getpid())

running_scenarios_file = os.path.join(scenario_queue_dir, job_id+"_running.txt")

# list of scenarios currently being run by this job (always just one with the current code)
running_scenarios = []

# make sure the scenario_queue_dir exists (marginally better to do this once
# rather than every time we need to write a file there)
try:
    os.makedirs(scenario_queue_dir)
except OSError:
    pass    # directory probably exists already

#import pdb; pdb.set_trace()

def main():
    # remove lock directories for any scenarios that were
    # previously being solved by this job but were interrupted
    unlock_running_scenarios()

    for (scenario_name, args) in scenarios_to_run():
        print(
            "\n\n=======================================================================\n"
            + "running scenario {s}\n".format(s=scenario_name)
            + "arguments: {}\n".format(args)
            + "=======================================================================\n"
        )

        # call the standard solve module with the arguments for this particular scenario
        solve.main(args=args)

        # another option:
        # subprocess.call(shlex.split("python -m solve") + args) <- omit args from options.txt
        # it should also be possible to use a solver server, but that's not really needed
        # since this script has built-in queue management.

        mark_completed(scenario_name)

def scenarios_to_run():
    """Generator function which returns argument lists for each scenario that should be run.

    Note: each time a new scenario is required, this re-reads the scenario_list file
    and then returns the first scenario that hasn't already started running.
    This allows multiple copies of the script to be run and allocate scenarios among
    themselves."""

    skipped = []
    ran = []

    if requested_scenarios:
        # user requested one or more scenarios
        # just run them in the order specified, with no queue-management
        for scenario_name in requested_scenarios:
            completed = False
            scenario_args = scenario_option_file_args + get_scenario_dict()[scenario_name] + scenario_cmd_line_args
            # flag the scenario as being run; then run it whether or not it was previously run
            checkout(scenario_name, force=True)
            yield (scenario_name, scenario_args)
        # no more scenarios to run
        return
    else:   # no specific scenarios requested
        # Run every scenario in the list, with queue management
        # This is done by repeatedly scanning the scenario list and choosing
        # the first scenario that hasn't been run. This way, users can edit the
        # list and this script will adapt to the changes as soon as it finishes
        # the current scenario.
        all_done = False
        while not all_done:
            all_done = True
            # cache a list of scenarios that have been run, to avoid trying to checkout every one.
            # This list is found by retrieving the names of the lock-directories.
            already_run = filter(os.path.isdir, os.listdir("."))
            for scenario_name, base_args in get_scenario_dict().items():
                if scenario_name not in already_run and checkout(scenario_name):
                    # run this scenario, then start again at the top of the list
                    ran.append(scenario_name)
                    scenario_args = scenario_option_file_args + base_args + scenario_cmd_line_args
                    yield (scenario_name, scenario_args)
                    all_done = False
                    break
                else:
                    if scenario_name not in skipped and scenario_name not in ran:
                        skipped.append(scenario_name)
                        print("Skipping {} because it was already run.".format(scenario_name))
                # move on to the next candidate
        # no more scenarios to run
        if skipped and not ran:
            print(
                "Please remove the {sq} directory or its contents if you would like to "
                "run these scenarios again. (rm -rf {sq})".format(sq=scenario_queue_dir)
            )
        return

def parse_arg(arg, args=sys.argv[1:], **parse_kw):
    """Parse one argument from the argument list, using options as specified for argparse"""
    parser = _ArgumentParser(allow_abbrev=False)
    # Set output destination to 'option', so we can retrieve the value predictably.
    # This is done by updating parse_kw, so it can't be overridden by callers.
    # (They have no reason to set the destination anyway.)
    # note: we use the term "option" so that parsing errors will make a little more
    # sense, e.g., if users call with "--suffixes <blank>" (instead of just omitting it)
    parse_kw["dest"]="option"
    parser.add_argument(arg, **parse_kw)
    return parser.parse_known_args(args)[0].option

def get_scenario_name(scenario_args):
    # use ad-hoc parsing to extract the scenario name from a scenario-definition string
    return parse_arg("--scenario-name", default="", args=scenario_args)

def get_scenario_dict():
    # note: we read the list from the disk each time so that we get a fresher version
    # if the standard list is changed during a long solution effort.
    with open(scenario_list_file, 'r') as f:
        scenario_list_text = f.read()
    # note: text.splitlines() omits newlines and ignores presence/absence of \n at end of the text
    # shlex.split() breaks an command-line-style argument string into a list like sys.argv
    scenario_list = [shlex.split(r) for r in scenario_list_text.splitlines()]
    return OrderedDict((get_scenario_name(s), s) for s in scenario_list)

def checkout(scenario_name, force=False):
    # write a flag that we are solving this scenario, before actually trying to lock it
    # this way, if the job gets interrupted in the middle of this function, the
    # worst that can happen is the scenario will be restarted then next time the job restarts
    # (if we locked the scenario and then got interrupted before setting the flag, then
    # the scenario would not be restarted when the job restarts, which is worse.)
    running_scenarios.append(scenario_name)
    write_running_scenarios_file()
    try:
        # create a lock directory for this scenario
        os.mkdir(os.path.join(scenario_queue_dir, scenario_name))
        locked = True
    except OSError as e:
        if e.errno != 17:     # File exists
            raise
        locked = False
    if locked or force:
        return True
    else:
        # remove the flag that we're running this scenario
        running_scenarios.remove(scenario_name)
        write_running_scenarios_file()
        return False

def mark_completed(scenario_name):
    # remove the scenario from the list of running scenarios (since it's been completed now)
    running_scenarios.remove(scenario_name)
    write_running_scenarios_file()
    # note: the scenario lock directory is left in place so the scenario won't get checked
    # out again

def write_running_scenarios_file():
    # write the list of scenarios currently being run by this job to disk
    # so they can be released back to the queue if the job is interrupted and restarted
    if running_scenarios:
        # note: we use open("r+") and truncate() instead of open("w")
        # to give a better chance of retaining the original file contents if this
        # job is interrupted in the middle of writing the file. This is not an issue unless
        # this job is running multiple scenarios at once
        # (If the file is only partial, then the queue will just think some jobs have been
        # done that actually haven't.)
        flags = "r+" if os.path.exists(running_scenarios_file) else "w"
        with open(running_scenarios_file, flags) as f:
            f.write("\n".join(running_scenarios))
            f.truncate()
    else:
        # remove the running_scenarios_file entirely if it would be empty
        try:
            os.remove(running_scenarios_file)
        except OSError as e:
            if e.errno != 2:    # no such file
                raise

def unlock_running_scenarios():
    # called during startup to remove lockfiles for any scenarios that were still running
    # when this job was interrupted
    if os.path.exists(running_scenarios_file):
        with open(running_scenarios_file) as f:
            interrupted = f.read().splitlines()
        for scenario_name in interrupted:
            try:
                os.rmdir(scenario_name)
            except OSError as e:
                if e.errno != 2:    # no such file
                    raise

# run the main function if called as a script
if __name__ == "__main__":
    main()
<commit_msg>add final newline to running_scenarios file for better display with 'cat'<commit_after>#!/usr/bin/env python

"""Scenario management module.
Reads scenario-related arguments from the command line and uses them to
setup scenarios_to_run(). For each scenario, this generator will yield a
tokenized list of arguments that define that scenario (similar to sys.argv,
but based on a line from a scenario definition file, followed by any options
specified on the command line). The arguments can be retrieved one-by-one
using parse_arg() or (once attached to a model as model.args) by using
get_arg(). These functions use the same syntax as the argparse module.
A queueing system (based on lock files in a queue directory) is used to
ensure that scenarios_to_run() will always return the next unsolved
scenario from the scenario list file, even if the file is edited while this
script is running. This makes it possible to amend the scenario list while
long solver jobs are running. Multiple solvers can also use scenarios_to_run()
in separate processes to select the next job to do.
"""

from __future__ import print_function
import sys, os, time
import argparse, shlex, socket, io, glob
from collections import OrderedDict

from utilities import _ArgumentParser

# load the main model-solver module
import solve

# retrieve base options and command-line arguments
option_file_args = solve.get_option_file_args()
cmd_line_args = sys.argv[1:]

# Parse scenario-manager-related command-line arguments.
# Other command-line arguments will be passed through to solve.py via scenario_cmd_line_args
parser = _ArgumentParser(
    allow_abbrev=False, description='Solve one or more SWITCH scenarios.'
)
parser.add_argument('--scenario', '--scenarios', nargs='+', dest='scenarios', default=[])
#parser.add_argument('--scenarios', nargs='+', default=[])
parser.add_argument("--scenario-list", default="scenarios.txt")
parser.add_argument("--scenario-queue", default="scenario_queue")
parser.add_argument("--job-id", default=None)

#import pdb; pdb.set_trace()
scenario_manager_args = parser.parse_known_args(args=option_file_args + cmd_line_args)[0]
scenario_option_file_args = parser.parse_known_args(args=option_file_args)[1]
scenario_cmd_line_args = parser.parse_known_args(args=cmd_line_args)[1]

requested_scenarios = scenario_manager_args.scenarios
scenario_list_file = scenario_manager_args.scenario_list
scenario_queue_dir = scenario_manager_args.scenario_queue
job_id = scenario_manager_args.job_id

# Note: we could use code like below to collect arguments ad-hoc,
# but then it is difficult to separate the scenario-related arguments from the
# pass-through arguments for solve.py.
# requested_scenarios = (
#     parse_arg("--scenario", nargs='+', default=[])
#     + parse_arg("--scenarios", nargs='+', default=[])
# )
# scenario_list_file = parse_arg("--scenario-list", default="scenarios.txt")
# scenario_queue_dir = parse_arg("--scenario-queue", default="scenario_queue")
#
# job_id = parse_arg("--job-id", default=None)
# scenario_cmd_line_args = sys.argv[1:] # note: this includes scenario-related args, which is bad

# note: we make a best effort to get a unique, persistent job_id for each job.
# this is used to clear the queue of running jobs if a job is stopped and
# restarted. (would be better if other jobs could do this when this job dies
# but it's hard to see how they can detect when this job fails.)
if job_id is None:
    job_id = os.environ.get('JOB_ID') # could be set by user
if job_id is None:
    job_id = os.environ.get('JOBID') # could be set by user
if job_id is None:
    job_id = os.environ.get('SLURM_JOBID')
if job_id is None:
    job_id = os.environ.get('OMPI_MCA_ess_base_jobid')
if job_id is None:
    # construct one
    job_id = socket.gethostname() + '_' + str(os.getpid())

running_scenarios_file = os.path.join(scenario_queue_dir, job_id+"_running.txt")

# list of scenarios currently being run by this job (always just one with the current code)
running_scenarios = []

# make sure the scenario_queue_dir exists (marginally better to do this once
# rather than every time we need to write a file there)
try:
    os.makedirs(scenario_queue_dir)
except OSError:
    pass    # directory probably exists already

#import pdb; pdb.set_trace()

def main():
    # remove lock directories for any scenarios that were
    # previously being solved by this job but were interrupted
    unlock_running_scenarios()

    for (scenario_name, args) in scenarios_to_run():
        print(
            "\n\n=======================================================================\n"
            + "running scenario {s}\n".format(s=scenario_name)
            + "arguments: {}\n".format(args)
            + "=======================================================================\n"
        )

        # call the standard solve module with the arguments for this particular scenario
        solve.main(args=args)

        # another option:
        # subprocess.call(shlex.split("python -m solve") + args) <- omit args from options.txt
        # it should also be possible to use a solver server, but that's not really needed
        # since this script has built-in queue management.

        mark_completed(scenario_name)

def scenarios_to_run():
    """Generator function which returns argument lists for each scenario that should be run.

    Note: each time a new scenario is required, this re-reads the scenario_list file
    and then returns the first scenario that hasn't already started running.
    This allows multiple copies of the script to be run and allocate scenarios among
    themselves."""

    skipped = []
    ran = []

    if requested_scenarios:
        # user requested one or more scenarios
        # just run them in the order specified, with no queue-management
        for scenario_name in requested_scenarios:
            completed = False
            scenario_args = scenario_option_file_args + get_scenario_dict()[scenario_name] + scenario_cmd_line_args
            # flag the scenario as being run; then run it whether or not it was previously run
            checkout(scenario_name, force=True)
            yield (scenario_name, scenario_args)
        # no more scenarios to run
        return
    else:   # no specific scenarios requested
        # Run every scenario in the list, with queue management
        # This is done by repeatedly scanning the scenario list and choosing
        # the first scenario that hasn't been run. This way, users can edit the
        # list and this script will adapt to the changes as soon as it finishes
        # the current scenario.
        all_done = False
        while not all_done:
            all_done = True
            # cache a list of scenarios that have been run, to avoid trying to checkout every one.
            # This list is found by retrieving the names of the lock-directories.
            already_run = filter(os.path.isdir, os.listdir("."))
            for scenario_name, base_args in get_scenario_dict().items():
                if scenario_name not in already_run and checkout(scenario_name):
                    # run this scenario, then start again at the top of the list
                    ran.append(scenario_name)
                    scenario_args = scenario_option_file_args + base_args + scenario_cmd_line_args
                    yield (scenario_name, scenario_args)
                    all_done = False
                    break
                else:
                    if scenario_name not in skipped and scenario_name not in ran:
                        skipped.append(scenario_name)
                        print("Skipping {} because it was already run.".format(scenario_name))
                # move on to the next candidate
        # no more scenarios to run
        if skipped and not ran:
            print(
                "Please remove the {sq} directory or its contents if you would like to "
                "run these scenarios again. (rm -rf {sq})".format(sq=scenario_queue_dir)
            )
        return

def parse_arg(arg, args=sys.argv[1:], **parse_kw):
    """Parse one argument from the argument list, using options as specified for argparse"""
    parser = _ArgumentParser(allow_abbrev=False)
    # Set output destination to 'option', so we can retrieve the value predictably.
    # This is done by updating parse_kw, so it can't be overridden by callers.
    # (They have no reason to set the destination anyway.)
    # note: we use the term "option" so that parsing errors will make a little more
    # sense, e.g., if users call with "--suffixes <blank>" (instead of just omitting it)
    parse_kw["dest"]="option"
    parser.add_argument(arg, **parse_kw)
    return parser.parse_known_args(args)[0].option

def get_scenario_name(scenario_args):
    # use ad-hoc parsing to extract the scenario name from a scenario-definition string
    return parse_arg("--scenario-name", default="", args=scenario_args)

def get_scenario_dict():
    # note: we read the list from the disk each time so that we get a fresher version
    # if the standard list is changed during a long solution effort.
    with open(scenario_list_file, 'r') as f:
        scenario_list_text = f.read()
    # note: text.splitlines() omits newlines and ignores presence/absence of \n at end of the text
    # shlex.split() breaks an command-line-style argument string into a list like sys.argv
    scenario_list = [shlex.split(r) for r in scenario_list_text.splitlines()]
    return OrderedDict((get_scenario_name(s), s) for s in scenario_list)

def checkout(scenario_name, force=False):
    # write a flag that we are solving this scenario, before actually trying to lock it
    # this way, if the job gets interrupted in the middle of this function, the
    # worst that can happen is the scenario will be restarted then next time the job restarts
    # (if we locked the scenario and then got interrupted before setting the flag, then
    # the scenario would not be restarted when the job restarts, which is worse.)
    running_scenarios.append(scenario_name)
    write_running_scenarios_file()
    try:
        # create a lock directory for this scenario
        os.mkdir(os.path.join(scenario_queue_dir, scenario_name))
        locked = True
    except OSError as e:
        if e.errno != 17:     # File exists
            raise
        locked = False
    if locked or force:
        return True
    else:
        # remove the flag that we're running this scenario
        running_scenarios.remove(scenario_name)
        write_running_scenarios_file()
        return False

def mark_completed(scenario_name):
    # remove the scenario from the list of running scenarios (since it's been completed now)
    running_scenarios.remove(scenario_name)
    write_running_scenarios_file()
    # note: the scenario lock directory is left in place so the scenario won't get checked
    # out again

def write_running_scenarios_file():
    # write the list of scenarios currently being run by this job to disk
    # so they can be released back to the queue if the job is interrupted and restarted
    if running_scenarios:
        # note: we use open("r+") and truncate() instead of open("w")
        # to give a better chance of retaining the original file contents if this
        # job is interrupted in the middle of writing the file. This is not an issue unless
        # this job is running multiple scenarios at once
        # (If the file is only partial, then the queue will just think some jobs have been
        # done that actually haven't.)
        flags = "r+" if os.path.exists(running_scenarios_file) else "w"
        with open(running_scenarios_file, flags) as f:
            f.write("\n".join(running_scenarios)+"\n")
            f.truncate()
    else:
        # remove the running_scenarios_file entirely if it would be empty
        try:
            os.remove(running_scenarios_file)
        except OSError as e:
            if e.errno != 2:    # no such file
                raise

def unlock_running_scenarios():
    # called during startup to remove lockfiles for any scenarios that were still running
    # when this job was interrupted
    if os.path.exists(running_scenarios_file):
        with open(running_scenarios_file) as f:
            interrupted = f.read().splitlines()
        for scenario_name in interrupted:
            try:
                os.rmdir(scenario_name)
            except OSError as e:
                if e.errno != 2:    # no such file
                    raise

# run the main function if called as a script
if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>import numpy

import chainer
from chainer.functions.activation import sigmoid
from chainer.functions.activation import tanh
from chainer import link
from chainer.links.connection import linear

class GRUBase(link.Chain):

    def __init__(self, n_units, n_inputs=None, init=None,
                 inner_init=None, bias_init=0):
        if n_inputs is None:
            n_inputs = n_units
        super(GRUBase, self).__init__(
            W_r=linear.Linear(n_inputs, n_units,
                              initialW=init, initial_bias=bias_init),
            U_r=linear.Linear(n_units, n_units,
                              initialW=inner_init, initial_bias=bias_init),
            W_z=linear.Linear(n_inputs, n_units,
                              initialW=init, initial_bias=bias_init),
            U_z=linear.Linear(n_units, n_units,
                              initialW=inner_init, initial_bias=bias_init),
            W=linear.Linear(n_inputs, n_units,
                            initialW=init, initial_bias=bias_init),
            U=linear.Linear(n_units, n_units,
                            initialW=inner_init, initial_bias=bias_init),
        )

class GRU(GRUBase):

    """Stateless Gated Recurrent Unit function (GRU).

    GRU function has six parameters :math:`W_r`, :math:`W_z`, :math:`W`,
    :math:`U_r`, :math:`U_z`, and :math:`U`. All these parameters are
    :math:`n \\times n` matrices, where :math:`n` is the dimension of
    hidden vectors.

    Given two inputs a previous hidden vector :math:`h` and an input vector
    :math:`x`, GRU returns the next hidden vector :math:`h'` defined as

    .. math::

       r &=& \\sigma(W_r x + U_r h), \\\\
       z &=& \\sigma(W_z x + U_z h), \\\\
       \\bar{h} &=& \\tanh(W x + U (r \\odot h)), \\\\
       h' &=& (1 - z) \\odot h + z \\odot \\bar{h},

    where :math:`\\sigma` is the sigmoid function, and :math:`\\odot` is the
    element-wise product.

    :class:`~chainer.links.GRU` does not hold the value of
    hidden vector :math:`h`. So this is *stateless*.
    Use :class:`~chainer.links.StatefulGRU` as a *stateful* GRU.

    Args:
        n_units(int): Dimension of hidden vector :math:`h`.
        n_inputs(int): Dimension of input vector :math:`x`. If ``None``,
            it is set to the same value as ``n_units``.

    See:
        - `On the Properties of Neural Machine Translation: Encoder-Decoder
          Approaches <http://www.aclweb.org/anthology/W14-4012>`_
          [Cho+, SSST2014].
        - `Empirical Evaluation of Gated Recurrent Neural Networks on Sequence
          Modeling <https://arxiv.org/abs/1412.3555>`_
          [Chung+NIPS2014 DLWorkshop].

    .. seealso:: :class:`~chainer.links.StatefulGRU`

    """

    def __call__(self, h, x):
        r = sigmoid.sigmoid(self.W_r(x) + self.U_r(h))
        z = sigmoid.sigmoid(self.W_z(x) + self.U_z(h))
        h_bar = tanh.tanh(self.W(x) + self.U(r * h))
        h_new = (1 - z) * h + z * h_bar
        return h_new

class StatefulGRU(GRUBase):
    """Stateful Gated Recurrent Unit function (GRU).

    Stateful GRU function has six parameters :math:`W_r`, :math:`W_z`,
    :math:`W`, :math:`U_r`, :math:`U_z`, and :math:`U`.
    All these parameters are :math:`n \\times n` matrices,
    where :math:`n` is the dimension of hidden vectors.

    Given input vector :math:`x`, Stateful GRU returns the next
    hidden vector :math:`h'` defined as

    .. math::

       r &=& \\sigma(W_r x + U_r h), \\\\
       z &=& \\sigma(W_z x + U_z h), \\\\
       \\bar{h} &=& \\tanh(W x + U (r \\odot h)), \\\\
       h' &=& (1 - z) \\odot h + z \\odot \\bar{h},

    where :math:`h` is current hidden vector.

    As the name indicates, :class:`~chainer.links.StatefulGRU` is *stateful*,
    meaning that it also holds the next hidden vector `h'` as a state.
    Use :class:`~chainer.links.GRU` as a stateless version of GRU.

    Args:
        in_size(int): Dimension of input vector :math:`x`.
        out_size(int): Dimension of hidden vector :math:`h`.
        init: A callable that takes ``numpy.ndarray`` or
            ``cupy.ndarray`` and edits its value.
            It is used for initialization of the
            GRU's input units (:math:`W`). Maybe be `None` to use default
            initialization.
        inner_init: A callable that takes ``numpy.ndarray`` or
            ``cupy.ndarray`` and edits its value.
            It is used for initialization of the GRU's inner
            recurrent units (:math:`U`).
            Maybe be ``None`` to use default initialization.
        bias_init: A callable or scalar used to initialize the bias values for
            both the GRU's inner and input units. Maybe be ``None`` to use
            default initialization.

    Attributes:
        h(~chainer.Variable): Hidden vector that indicates the state of
            :class:`~chainer.links.StatefulGRU`.

    .. seealso:: :class:`~chainer.functions.GRU`

    """

    def __init__(self, in_size, out_size, init=None,
                 inner_init=None, bias_init=0):
        super(StatefulGRU, self).__init__(
            out_size, in_size, init, inner_init, bias_init)
        self.state_size = out_size
        self.reset_state()

    def to_cpu(self):
        super(StatefulGRU, self).to_cpu()
        if self.h is not None:
            self.h.to_cpu()

    def to_gpu(self, device=None):
        super(StatefulGRU, self).to_gpu(device)
        if self.h is not None:
            self.h.to_gpu(device)

    def set_state(self, h):
        assert isinstance(h, chainer.Variable)
        h_ = h
        if self.xp == numpy:
            h_.to_cpu()
        else:
            h_.to_gpu(self._device_id)
        self.h = h_

    def reset_state(self):
        self.h = None

    def __call__(self, x):
        z = self.W_z(x)
        h_bar = self.W(x)
        if self.h is not None:
            r = sigmoid.sigmoid(self.W_r(x) + self.U_r(self.h))
            z += self.U_z(self.h)
            h_bar += self.U(r * self.h)
        z = sigmoid.sigmoid(z)
        h_bar = tanh.tanh(h_bar)

        h_new = z * h_bar
        if self.h is not None:
            h_new += (1 - z) * self.h
        self.h = h_new
        return self.h
<commit_msg>Use linear-interpolate in GRU<commit_after>import numpy

import chainer
from chainer.functions.activation import sigmoid
from chainer.functions.activation import tanh
from chainer.functions.math import linear_interpolate
from chainer import link
from chainer.links.connection import linear

class GRUBase(link.Chain):

    def __init__(self, n_units, n_inputs=None, init=None,
                 inner_init=None, bias_init=0):
        if n_inputs is None:
            n_inputs = n_units
        super(GRUBase, self).__init__(
            W_r=linear.Linear(n_inputs, n_units,
                              initialW=init, initial_bias=bias_init),
            U_r=linear.Linear(n_units, n_units,
                              initialW=inner_init, initial_bias=bias_init),
            W_z=linear.Linear(n_inputs, n_units,
                              initialW=init, initial_bias=bias_init),
            U_z=linear.Linear(n_units, n_units,
                              initialW=inner_init, initial_bias=bias_init),
            W=linear.Linear(n_inputs, n_units,
                            initialW=init, initial_bias=bias_init),
            U=linear.Linear(n_units, n_units,
                            initialW=inner_init, initial_bias=bias_init),
        )

class GRU(GRUBase):

    """Stateless Gated Recurrent Unit function (GRU).

    GRU function has six parameters :math:`W_r`, :math:`W_z`, :math:`W`,
    :math:`U_r`, :math:`U_z`, and :math:`U`. All these parameters are
    :math:`n \\times n` matrices, where :math:`n` is the dimension of
    hidden vectors.

    Given two inputs a previous hidden vector :math:`h` and an input vector
    :math:`x`, GRU returns the next hidden vector :math:`h'` defined as

    .. math::

       r &=& \\sigma(W_r x + U_r h), \\\\
       z &=& \\sigma(W_z x + U_z h), \\\\
       \\bar{h} &=& \\tanh(W x + U (r \\odot h)), \\\\
       h' &=& (1 - z) \\odot h + z \\odot \\bar{h},

    where :math:`\\sigma` is the sigmoid function, and :math:`\\odot` is the
    element-wise product.

    :class:`~chainer.links.GRU` does not hold the value of
    hidden vector :math:`h`. So this is *stateless*.
    Use :class:`~chainer.links.StatefulGRU` as a *stateful* GRU.

    Args:
        n_units(int): Dimension of hidden vector :math:`h`.
        n_inputs(int): Dimension of input vector :math:`x`. If ``None``,
            it is set to the same value as ``n_units``.

    See:
        - `On the Properties of Neural Machine Translation: Encoder-Decoder
          Approaches <http://www.aclweb.org/anthology/W14-4012>`_
          [Cho+, SSST2014].
        - `Empirical Evaluation of Gated Recurrent Neural Networks on Sequence
          Modeling <https://arxiv.org/abs/1412.3555>`_
          [Chung+NIPS2014 DLWorkshop].

    .. seealso:: :class:`~chainer.links.StatefulGRU`

    """

    def __call__(self, h, x):
        r = sigmoid.sigmoid(self.W_r(x) + self.U_r(h))
        z = sigmoid.sigmoid(self.W_z(x) + self.U_z(h))
        h_bar = tanh.tanh(self.W(x) + self.U(r * h))
        h_new = (1 - z) * h + z * h_bar
        return h_new

class StatefulGRU(GRUBase):
    """Stateful Gated Recurrent Unit function (GRU).

    Stateful GRU function has six parameters :math:`W_r`, :math:`W_z`,
    :math:`W`, :math:`U_r`, :math:`U_z`, and :math:`U`.
    All these parameters are :math:`n \\times n` matrices,
    where :math:`n` is the dimension of hidden vectors.

    Given input vector :math:`x`, Stateful GRU returns the next
    hidden vector :math:`h'` defined as

    .. math::

       r &=& \\sigma(W_r x + U_r h), \\\\
       z &=& \\sigma(W_z x + U_z h), \\\\
       \\bar{h} &=& \\tanh(W x + U (r \\odot h)), \\\\
       h' &=& (1 - z) \\odot h + z \\odot \\bar{h},

    where :math:`h` is current hidden vector.

    As the name indicates, :class:`~chainer.links.StatefulGRU` is *stateful*,
    meaning that it also holds the next hidden vector `h'` as a state.
    Use :class:`~chainer.links.GRU` as a stateless version of GRU.

    Args:
        in_size(int): Dimension of input vector :math:`x`.
        out_size(int): Dimension of hidden vector :math:`h`.
        init: A callable that takes ``numpy.ndarray`` or
            ``cupy.ndarray`` and edits its value.
            It is used for initialization of the
            GRU's input units (:math:`W`). Maybe be `None` to use default
            initialization.
        inner_init: A callable that takes ``numpy.ndarray`` or
            ``cupy.ndarray`` and edits its value.
            It is used for initialization of the GRU's inner
            recurrent units (:math:`U`).
            Maybe be ``None`` to use default initialization.
        bias_init: A callable or scalar used to initialize the bias values for
            both the GRU's inner and input units. Maybe be ``None`` to use
            default initialization.

    Attributes:
        h(~chainer.Variable): Hidden vector that indicates the state of
            :class:`~chainer.links.StatefulGRU`.

    .. seealso:: :class:`~chainer.functions.GRU`

    """

    def __init__(self, in_size, out_size, init=None,
                 inner_init=None, bias_init=0):
        super(StatefulGRU, self).__init__(
            out_size, in_size, init, inner_init, bias_init)
        self.state_size = out_size
        self.reset_state()

    def to_cpu(self):
        super(StatefulGRU, self).to_cpu()
        if self.h is not None:
            self.h.to_cpu()

    def to_gpu(self, device=None):
        super(StatefulGRU, self).to_gpu(device)
        if self.h is not None:
            self.h.to_gpu(device)

    def set_state(self, h):
        assert isinstance(h, chainer.Variable)
        h_ = h
        if self.xp == numpy:
            h_.to_cpu()
        else:
            h_.to_gpu(self._device_id)
        self.h = h_

    def reset_state(self):
        self.h = None

    def __call__(self, x):
        z = self.W_z(x)
        h_bar = self.W(x)
        if self.h is not None:
            r = sigmoid.sigmoid(self.W_r(x) + self.U_r(self.h))
            z += self.U_z(self.h)
            h_bar += self.U(r * self.h)
        z = sigmoid.sigmoid(z)
        h_bar = tanh.tanh(h_bar)

        if self.h is not None:
            h_new = linear_interpolate.linear_interpolate(z, h_bar, self.h)
        else:
            h_new = z * h_bar
        self.h = h_new
        return self.h
<|endoftext|>


<commit_before>from __future__ import absolute_import, division, unicode_literals

from flask.ext.restful import reqparse
from sqlalchemy.orm import contains_eager

from changes.api.base import APIView
from changes.models import Build, TestCase, Job

SORT_CHOICES = (
    'duration',
    'name',
    'retries'
)

class BuildTestIndexAPIView(APIView):
    parser = reqparse.RequestParser()
    parser.add_argument('sort', type=unicode, location='args',
                        choices=SORT_CHOICES, default='duration')

    def get(self, build_id):
        build = Build.query.get(build_id)
        if build is None:
            return '', 404

        args = self.parser.parse_args()

        test_list = TestCase.query.options(
            contains_eager('job')
        ).join(
            Job, TestCase.job_id == Job.id,
        ).filter(
            Job.build_id == build.id,
        )

        if args.sort == 'duration':
            sort_by = TestCase.duration.desc()
        elif args.sort == 'name':
            sort_by = TestCase.name.asc()
        elif args.sort == 'retries':
            sort_by = TestCase.reruns.desc()

        test_list = list(test_list.order_by(sort_by))

        return self.paginate(test_list)
<commit_msg>Dont coerce test results to a list before pagination<commit_after>from __future__ import absolute_import, division, unicode_literals

from flask.ext.restful import reqparse
from sqlalchemy.orm import contains_eager

from changes.api.base import APIView
from changes.models import Build, TestCase, Job

SORT_CHOICES = (
    'duration',
    'name',
    'retries'
)

class BuildTestIndexAPIView(APIView):
    parser = reqparse.RequestParser()
    parser.add_argument('sort', type=unicode, location='args',
                        choices=SORT_CHOICES, default='duration')

    def get(self, build_id):
        build = Build.query.get(build_id)
        if build is None:
            return '', 404

        args = self.parser.parse_args()

        test_list = TestCase.query.options(
            contains_eager('job')
        ).join(
            Job, TestCase.job_id == Job.id,
        ).filter(
            Job.build_id == build.id,
        )

        if args.sort == 'duration':
            sort_by = TestCase.duration.desc()
        elif args.sort == 'name':
            sort_by = TestCase.name.asc()
        elif args.sort == 'retries':
            sort_by = TestCase.reruns.desc()

        test_list = test_list.order_by(sort_by)

        return self.paginate(test_list)
<|endoftext|>


<commit_before>from django.test import TestCase
from promrep.models import Person

class PersonTest(TestCase):

    def setUp(self):
        self.filiations = [
            {'text': None, 'f': None, 'n': None},
            {'text': '', 'f': None, 'n': None},
            {'text': 'ABC', 'f': None, 'n': None},
            {'text': 'A. f.', 'f': 'A.', 'n': None},
            {'text': 'A. n.', 'f': None, 'n': 'A.'},
            {'text': 'A. or B. f.', 'f': 'A. or B.', 'n': None},
            {'text': 'A. or B. n.', 'f': None, 'n': 'A. or B.'},
            {'text': '- f.', 'f': None, 'n': None},
            {'text': '- n.', 'f': None, 'n': None},
            {'text': 'A. f. B. n.', 'f': 'A.', 'n': 'B.'},
            {'text': 'A. or C. f. B. n.', 'f': 'A. or C.', 'n': 'B.'},
            {'text': '- f. B. n.', 'f': None, 'n': 'B.'},
            {'text': 'A. or C. f. B. or D. n.', 'f': 'A. or C.',
             'n': 'B. or D.'},
            {'text': 'A. or C. f. - n.', 'f': 'A. or C.', 'n': None},
            {'text': '- f. - n.', 'f': None, 'n': None},
        ]
        self.person = Person()

    def test__f(self):
        # 1. fail the test
        # self.fail()
        # 2. easy pass of the test
        # self.assertIsNone(self.person.f)
        # 3. add test cases and iteratively implement the functionality to pass
        for filiation in self.filiations:
            self.person.filiation = filiation['text']
            self.assertEqual(filiation['f'], self.person.f)

    def test__n(self):
        for filiation in self.filiations:
            self.person.filiation = filiation['text']
            self.assertEqual(filiation['n'], self.person.n)
<commit_msg>added a simple test to PromrepFacetedSearchForm<commit_after>from django.test import TestCase
from promrep.models import Person

from promrep.forms import PromrepFacetedSearchForm

class PersonTest(TestCase):

    def setUp(self):
        self.filiations = [
            {'text': None, 'f': None, 'n': None},
            {'text': '', 'f': None, 'n': None},
            {'text': 'ABC', 'f': None, 'n': None},
            {'text': 'A. f.', 'f': 'A.', 'n': None},
            {'text': 'A. n.', 'f': None, 'n': 'A.'},
            {'text': 'A. or B. f.', 'f': 'A. or B.', 'n': None},
            {'text': 'A. or B. n.', 'f': None, 'n': 'A. or B.'},
            {'text': '- f.', 'f': None, 'n': None},
            {'text': '- n.', 'f': None, 'n': None},
            {'text': 'A. f. B. n.', 'f': 'A.', 'n': 'B.'},
            {'text': 'A. or C. f. B. n.', 'f': 'A. or C.', 'n': 'B.'},
            {'text': '- f. B. n.', 'f': None, 'n': 'B.'},
            {'text': 'A. or C. f. B. or D. n.', 'f': 'A. or C.',
             'n': 'B. or D.'},
            {'text': 'A. or C. f. - n.', 'f': 'A. or C.', 'n': None},
            {'text': '- f. - n.', 'f': None, 'n': None},
        ]
        self.person = Person()

    def test__f(self):
        # 1. fail the test
        # self.fail()
        # 2. easy pass of the test
        # self.assertIsNone(self.person.f)
        # 3. add test cases and iteratively implement the functionality to pass
        for filiation in self.filiations:
            self.person.filiation = filiation['text']
            self.assertEqual(filiation['f'], self.person.f)

    def test__n(self):
        for filiation in self.filiations:
            self.person.filiation = filiation['text']
            self.assertEqual(filiation['n'], self.person.n)

class PromrepFacetedSearchFormTest(TestCase):

    def setUp(self):
        self.form = PromrepFacetedSearchForm()

    def test_no_query_found(self):
        self.assertIsNotNone(self.form.no_query_found())

    def test_empty_search(self):
        # cannot compare objects
        # naively comparing first element and count of elements

        expected = self.form.no_query_found()
        result = self.form.search()

        self.assertEqual(expected.count(), result.count())
        self.assertEqual(expected[0].id, result[0].id)
<|endoftext|>


<commit_before>
__copyright__ = "Copyright 2017-2018, http://radical.rutgers.edu"
__author__    = "Vivek Balasubramanian <vivek.balasubramanian@rutgers.edu>"
__license__   = "MIT"

import os
import json
import pika
import uuid

import threading     as mt
import radical.utils as ru

from ...exceptions import EnTKError, TypeError

from .resource_manager import Base_ResourceManager

# ------------------------------------------------------------------------------
#
class Base_TaskManager(object):
    """
    A Task Manager takes the responsibility of dispatching tasks it receives
    from a pending_queue for execution on to the available resources using a
    runtime system. Once the tasks have completed execution, they are pushed
    on to the completed_queue for other components of EnTK to process.

    :arguments:
        :pending_queue:     (list) List of queue(s) with tasks ready to be
                            executed. Currently, only one queue.
        :completed_queue:   (list) List of queue(s) with tasks that have
                            finished execution. Currently, only one queue.
        :rmgr:              (ResourceManager) Object to be used to access the
                            Pilot where the tasks can be submitted
        :rmq_conn_params:   (pika.connection.ConnectionParameters) object of
                            parameters necessary to connect to RabbitMQ

    Currently, EnTK is configured to work with one pending queue and one
    completed queue. In the future, the number of queues can be varied for
    different throughput requirements at the cost of additional Memory and CPU
    consumption.
    """

    # --------------------------------------------------------------------------
    #
    def __init__(self, sid, pending_queue, completed_queue, rmgr,
                       rmq_conn_params, rts):

        if not isinstance(sid, str):
            raise TypeError(expected_type=str,
                            actual_type=type(sid))

        if not isinstance(pending_queue, list):
            raise TypeError(expected_type=str,
                            actual_type=type(pending_queue))

        if not isinstance(completed_queue, list):
            raise TypeError(expected_type=str,
                            actual_type=type(completed_queue))

        if not isinstance(rmgr, Base_ResourceManager):
            raise TypeError(expected_type=Base_ResourceManager,
                            actual_type=type(rmgr))

        if not isinstance(rmq_conn_params,
                            pika.connection.ConnectionParameters):
            raise TypeError(expected_type=pika.connection.ConnectionParameters,
                            actual_type=type(rmq_conn_params))

        self._sid             = sid
        self._pending_queue   = pending_queue
        self._completed_queue = completed_queue
        self._rmgr            = rmgr
        self._rts             = rts
        self._rmq_conn_params = rmq_conn_params

        # Utility parameters
        self._uid  = ru.generate_id('task_manager.%(counter)04d', ru.ID_CUSTOM)
        self._path = os.getcwd() + '/' + self._sid

        name = 'radical.entk.%s' % self._uid
        self._log  = ru.Logger  (name, path=self._path)
        self._prof = ru.Profiler(name, path=self._path)
        self._dh = ru.DebugHelper(name=name)

        # Thread should run till terminate condtion is encountered
        mq_connection = pika.BlockingConnection(rmq_conn_params)

        self._hb_request_q  = '%s-hb-request'  % self._sid
        self._hb_response_q = '%s-hb-response' % self._sid

        mq_channel = mq_connection.channel()

        # To respond to heartbeat - get request from rpc_queue
        mq_channel.queue_delete(queue=self._hb_response_q)
        mq_channel.queue_declare(queue=self._hb_response_q)

        # To respond to heartbeat - get request from rpc_queue
        mq_channel.queue_delete(queue=self._hb_request_q)
        mq_channel.queue_declare(queue=self._hb_request_q)

        self._tmgr_process = None
        self._tmgr_terminate = None
        self._hb_thread    = None
        self._hb_terminate = None
        self._hb_interval  = int(os.getenv('ENTK_HB_INTERVAL', '30'))

        mq_connection.close()

    # --------------------------------------------------------------------------
    #
    def _tmgr(self, uid, rmgr, pending_queue, completed_queue,
                    rmq_conn_params):
        """
        **Purpose**: Method to be run by the tmgr process. This method receives
                     a Task from the pending_queue and submits it to the RTS.
                     At all state transititons, they are synced (blocking) with
                     the AppManager in the master process.

                     In addition, the tmgr also receives heartbeat 'request'
                     msgs from the heartbeat-req queue. It responds with
                     a 'response' message to the 'heartbeart-res' queue.

        **Details**: The AppManager can re-invoke the tmgr process with this
                     function if the execution of the workflow is still
                     incomplete. There is also population of a dictionary,
                     placeholder_dict, which stores the path of each of the
                     tasks on the remote machine.
        """

        raise NotImplementedError('_tmgr() method not implemented in '
                                  'TaskManager for %s' % self._rts)

    # --------------------------------------------------------------------------
    #
    def _sync_with_master(self, obj, obj_type, channel, conn_params, queue):

        corr_id = str(uuid.uuid4())
        body    = json.dumps({'object': obj.to_dict(),
                              'type'  : obj_type})

        if   obj_type == 'Task' : msg = obj.parent_stage['uid']
        elif obj_type == 'Stage': msg = obj.parent_pipeline['uid']
        else                    : msg = ''

        self._prof.prof('pub_sync', state=obj.state, uid=obj.uid, msg=msg)
        self._log.debug('%s (%s) to sync with amgr', obj.uid, obj.state)
        try:
            channel.basic_publish(exchange='', routing_key=queue, body=body,
                        properties=pika.BasicProperties(correlation_id=corr_id))
        except (pika.exceptions.ConnectionClosed,
                pika.exceptions.ChannelClosed):
            connection = pika.BlockingConnection(conn_params)
            channel = connection.channel()
            channel.basic_publish(exchange='', routing_key=queue, body=body,
                        properties=pika.BasicProperties(correlation_id=corr_id))

        # all queue name parts up to the last three are used as sid, the last
        # three parts are channel specifiers which need to be inversed to obtain
        # the target channel.
        sid         = '-'.join(queue.split('-')[:-3])
        qname       = queue.split('-')[-3:]
        reply_queue = '-'.join(list(reversed(qname)))
        reply_queue = sid + '-' + reply_queue

        # The `while` loop is diabled with PR #466, and the explanation is:
        #  The task manager and app manager continue to have ack semantics with
        #  this PR. They exchange messages to sync through two channels,
        #  sync-to-tmgr and sync-to-cb, these have not been touched. The ack
        #  that located here is to  send an acknowledgment from the app manager
        #  to the task manager although it doesn't take any further action.
        #  This is redundant and doesn't break or reduce reliability by
        #  commenting out.

        # while True:

        #     # FIXME: is this a busy loop?

        #     method_frame, props, body = channel.basic_get(queue=reply_queue)

        #     if not body:
        #         continue

        #     if corr_id != props.correlation_id:
        #         continue

        #     channel.basic_ack(delivery_tag=method_frame.delivery_tag)

        #     self._prof.prof('sync', state=obj.state, uid=obj.uid, msg=msg)
        #     self._log.debug('%s (%s) synced with amgr', obj.uid, obj.state)

        #     break

    # --------------------------------------------------------------------------
    #
    def _advance(self, obj, obj_type, new_state, channel, conn_params, queue):

        try:
            old_state = obj.state
            obj.state = new_state

            if   obj_type == 'Task' : msg = obj.parent_stage['uid']
            elif obj_type == 'Stage': msg = obj.parent_pipeline['uid']
            else                    : msg = None

            self._prof.prof('advance', uid=obj.uid, state=obj.state, msg=msg)
            self._log.info('Transition %s to %s', obj.uid, new_state)

            self._sync_with_master(obj, obj_type, channel, conn_params, queue)

        except Exception as ex:
            self._log.exception('Transition %s to state %s failed, error: %s',
                                obj.uid, new_state, ex)
            obj.state = old_state
            self._sync_with_master(obj, obj_type, channel, conn_params, queue)
            raise EnTKError(ex) from ex

    # --------------------------------------------------------------------------
    #
    def _heartbeat(self):
        """
        **Purpose**: Method to be executed in the heartbeat thread. This method
                     sends a 'request' to the heartbeat-req queue. It expects
                     a 'response' message from the 'heartbeart-res' queue within
                     10 seconds. This message should contain the same
                     correlation id. If no message if received in 10 seconds,
                     the tmgr is assumed dead. The end_manager() is called to
                     cleanly terminate tmgr process and the heartbeat thread is
                     also terminated.

        **Details**: The AppManager can re-invoke both if the execution is still
                     not complete.
        """

        try:
            self._prof.prof('hbeat_start', uid=self._uid)

            mq_connection = pika.BlockingConnection(self._rmq_conn_params)
            mq_channel = mq_connection.channel()

            while not self._hb_terminate.is_set():

                corr_id  = str(uuid.uuid4())
                try:
                    # Heartbeat request signal sent to task manager via rpc-queue
                    props = pika.BasicProperties(reply_to=self._hb_response_q,
                                                 correlation_id=corr_id)
                    mq_channel.basic_publish(exchange='',
                                             routing_key=self._hb_request_q,
                                             properties=props,
                                             body='request')
                except (pika.exceptions.ConnectionClosed,
                        pika.exceptions.ChannelClosed):
                    mq_connection = pika.BlockingConnection(self._rmq_conn_params)
                    mq_channel = mq_connection.channel()
                    props = pika.BasicProperties(reply_to=self._hb_response_q,
                                                 correlation_id=corr_id)
                    mq_channel.basic_publish(exchange='',
                                             routing_key=self._hb_request_q,
                                             properties=props,
                                             body='request')

                self._log.info('Sent heartbeat request')

                # Sleep for hb_interval and then check if tmgr responded
                mq_connection.sleep(self._hb_interval)

                method_frame, props, body = mq_channel.basic_get(
                                                      queue=self._hb_response_q)
                if not body:
                    # no usable response
                    self._log.error('Heartbeat response no body')
                    return

                if corr_id != props.correlation_id:
                    # incorrect response
                    self._log.error('Heartbeat response wrong correlation')
                    return

                self._log.info('Received heartbeat response')
                mq_channel.basic_ack(delivery_tag=method_frame.delivery_tag)

                # Appease pika cos it thinks the connection is dead
                # mq_connection.close()

        except KeyboardInterrupt as ex:
            self._log.exception('Execution interrupted by user (probably '
                                   ' hit Ctrl+C), cancel tmgr gracefully...')
            raise KeyboardInterrupt from ex

        finally:
            try:
                mq_connection.close()
            except:
                self._log.warning('mq_connection not closed')

            self._prof.prof('hbeat_stop', uid=self._uid)

    # --------------------------------------------------------------------------
    #
    def start_heartbeat(self):
        """
        **Purpose**: Method to start the heartbeat thread. The heartbeat
                     function is not to be accessed directly. The function
                     is started in a separate thread using this method.
        """

        if self._hb_thread:
            self._log.warn('Heartbeat thread already running!')
            return

        try:
            self._log.info('Starting heartbeat thread')
            self._prof.prof('hbeat_create', uid=self._uid)

            self._hb_terminate = mt.Event()
            self._hb_thread    = mt.Thread(target=self._heartbeat,
                                           name='heartbeat')

            self._prof.prof('starting heartbeat thread', uid=self._uid)
            self._hb_thread.start()

            return True

        except Exception as ex:

            self._log.exception('Heartbeat not started, error: %s', ex)
            self.terminate_heartbeat()
            raise EnTKError(ex) from ex

    # --------------------------------------------------------------------------
    #
    def terminate_heartbeat(self):
        """
        **Purpose**: Method to terminate the heartbeat thread. This method is
        blocking as it waits for the heartbeat thread to terminate (aka join).

        This is the last method that is executed from the TaskManager and
        hence closes the profiler.
        """

        try:

            if self._hb_thread:

                self._hb_terminate.set()

                if self.check_heartbeat():
                    self._hb_thread.join()

                self._hb_thread = None

                self._log.info('Hearbeat thread terminated')
                self._prof.prof('hbeat_term', uid=self._uid)

        except Exception as ex:
            self._log.exception('Could not terminate heartbeat thread')
            raise EnTKError(ex) from ex

        finally:

            if not self.check_heartbeat() or self.check_manager():

                conn = pika.BlockingConnection(self._rmq_conn_params)
                mq_channel = conn.channel()

                # To respond to heartbeat - get request from rpc_queue
                mq_channel.queue_delete(queue=self._hb_response_q)
                mq_channel.queue_delete(queue=self._hb_request_q)

                conn.close()

    # --------------------------------------------------------------------------
    #
    def start_manager(self):
        """
        **Purpose**: Method to start the tmgr process. The tmgr function
        is not to be accessed directly. The function is started in a separate
        thread using this method.
        """

        raise NotImplementedError('start_manager() method not implemented in '
                                  'TaskManager for %s' % self._rts)

    # --------------------------------------------------------------------------
    #
    def terminate_manager(self):
        """
        **Purpose**: Method to terminate the tmgr process. This method is
        blocking as it waits for the tmgr process to terminate (aka join).
        """

        try:
            if self._tmgr_process:
                self._log.debug('Trying to terminate task manager.')
                if self._tmgr_terminate is not None:
                    if not self._tmgr_terminate.is_set():
                        self._tmgr_terminate.set()
                    self._log.debug('TMGR terminate is set %s' % self._tmgr_terminate.is_set())
                if self.check_manager():
                    self._log.debug('TMGR process is alive')
                    self._tmgr_process.join(30)
                self._log.debug('TMGR process joined')
                self._tmgr_process = None

                self._log.info('Task manager process closed')
                self._prof.prof('tmgr_term', uid=self._uid)

        except Exception as ex:
            self._log.exception('Could not terminate task manager process')
            raise EnTKError(ex) from ex

    # --------------------------------------------------------------------------
    #
    def check_heartbeat(self):
        """
        **Purpose**: Check if the heartbeat thread is alive and running
        """

        if self._hb_thread:
            return self._hb_thread.is_alive()
        else:
            return False

    # --------------------------------------------------------------------------
    #
    def check_manager(self):
        """
        **Purpose**: Check if the tmgr process is alive and running
        """

        # Return False if the process does not exist
        if self._tmgr_process:
            return self._tmgr_process.is_alive()
        else:
            return False

# ------------------------------------------------------------------------------

<commit_msg>Restoring wrong del<commit_after>
__copyright__ = "Copyright 2017-2018, http://radical.rutgers.edu"
__author__    = "Vivek Balasubramanian <vivek.balasubramanian@rutgers.edu>"
__license__   = "MIT"

import os
import json
import pika
import uuid

import threading     as mt
import radical.utils as ru

from ...exceptions import EnTKError, TypeError

from .resource_manager import Base_ResourceManager

# ------------------------------------------------------------------------------
#
class Base_TaskManager(object):
    """
    A Task Manager takes the responsibility of dispatching tasks it receives
    from a pending_queue for execution on to the available resources using a
    runtime system. Once the tasks have completed execution, they are pushed
    on to the completed_queue for other components of EnTK to process.

    :arguments:
        :pending_queue:     (list) List of queue(s) with tasks ready to be
                            executed. Currently, only one queue.
        :completed_queue:   (list) List of queue(s) with tasks that have
                            finished execution. Currently, only one queue.
        :rmgr:              (ResourceManager) Object to be used to access the
                            Pilot where the tasks can be submitted
        :rmq_conn_params:   (pika.connection.ConnectionParameters) object of
                            parameters necessary to connect to RabbitMQ

    Currently, EnTK is configured to work with one pending queue and one
    completed queue. In the future, the number of queues can be varied for
    different throughput requirements at the cost of additional Memory and CPU
    consumption.
    """

    # --------------------------------------------------------------------------
    #
    def __init__(self, sid, pending_queue, completed_queue, rmgr,
                       rmq_conn_params, rts):

        if not isinstance(sid, str):
            raise TypeError(expected_type=str,
                            actual_type=type(sid))

        if not isinstance(pending_queue, list):
            raise TypeError(expected_type=str,
                            actual_type=type(pending_queue))

        if not isinstance(completed_queue, list):
            raise TypeError(expected_type=str,
                            actual_type=type(completed_queue))

        if not isinstance(rmgr, Base_ResourceManager):
            raise TypeError(expected_type=Base_ResourceManager,
                            actual_type=type(rmgr))

        if not isinstance(rmq_conn_params,
                            pika.connection.ConnectionParameters):
            raise TypeError(expected_type=pika.connection.ConnectionParameters,
                            actual_type=type(rmq_conn_params))

        self._sid             = sid
        self._pending_queue   = pending_queue
        self._completed_queue = completed_queue
        self._rmgr            = rmgr
        self._rts             = rts
        self._rmq_conn_params = rmq_conn_params

        # Utility parameters
        self._uid  = ru.generate_id('task_manager.%(counter)04d', ru.ID_CUSTOM)
        self._path = os.getcwd() + '/' + self._sid

        name = 'radical.entk.%s' % self._uid
        self._log  = ru.Logger  (name, path=self._path)
        self._prof = ru.Profiler(name, path=self._path)
        self._dh = ru.DebugHelper(name=name)

        # Thread should run till terminate condtion is encountered
        mq_connection = pika.BlockingConnection(rmq_conn_params)

        self._hb_request_q  = '%s-hb-request'  % self._sid
        self._hb_response_q = '%s-hb-response' % self._sid

        mq_channel = mq_connection.channel()

        # To respond to heartbeat - get request from rpc_queue
        mq_channel.queue_delete(queue=self._hb_response_q)
        mq_channel.queue_declare(queue=self._hb_response_q)

        # To respond to heartbeat - get request from rpc_queue
        mq_channel.queue_delete(queue=self._hb_request_q)
        mq_channel.queue_declare(queue=self._hb_request_q)

        self._tmgr_process = None
        self._tmgr_terminate = None
        self._hb_thread    = None
        self._hb_terminate = None
        self._hb_interval  = int(os.getenv('ENTK_HB_INTERVAL', '30'))

        mq_connection.close()

    # --------------------------------------------------------------------------
    #
    def _tmgr(self, uid, rmgr, pending_queue, completed_queue,
                    rmq_conn_params):
        """
        **Purpose**: Method to be run by the tmgr process. This method receives
                     a Task from the pending_queue and submits it to the RTS.
                     At all state transititons, they are synced (blocking) with
                     the AppManager in the master process.

                     In addition, the tmgr also receives heartbeat 'request'
                     msgs from the heartbeat-req queue. It responds with
                     a 'response' message to the 'heartbeart-res' queue.

        **Details**: The AppManager can re-invoke the tmgr process with this
                     function if the execution of the workflow is still
                     incomplete. There is also population of a dictionary,
                     placeholder_dict, which stores the path of each of the
                     tasks on the remote machine.
        """

        raise NotImplementedError('_tmgr() method not implemented in '
                                  'TaskManager for %s' % self._rts)

    # --------------------------------------------------------------------------
    #
    def _sync_with_master(self, obj, obj_type, channel, conn_params, queue):

        corr_id = str(uuid.uuid4())
        body    = json.dumps({'object': obj.to_dict(),
                              'type'  : obj_type})

        if   obj_type == 'Task' : msg = obj.parent_stage['uid']
        elif obj_type == 'Stage': msg = obj.parent_pipeline['uid']
        else                    : msg = ''

        self._prof.prof('pub_sync', state=obj.state, uid=obj.uid, msg=msg)
        self._log.debug('%s (%s) to sync with amgr', obj.uid, obj.state)
        try:
            channel.basic_publish(exchange='', routing_key=queue, body=body,
                        properties=pika.BasicProperties(correlation_id=corr_id))
        except (pika.exceptions.ConnectionClosed,
                pika.exceptions.ChannelClosed):
            connection = pika.BlockingConnection(conn_params)
            channel = connection.channel()
            channel.basic_publish(exchange='', routing_key=queue, body=body,
                        properties=pika.BasicProperties(correlation_id=corr_id))

        # all queue name parts up to the last three are used as sid, the last
        # three parts are channel specifiers which need to be inversed to obtain
        # the target channel.
        sid         = '-'.join(queue.split('-')[:-3])
        qname       = queue.split('-')[-3:]
        reply_queue = '-'.join(list(reversed(qname)))
        reply_queue = sid + '-' + reply_queue

        # The `while` loop is diabled with PR #466, and the explanation is:
        #  The task manager and app manager continue to have ack semantics with
        #  this PR. They exchange messages to sync through two channels,
        #  sync-to-tmgr and sync-to-cb, these have not been touched. The ack
        #  that located here is to  send an acknowledgment from the app manager
        #  to the task manager although it doesn't take any further action.
        #  This is redundant and doesn't break or reduce reliability by
        #  commenting out.

        # while True:

        #     # FIXME: is this a busy loop?

        #     method_frame, props, body = channel.basic_get(queue=reply_queue)

        #     if not body:
        #         continue

        #     if corr_id != props.correlation_id:
        #         continue

        #     channel.basic_ack(delivery_tag=method_frame.delivery_tag)

        #     self._prof.prof('sync', state=obj.state, uid=obj.uid, msg=msg)
        #     self._log.debug('%s (%s) synced with amgr', obj.uid, obj.state)

        #     break

    # --------------------------------------------------------------------------
    #
    def _advance(self, obj, obj_type, new_state, channel, conn_params, queue):

        try:
            old_state = obj.state
            obj.state = new_state

            if   obj_type == 'Task' : msg = obj.parent_stage['uid']
            elif obj_type == 'Stage': msg = obj.parent_pipeline['uid']
            else                    : msg = None

            self._prof.prof('advance', uid=obj.uid, state=obj.state, msg=msg)
            self._log.info('Transition %s to %s', obj.uid, new_state)

            self._sync_with_master(obj, obj_type, channel, conn_params, queue)

        except Exception as ex:
            self._log.exception('Transition %s to state %s failed, error: %s',
                                obj.uid, new_state, ex)
            obj.state = old_state
            self._sync_with_master(obj, obj_type, channel, conn_params, queue)
            raise EnTKError(ex) from ex

    # --------------------------------------------------------------------------
    #
    def _heartbeat(self):
        """
        **Purpose**: Method to be executed in the heartbeat thread. This method
                     sends a 'request' to the heartbeat-req queue. It expects
                     a 'response' message from the 'heartbeart-res' queue within
                     10 seconds. This message should contain the same
                     correlation id. If no message if received in 10 seconds,
                     the tmgr is assumed dead. The end_manager() is called to
                     cleanly terminate tmgr process and the heartbeat thread is
                     also terminated.

        **Details**: The AppManager can re-invoke both if the execution is still
                     not complete.
        """

        try:
            self._prof.prof('hbeat_start', uid=self._uid)

            mq_connection = pika.BlockingConnection(self._rmq_conn_params)
            mq_channel = mq_connection.channel()

            while not self._hb_terminate.is_set():

                corr_id  = str(uuid.uuid4())
                try:
                    # Heartbeat request signal sent to task manager via rpc-queue
                    props = pika.BasicProperties(reply_to=self._hb_response_q,
                                                 correlation_id=corr_id)
                    mq_channel.basic_publish(exchange='',
                                             routing_key=self._hb_request_q,
                                             properties=props,
                                             body='request')
                except (pika.exceptions.ConnectionClosed,
                        pika.exceptions.ChannelClosed):
                    mq_connection = pika.BlockingConnection(self._rmq_conn_params)
                    mq_channel = mq_connection.channel()
                    props = pika.BasicProperties(reply_to=self._hb_response_q,
                                                 correlation_id=corr_id)
                    mq_channel.basic_publish(exchange='',
                                             routing_key=self._hb_request_q,
                                             properties=props,
                                             body='request')

                self._log.info('Sent heartbeat request')

                # Sleep for hb_interval and then check if tmgr responded
                mq_connection.sleep(self._hb_interval)

                method_frame, props, body = mq_channel.basic_get(
                                                      queue=self._hb_response_q)
                if not body:
                    # no usable response
                    self._log.error('Heartbeat response no body')
                    return

                if corr_id != props.correlation_id:
                    # incorrect response
                    self._log.error('Heartbeat response wrong correlation')
                    return

                self._log.info('Received heartbeat response')
                mq_channel.basic_ack(delivery_tag=method_frame.delivery_tag)

                # Appease pika cos it thinks the connection is dead
                # mq_connection.close()

        except KeyboardInterrupt as ex:
            self._log.exception('Execution interrupted by user (probably '
                                   ' hit Ctrl+C), cancel tmgr gracefully...')
            raise KeyboardInterrupt from ex
        except Exception as ex:
            self._log.exception('Heartbeat failed with error: %s', ex)
            raise EnTKError(ex) from ex
        finally:
            try:
                mq_connection.close()
            except:
                self._log.warning('mq_connection not closed')

            self._prof.prof('hbeat_stop', uid=self._uid)

    # --------------------------------------------------------------------------
    #
    def start_heartbeat(self):
        """
        **Purpose**: Method to start the heartbeat thread. The heartbeat
                     function is not to be accessed directly. The function
                     is started in a separate thread using this method.
        """

        if self._hb_thread:
            self._log.warn('Heartbeat thread already running!')
            return

        try:
            self._log.info('Starting heartbeat thread')
            self._prof.prof('hbeat_create', uid=self._uid)

            self._hb_terminate = mt.Event()
            self._hb_thread    = mt.Thread(target=self._heartbeat,
                                           name='heartbeat')

            self._prof.prof('starting heartbeat thread', uid=self._uid)
            self._hb_thread.start()

            return True

        except Exception as ex:

            self._log.exception('Heartbeat not started, error: %s', ex)
            self.terminate_heartbeat()
            raise EnTKError(ex) from ex

    # --------------------------------------------------------------------------
    #
    def terminate_heartbeat(self):
        """
        **Purpose**: Method to terminate the heartbeat thread. This method is
        blocking as it waits for the heartbeat thread to terminate (aka join).

        This is the last method that is executed from the TaskManager and
        hence closes the profiler.
        """

        try:

            if self._hb_thread:

                self._hb_terminate.set()

                if self.check_heartbeat():
                    self._hb_thread.join()

                self._hb_thread = None

                self._log.info('Hearbeat thread terminated')
                self._prof.prof('hbeat_term', uid=self._uid)

        except Exception as ex:
            self._log.exception('Could not terminate heartbeat thread')
            raise EnTKError(ex) from ex

        finally:

            if not self.check_heartbeat() or self.check_manager():

                conn = pika.BlockingConnection(self._rmq_conn_params)
                mq_channel = conn.channel()

                # To respond to heartbeat - get request from rpc_queue
                mq_channel.queue_delete(queue=self._hb_response_q)
                mq_channel.queue_delete(queue=self._hb_request_q)

                conn.close()

    # --------------------------------------------------------------------------
    #
    def start_manager(self):
        """
        **Purpose**: Method to start the tmgr process. The tmgr function
        is not to be accessed directly. The function is started in a separate
        thread using this method.
        """

        raise NotImplementedError('start_manager() method not implemented in '
                                  'TaskManager for %s' % self._rts)

    # --------------------------------------------------------------------------
    #
    def terminate_manager(self):
        """
        **Purpose**: Method to terminate the tmgr process. This method is
        blocking as it waits for the tmgr process to terminate (aka join).
        """

        try:
            if self._tmgr_process:
                self._log.debug('Trying to terminate task manager.')
                if self._tmgr_terminate is not None:
                    if not self._tmgr_terminate.is_set():
                        self._tmgr_terminate.set()
                    self._log.debug('TMGR terminate is set %s' % self._tmgr_terminate.is_set())
                if self.check_manager():
                    self._log.debug('TMGR process is alive')
                    self._tmgr_process.join(30)
                self._log.debug('TMGR process joined')
                self._tmgr_process = None

                self._log.info('Task manager process closed')
                self._prof.prof('tmgr_term', uid=self._uid)

        except Exception as ex:
            self._log.exception('Could not terminate task manager process')
            raise EnTKError(ex) from ex

    # --------------------------------------------------------------------------
    #
    def check_heartbeat(self):
        """
        **Purpose**: Check if the heartbeat thread is alive and running
        """

        if self._hb_thread:
            return self._hb_thread.is_alive()
        else:
            return False

    # --------------------------------------------------------------------------
    #
    def check_manager(self):
        """
        **Purpose**: Check if the tmgr process is alive and running
        """

        # Return False if the process does not exist
        if self._tmgr_process:
            return self._tmgr_process.is_alive()
        else:
            return False

# ------------------------------------------------------------------------------

<|endoftext|>