<commit_before># Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of the Neural Programmer model described in https://openreview.net/pdf?id=ry2YOrcge

This file calls functions to load & pre-process data, construct the TF graph
and performs training or evaluation as specified by the flag evaluator_job
Author: aneelakantan (Arvind Neelakantan)
"""
import time
from random import Random
import numpy as np
import tensorflow as tf
import model
import wiki_data
import parameters
import data_utils

tf.flags.DEFINE_integer("train_steps", 100001, "Number of steps to train")
tf.flags.DEFINE_integer("eval_cycle", 500,
                        "Evaluate model at every eval_cycle steps")
tf.flags.DEFINE_integer("max_elements", 100,
                        "maximum rows that are  considered for processing")
tf.flags.DEFINE_integer(
    "max_number_cols", 15,
    "maximum number columns that are considered for processing")
tf.flags.DEFINE_integer(
    "max_word_cols", 25,
    "maximum number columns that are considered for processing")
tf.flags.DEFINE_integer("question_length", 62, "maximum question length")
tf.flags.DEFINE_integer("max_entry_length", 1, "")
tf.flags.DEFINE_integer("max_passes", 4, "number of operation passes")
tf.flags.DEFINE_integer("embedding_dims", 256, "")
tf.flags.DEFINE_integer("batch_size", 20, "")
tf.flags.DEFINE_float("clip_gradients", 1.0, "")
tf.flags.DEFINE_float("eps", 1e-6, "")
tf.flags.DEFINE_float("param_init", 0.1, "")
tf.flags.DEFINE_float("learning_rate", 0.001, "")
tf.flags.DEFINE_float("l2_regularizer", 0.0001, "")
tf.flags.DEFINE_float("print_cost", 50.0,
                      "weighting factor in the objective function")
tf.flags.DEFINE_string("job_id", "temp", """job id""")
tf.flags.DEFINE_string("output_dir", "../model/",
                       """output_dir""")
tf.flags.DEFINE_string("data_dir", "../data/",
                       """data_dir""")
tf.flags.DEFINE_integer("write_every", 500, "wrtie every N")
tf.flags.DEFINE_integer("param_seed", 150, "")
tf.flags.DEFINE_integer("python_seed", 200, "")
tf.flags.DEFINE_float("dropout", 0.8, "dropout keep probability")
tf.flags.DEFINE_float("rnn_dropout", 0.9,
                      "dropout keep probability for rnn connections")
tf.flags.DEFINE_float("pad_int", -20000.0,
                      "number columns are padded with pad_int")
tf.flags.DEFINE_string("data_type", "double", "float or double")
tf.flags.DEFINE_float("word_dropout_prob", 0.9, "word dropout keep prob")
tf.flags.DEFINE_integer("word_cutoff", 10, "")
tf.flags.DEFINE_integer("vocab_size", 10800, "")
tf.flags.DEFINE_boolean("evaluator_job", False,
                        "wehther to run as trainer/evaluator")
tf.flags.DEFINE_float(
    "bad_number_pre_process", -200000.0,
    "number that is added to a corrupted table entry in a number column")
tf.flags.DEFINE_float("max_math_error", 3.0,
                      "max square loss error that is considered")
tf.flags.DEFINE_float("soft_min_value", 5.0, "")
FLAGS = tf.flags.FLAGS

class Utility:
  #holds FLAGS and other variables that are used in different files
  def __init__(self):
    global FLAGS
    self.FLAGS = FLAGS
    self.unk_token = "UNK"
    self.entry_match_token = "entry_match"
    self.column_match_token = "column_match"
    self.dummy_token = "dummy_token"
    self.tf_data_type = {}
    self.tf_data_type["double"] = tf.float64
    self.tf_data_type["float"] = tf.float32
    self.np_data_type = {}
    self.np_data_type["double"] = np.float64
    self.np_data_type["float"] = np.float32
    self.operations_set = ["count"] + [
        "prev", "next", "first_rs", "last_rs", "group_by_max", "greater",
        "lesser", "geq", "leq", "max", "min", "word-match"
    ] + ["reset_select"] + ["print"]
    self.word_ids = {}
    self.reverse_word_ids = {}
    self.word_count = {}
    self.random = Random(FLAGS.python_seed)

def evaluate(sess, data, batch_size, graph, i):
  #computes accuracy
  num_examples = 0.0
  gc = 0.0
  for j in range(0, len(data) - batch_size + 1, batch_size):
    [ct] = sess.run([graph.final_correct],
                    feed_dict=data_utils.generate_feed_dict(data, j, batch_size,
                                                            graph))
    gc += ct * batch_size
    num_examples += batch_size
  print "dev set accuracy   after ", i, " : ", gc / num_examples
  print num_examples, len(data)
  print "--------"

def Train(graph, utility, batch_size, train_data, sess, model_dir,
          saver):
  #performs training
  curr = 0
  train_set_loss = 0.0
  utility.random.shuffle(train_data)
  start = time.time()
  for i in range(utility.FLAGS.train_steps):
    curr_step = i
    if (i > 0 and i % FLAGS.write_every == 0):
      model_file = model_dir + "/model_" + str(i)
      saver.save(sess, model_file)
    if curr + batch_size >= len(train_data):
      curr = 0
      utility.random.shuffle(train_data)
    step, cost_value = sess.run(
        [graph.step, graph.total_cost],
        feed_dict=data_utils.generate_feed_dict(
            train_data, curr, batch_size, graph, train=True, utility=utility))
    curr = curr + batch_size
    train_set_loss += cost_value
    if (i > 0 and i % FLAGS.eval_cycle == 0):
      end = time.time()
      time_taken = end - start
      print "step ", i, " ", time_taken, " seconds "
      start = end
      print " printing train set loss: ", train_set_loss / utility.FLAGS.eval_cycle
      train_set_loss = 0.0

def master(train_data, dev_data, utility):
  #creates TF graph and calls trainer or evaluator
  batch_size = utility.FLAGS.batch_size
  model_dir = utility.FLAGS.output_dir + "/model" + utility.FLAGS.job_id + "/"
  #create all paramters of the model
  param_class = parameters.Parameters(utility)
  params, global_step, init = param_class.parameters(utility)
  key = "test" if (FLAGS.evaluator_job) else "train"
  graph = model.Graph(utility, batch_size, utility.FLAGS.max_passes, mode=key)
  graph.create_graph(params, global_step)
  prev_dev_error = 0.0
  final_loss = 0.0
  final_accuracy = 0.0
  #start session
  with tf.Session() as sess:
    sess.run(init.name)
    sess.run(graph.init_op.name)
    to_save = params.copy()
    saver = tf.train.Saver(to_save, max_to_keep=500)
    if (FLAGS.evaluator_job):
      while True:
        selected_models = {}
        file_list = tf.gfile.ListDirectory(model_dir)
        for model_file in file_list:
          if ("checkpoint" in model_file or "index" in model_file or
              "meta" in model_file):
            continue
          if ("data" in model_file):
            model_file = model_file.split(".")[0]
          model_step = int(
              model_file.split("_")[len(model_file.split("_")) - 1])
          selected_models[model_step] = model_file
        file_list = sorted(selected_models.items(), key=lambda x: x[0])
        if (len(file_list) > 0):
          file_list = file_list[0:len(file_list) - 1]
	print "list of models: ", file_list
        for model_file in file_list:
          model_file = model_file[1]
          print "restoring: ", model_file
          saver.restore(sess, model_dir + "/" + model_file)
          model_step = int(
              model_file.split("_")[len(model_file.split("_")) - 1])
          print "evaluating on dev ", model_file, model_step
          evaluate(sess, dev_data, batch_size, graph, model_step)
    else:
      ckpt = tf.train.get_checkpoint_state(model_dir)
      print "model dir: ", model_dir
      if (not (tf.gfile.IsDirectory(model_dir))):
        print "create dir: ", model_dir
        tf.gfile.MkDir(model_dir)
      Train(graph, utility, batch_size, train_data, sess, model_dir,
            saver)

def main(args):
  utility = Utility()
  train_name = "random-split-1-train.examples"
  dev_name = "random-split-1-dev.examples"
  test_name = "pristine-unseen-tables.examples"
  #load data
  dat = wiki_data.WikiQuestionGenerator(train_name, dev_name, test_name, FLAGS.data_dir)
  train_data, dev_data, test_data = dat.load()
  utility.words = []
  utility.word_ids = {}
  utility.reverse_word_ids = {}
  #construct vocabulary
  data_utils.construct_vocab(train_data, utility)
  data_utils.construct_vocab(dev_data, utility, True)
  data_utils.construct_vocab(test_data, utility, True)
  data_utils.add_special_words(utility)
  data_utils.perform_word_cutoff(utility)
  #convert data to int format and pad the inputs
  train_data = data_utils.complete_wiki_processing(train_data, utility, True)
  dev_data = data_utils.complete_wiki_processing(dev_data, utility, False)
  test_data = data_utils.complete_wiki_processing(test_data, utility, False)
  print "# train examples ", len(train_data)
  print "# dev examples ", len(dev_data)
  print "# test examples ", len(test_data)
  print "running open source"
  #construct TF graph and train or evaluate
  master(train_data, dev_data, utility)

if __name__ == "__main__":
  tf.app.run()
<commit_msg>ensure output directory exists<commit_after># Copyright 2016 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Implementation of the Neural Programmer model described in https://openreview.net/pdf?id=ry2YOrcge

This file calls functions to load & pre-process data, construct the TF graph
and performs training or evaluation as specified by the flag evaluator_job
Author: aneelakantan (Arvind Neelakantan)
"""
import time
from random import Random
import numpy as np
import tensorflow as tf
import model
import wiki_data
import parameters
import data_utils

tf.flags.DEFINE_integer("train_steps", 100001, "Number of steps to train")
tf.flags.DEFINE_integer("eval_cycle", 500,
                        "Evaluate model at every eval_cycle steps")
tf.flags.DEFINE_integer("max_elements", 100,
                        "maximum rows that are  considered for processing")
tf.flags.DEFINE_integer(
    "max_number_cols", 15,
    "maximum number columns that are considered for processing")
tf.flags.DEFINE_integer(
    "max_word_cols", 25,
    "maximum number columns that are considered for processing")
tf.flags.DEFINE_integer("question_length", 62, "maximum question length")
tf.flags.DEFINE_integer("max_entry_length", 1, "")
tf.flags.DEFINE_integer("max_passes", 4, "number of operation passes")
tf.flags.DEFINE_integer("embedding_dims", 256, "")
tf.flags.DEFINE_integer("batch_size", 20, "")
tf.flags.DEFINE_float("clip_gradients", 1.0, "")
tf.flags.DEFINE_float("eps", 1e-6, "")
tf.flags.DEFINE_float("param_init", 0.1, "")
tf.flags.DEFINE_float("learning_rate", 0.001, "")
tf.flags.DEFINE_float("l2_regularizer", 0.0001, "")
tf.flags.DEFINE_float("print_cost", 50.0,
                      "weighting factor in the objective function")
tf.flags.DEFINE_string("job_id", "temp", """job id""")
tf.flags.DEFINE_string("output_dir", "../model/",
                       """output_dir""")
tf.flags.DEFINE_string("data_dir", "../data/",
                       """data_dir""")
tf.flags.DEFINE_integer("write_every", 500, "wrtie every N")
tf.flags.DEFINE_integer("param_seed", 150, "")
tf.flags.DEFINE_integer("python_seed", 200, "")
tf.flags.DEFINE_float("dropout", 0.8, "dropout keep probability")
tf.flags.DEFINE_float("rnn_dropout", 0.9,
                      "dropout keep probability for rnn connections")
tf.flags.DEFINE_float("pad_int", -20000.0,
                      "number columns are padded with pad_int")
tf.flags.DEFINE_string("data_type", "double", "float or double")
tf.flags.DEFINE_float("word_dropout_prob", 0.9, "word dropout keep prob")
tf.flags.DEFINE_integer("word_cutoff", 10, "")
tf.flags.DEFINE_integer("vocab_size", 10800, "")
tf.flags.DEFINE_boolean("evaluator_job", False,
                        "wehther to run as trainer/evaluator")
tf.flags.DEFINE_float(
    "bad_number_pre_process", -200000.0,
    "number that is added to a corrupted table entry in a number column")
tf.flags.DEFINE_float("max_math_error", 3.0,
                      "max square loss error that is considered")
tf.flags.DEFINE_float("soft_min_value", 5.0, "")
FLAGS = tf.flags.FLAGS

class Utility:
  #holds FLAGS and other variables that are used in different files
  def __init__(self):
    global FLAGS
    self.FLAGS = FLAGS
    self.unk_token = "UNK"
    self.entry_match_token = "entry_match"
    self.column_match_token = "column_match"
    self.dummy_token = "dummy_token"
    self.tf_data_type = {}
    self.tf_data_type["double"] = tf.float64
    self.tf_data_type["float"] = tf.float32
    self.np_data_type = {}
    self.np_data_type["double"] = np.float64
    self.np_data_type["float"] = np.float32
    self.operations_set = ["count"] + [
        "prev", "next", "first_rs", "last_rs", "group_by_max", "greater",
        "lesser", "geq", "leq", "max", "min", "word-match"
    ] + ["reset_select"] + ["print"]
    self.word_ids = {}
    self.reverse_word_ids = {}
    self.word_count = {}
    self.random = Random(FLAGS.python_seed)

def evaluate(sess, data, batch_size, graph, i):
  #computes accuracy
  num_examples = 0.0
  gc = 0.0
  for j in range(0, len(data) - batch_size + 1, batch_size):
    [ct] = sess.run([graph.final_correct],
                    feed_dict=data_utils.generate_feed_dict(data, j, batch_size,
                                                            graph))
    gc += ct * batch_size
    num_examples += batch_size
  print "dev set accuracy   after ", i, " : ", gc / num_examples
  print num_examples, len(data)
  print "--------"

def Train(graph, utility, batch_size, train_data, sess, model_dir,
          saver):
  #performs training
  curr = 0
  train_set_loss = 0.0
  utility.random.shuffle(train_data)
  start = time.time()
  for i in range(utility.FLAGS.train_steps):
    curr_step = i
    if (i > 0 and i % FLAGS.write_every == 0):
      model_file = model_dir + "/model_" + str(i)
      saver.save(sess, model_file)
    if curr + batch_size >= len(train_data):
      curr = 0
      utility.random.shuffle(train_data)
    step, cost_value = sess.run(
        [graph.step, graph.total_cost],
        feed_dict=data_utils.generate_feed_dict(
            train_data, curr, batch_size, graph, train=True, utility=utility))
    curr = curr + batch_size
    train_set_loss += cost_value
    if (i > 0 and i % FLAGS.eval_cycle == 0):
      end = time.time()
      time_taken = end - start
      print "step ", i, " ", time_taken, " seconds "
      start = end
      print " printing train set loss: ", train_set_loss / utility.FLAGS.eval_cycle
      train_set_loss = 0.0

def master(train_data, dev_data, utility):
  #creates TF graph and calls trainer or evaluator
  batch_size = utility.FLAGS.batch_size
  model_dir = utility.FLAGS.output_dir + "/model" + utility.FLAGS.job_id + "/"
  #create all paramters of the model
  param_class = parameters.Parameters(utility)
  params, global_step, init = param_class.parameters(utility)
  key = "test" if (FLAGS.evaluator_job) else "train"
  graph = model.Graph(utility, batch_size, utility.FLAGS.max_passes, mode=key)
  graph.create_graph(params, global_step)
  prev_dev_error = 0.0
  final_loss = 0.0
  final_accuracy = 0.0
  #start session
  with tf.Session() as sess:
    sess.run(init.name)
    sess.run(graph.init_op.name)
    to_save = params.copy()
    saver = tf.train.Saver(to_save, max_to_keep=500)
    if (FLAGS.evaluator_job):
      while True:
        selected_models = {}
        file_list = tf.gfile.ListDirectory(model_dir)
        for model_file in file_list:
          if ("checkpoint" in model_file or "index" in model_file or
              "meta" in model_file):
            continue
          if ("data" in model_file):
            model_file = model_file.split(".")[0]
          model_step = int(
              model_file.split("_")[len(model_file.split("_")) - 1])
          selected_models[model_step] = model_file
        file_list = sorted(selected_models.items(), key=lambda x: x[0])
        if (len(file_list) > 0):
          file_list = file_list[0:len(file_list) - 1]
	print "list of models: ", file_list
        for model_file in file_list:
          model_file = model_file[1]
          print "restoring: ", model_file
          saver.restore(sess, model_dir + "/" + model_file)
          model_step = int(
              model_file.split("_")[len(model_file.split("_")) - 1])
          print "evaluating on dev ", model_file, model_step
          evaluate(sess, dev_data, batch_size, graph, model_step)
    else:
      ckpt = tf.train.get_checkpoint_state(model_dir)
      print "model dir: ", model_dir
      if (not (tf.gfile.IsDirectory(utility.FLAGS.output_dir))):
        print "create dir: ", utility.FLAGS.output_dir
        tf.gfile.MkDir(utility.FLAGS.output_dir)
      if (not (tf.gfile.IsDirectory(model_dir))):
        print "create dir: ", model_dir
        tf.gfile.MkDir(model_dir)
      Train(graph, utility, batch_size, train_data, sess, model_dir,
            saver)

def main(args):
  utility = Utility()
  train_name = "random-split-1-train.examples"
  dev_name = "random-split-1-dev.examples"
  test_name = "pristine-unseen-tables.examples"
  #load data
  dat = wiki_data.WikiQuestionGenerator(train_name, dev_name, test_name, FLAGS.data_dir)
  train_data, dev_data, test_data = dat.load()
  utility.words = []
  utility.word_ids = {}
  utility.reverse_word_ids = {}
  #construct vocabulary
  data_utils.construct_vocab(train_data, utility)
  data_utils.construct_vocab(dev_data, utility, True)
  data_utils.construct_vocab(test_data, utility, True)
  data_utils.add_special_words(utility)
  data_utils.perform_word_cutoff(utility)
  #convert data to int format and pad the inputs
  train_data = data_utils.complete_wiki_processing(train_data, utility, True)
  dev_data = data_utils.complete_wiki_processing(dev_data, utility, False)
  test_data = data_utils.complete_wiki_processing(test_data, utility, False)
  print "# train examples ", len(train_data)
  print "# dev examples ", len(dev_data)
  print "# test examples ", len(test_data)
  print "running open source"
  #construct TF graph and train or evaluate
  master(train_data, dev_data, utility)

if __name__ == "__main__":
  tf.app.run()
<|endoftext|>


<commit_before># vim: tabstop=4 shiftwidth=4 softtabstop=4

#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import datetime
import time

from oslo.config import cfg

from nova.openstack.common.gettextutils import _
from nova.openstack.common import log as logging
from nova.openstack.common import timeutils

periodic_opts = [
    cfg.BoolOpt('run_external_periodic_tasks',
                default=True,
                help=('Some periodic tasks can be run in a separate process. '
                      'Should we run them here?')),
]

CONF = cfg.CONF
CONF.register_opts(periodic_opts)

LOG = logging.getLogger(__name__)

DEFAULT_INTERVAL = 60.0

class InvalidPeriodicTaskArg(Exception):
    message = _("Unexpected argument for periodic task creation: %(arg)s.")

def periodic_task(*args, **kwargs):
    """Decorator to indicate that a method is a periodic task.

    This decorator can be used in two ways:

        1. Without arguments '@periodic_task', this will be run on every cycle
           of the periodic scheduler.

        2. With arguments:
           @periodic_task(spacing=N [, run_immediately=[True|False]])
           this will be run on approximately every N seconds. If this number is
           negative the periodic task will be disabled. If the run_immediately
           argument is provided and has a value of 'True', the first run of the
           task will be shortly after task scheduler starts.  If
           run_immediately is omitted or set to 'False', the first time the
           task runs will be approximately N seconds after the task scheduler
           starts.
    """
    def decorator(f):
        # Test for old style invocation
        if 'ticks_between_runs' in kwargs:
            raise InvalidPeriodicTaskArg(arg='ticks_between_runs')

        # Control if run at all
        f._periodic_task = True
        f._periodic_external_ok = kwargs.pop('external_process_ok', False)
        if f._periodic_external_ok and not CONF.run_external_periodic_tasks:
            f._periodic_enabled = False
        else:
            f._periodic_enabled = kwargs.pop('enabled', True)

        # Control frequency
        f._periodic_spacing = kwargs.pop('spacing', 0)
        f._periodic_immediate = kwargs.pop('run_immediately', False)
        if f._periodic_immediate:
            f._periodic_last_run = None
        else:
            f._periodic_last_run = timeutils.utcnow()
        return f

    # NOTE(sirp): The `if` is necessary to allow the decorator to be used with
    # and without parens.
    #
    # In the 'with-parens' case (with kwargs present), this function needs to
    # return a decorator function since the interpreter will invoke it like:
    #
    #   periodic_task(*args, **kwargs)(f)
    #
    # In the 'without-parens' case, the original function will be passed
    # in as the first argument, like:
    #
    #   periodic_task(f)
    if kwargs:
        return decorator
    else:
        return decorator(args[0])

class _PeriodicTasksMeta(type):
    def __init__(cls, names, bases, dict_):
        """Metaclass that allows us to collect decorated periodic tasks."""
        super(_PeriodicTasksMeta, cls).__init__(names, bases, dict_)

        # NOTE(sirp): if the attribute is not present then we must be the base
        # class, so, go ahead an initialize it. If the attribute is present,
        # then we're a subclass so make a copy of it so we don't step on our
        # parent's toes.
        try:
            cls._periodic_tasks = cls._periodic_tasks[:]
        except AttributeError:
            cls._periodic_tasks = []

        try:
            cls._periodic_last_run = cls._periodic_last_run.copy()
        except AttributeError:
            cls._periodic_last_run = {}

        try:
            cls._periodic_spacing = cls._periodic_spacing.copy()
        except AttributeError:
            cls._periodic_spacing = {}

        for value in cls.__dict__.values():
            if getattr(value, '_periodic_task', False):
                task = value
                name = task.__name__

                if task._periodic_spacing < 0:
                    LOG.info(_('Skipping periodic task %(task)s because '
                               'its interval is negative'),
                             {'task': name})
                    continue
                if not task._periodic_enabled:
                    LOG.info(_('Skipping periodic task %(task)s because '
                               'it is disabled'),
                             {'task': name})
                    continue

                # A periodic spacing of zero indicates that this task should
                # be run every pass
                if task._periodic_spacing == 0:
                    task._periodic_spacing = None

                cls._periodic_tasks.append((name, task))
                cls._periodic_spacing[name] = task._periodic_spacing
                cls._periodic_last_run[name] = task._periodic_last_run

class PeriodicTasks(object):
    __metaclass__ = _PeriodicTasksMeta

    def run_periodic_tasks(self, context, raise_on_error=False):
        """Tasks to be run at a periodic interval."""
        idle_for = DEFAULT_INTERVAL
        for task_name, task in self._periodic_tasks:
            full_task_name = '.'.join([self.__class__.__name__, task_name])

            now = timeutils.utcnow()
            spacing = self._periodic_spacing[task_name]
            last_run = self._periodic_last_run[task_name]

            # If a periodic task is _nearly_ due, then we'll run it early
            if spacing is not None and last_run is not None:
                due = last_run + datetime.timedelta(seconds=spacing)
                if not timeutils.is_soon(due, 0.2):
                    idle_for = min(idle_for, timeutils.delta_seconds(now, due))
                    continue

            if spacing is not None:
                idle_for = min(idle_for, spacing)

            LOG.debug(_("Running periodic task %(full_task_name)s"), locals())
            self._periodic_last_run[task_name] = timeutils.utcnow()

            try:
                task(self, context)
            except Exception as e:
                if raise_on_error:
                    raise
                LOG.exception(_("Error during %(full_task_name)s: %(e)s"),
                              locals())
            time.sleep(0)

        return idle_for
<commit_msg>Update openstack/common/periodic_task<commit_after># vim: tabstop=4 shiftwidth=4 softtabstop=4

#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import datetime
import time

from oslo.config import cfg
import six

from nova.openstack.common.gettextutils import _  # noqa
from nova.openstack.common import log as logging
from nova.openstack.common import timeutils

periodic_opts = [
    cfg.BoolOpt('run_external_periodic_tasks',
                default=True,
                help=('Some periodic tasks can be run in a separate process. '
                      'Should we run them here?')),
]

CONF = cfg.CONF
CONF.register_opts(periodic_opts)

LOG = logging.getLogger(__name__)

DEFAULT_INTERVAL = 60.0

class InvalidPeriodicTaskArg(Exception):
    message = _("Unexpected argument for periodic task creation: %(arg)s.")

def periodic_task(*args, **kwargs):
    """Decorator to indicate that a method is a periodic task.

    This decorator can be used in two ways:

        1. Without arguments '@periodic_task', this will be run on every cycle
           of the periodic scheduler.

        2. With arguments:
           @periodic_task(spacing=N [, run_immediately=[True|False]])
           this will be run on approximately every N seconds. If this number is
           negative the periodic task will be disabled. If the run_immediately
           argument is provided and has a value of 'True', the first run of the
           task will be shortly after task scheduler starts.  If
           run_immediately is omitted or set to 'False', the first time the
           task runs will be approximately N seconds after the task scheduler
           starts.
    """
    def decorator(f):
        # Test for old style invocation
        if 'ticks_between_runs' in kwargs:
            raise InvalidPeriodicTaskArg(arg='ticks_between_runs')

        # Control if run at all
        f._periodic_task = True
        f._periodic_external_ok = kwargs.pop('external_process_ok', False)
        if f._periodic_external_ok and not CONF.run_external_periodic_tasks:
            f._periodic_enabled = False
        else:
            f._periodic_enabled = kwargs.pop('enabled', True)

        # Control frequency
        f._periodic_spacing = kwargs.pop('spacing', 0)
        f._periodic_immediate = kwargs.pop('run_immediately', False)
        if f._periodic_immediate:
            f._periodic_last_run = None
        else:
            f._periodic_last_run = timeutils.utcnow()
        return f

    # NOTE(sirp): The `if` is necessary to allow the decorator to be used with
    # and without parens.
    #
    # In the 'with-parens' case (with kwargs present), this function needs to
    # return a decorator function since the interpreter will invoke it like:
    #
    #   periodic_task(*args, **kwargs)(f)
    #
    # In the 'without-parens' case, the original function will be passed
    # in as the first argument, like:
    #
    #   periodic_task(f)
    if kwargs:
        return decorator
    else:
        return decorator(args[0])

class _PeriodicTasksMeta(type):
    def __init__(cls, names, bases, dict_):
        """Metaclass that allows us to collect decorated periodic tasks."""
        super(_PeriodicTasksMeta, cls).__init__(names, bases, dict_)

        # NOTE(sirp): if the attribute is not present then we must be the base
        # class, so, go ahead an initialize it. If the attribute is present,
        # then we're a subclass so make a copy of it so we don't step on our
        # parent's toes.
        try:
            cls._periodic_tasks = cls._periodic_tasks[:]
        except AttributeError:
            cls._periodic_tasks = []

        try:
            cls._periodic_last_run = cls._periodic_last_run.copy()
        except AttributeError:
            cls._periodic_last_run = {}

        try:
            cls._periodic_spacing = cls._periodic_spacing.copy()
        except AttributeError:
            cls._periodic_spacing = {}

        for value in cls.__dict__.values():
            if getattr(value, '_periodic_task', False):
                task = value
                name = task.__name__

                if task._periodic_spacing < 0:
                    LOG.info(_('Skipping periodic task %(task)s because '
                               'its interval is negative'),
                             {'task': name})
                    continue
                if not task._periodic_enabled:
                    LOG.info(_('Skipping periodic task %(task)s because '
                               'it is disabled'),
                             {'task': name})
                    continue

                # A periodic spacing of zero indicates that this task should
                # be run every pass
                if task._periodic_spacing == 0:
                    task._periodic_spacing = None

                cls._periodic_tasks.append((name, task))
                cls._periodic_spacing[name] = task._periodic_spacing
                cls._periodic_last_run[name] = task._periodic_last_run

@six.add_metaclass(_PeriodicTasksMeta)
class PeriodicTasks(object):

    def run_periodic_tasks(self, context, raise_on_error=False):
        """Tasks to be run at a periodic interval."""
        idle_for = DEFAULT_INTERVAL
        for task_name, task in self._periodic_tasks:
            full_task_name = '.'.join([self.__class__.__name__, task_name])

            now = timeutils.utcnow()
            spacing = self._periodic_spacing[task_name]
            last_run = self._periodic_last_run[task_name]

            # If a periodic task is _nearly_ due, then we'll run it early
            if spacing is not None and last_run is not None:
                due = last_run + datetime.timedelta(seconds=spacing)
                if not timeutils.is_soon(due, 0.2):
                    idle_for = min(idle_for, timeutils.delta_seconds(now, due))
                    continue

            if spacing is not None:
                idle_for = min(idle_for, spacing)

            LOG.debug(_("Running periodic task %(full_task_name)s"),
                      {"full_task_name": full_task_name})
            self._periodic_last_run[task_name] = timeutils.utcnow()

            try:
                task(self, context)
            except Exception as e:
                if raise_on_error:
                    raise
                LOG.exception(_("Error during %(full_task_name)s: %(e)s"),
                              {"full_task_name": full_task_name, "e": e})
            time.sleep(0)

        return idle_for
<|endoftext|>


<commit_before># Copyright (c) 2011 Openstack, LLC.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""
The Zone Aware Scheduler is a base class Scheduler for creating instances
across zones. There are two expansion points to this class for:
1. Assigning Weights to hosts for requested instances
2. Filtering Hosts based on required instance capabilities
"""

import operator
import json

import M2Crypto
import novaclient

from nova import crypto
from nova import db
from nova import exception
from nova import flags
from nova import log as logging
from nova import rpc

from nova.scheduler import api
from nova.scheduler import driver

FLAGS = flags.FLAGS
LOG = logging.getLogger('nova.scheduler.zone_aware_scheduler')

class InvalidBlob(exception.NovaException):
    message = _("Ill-formed or incorrectly routed 'blob' data sent "
                "to instance create request.")

class ZoneAwareScheduler(driver.Scheduler):
    """Base class for creating Zone Aware Schedulers."""

    def _call_zone_method(self, context, method, specs):
        """Call novaclient zone method. Broken out for testing."""
        return api.call_zone_method(context, method, specs=specs)

    def _provision_resource_locally(self, context, item, instance_id, kwargs):
        """Create the requested resource in this Zone."""
        host = item['hostname']
        kwargs['instance_id'] = instance_id
        rpc.cast(context,
                 db.queue_get_for(context, "compute", host),
                 {"method": "run_instance",
                  "args": kwargs})
        LOG.debug(_("Provisioning locally via compute node %(host)s")
                            % locals())

    def _decrypt_blob(self, blob):
        """Returns the decrypted blob or None if invalid. Broken out
        for testing."""
        decryptor = crypto.decryptor(FLAGS.build_plan_encryption_key)
        try:
            json_entry = decryptor(blob)
            return json.dumps(entry)
        except M2Crypto.EVP.EVPError:
            pass
        return None

    def _ask_child_zone_to_create_instance(self, context, zone_info,
                                           request_spec, kwargs):
        """Once we have determined that the request should go to one
        of our children, we need to fabricate a new POST /servers/
        call with the same parameters that were passed into us.

        Note that we have to reverse engineer from our args to get back the
        image, flavor, ipgroup, etc. since the original call could have
        come in from EC2 (which doesn't use these things)."""

        instance_type = request_spec['instance_type']
        instance_properties = request_spec['instance_properties']

        name = instance_properties['display_name']
        image_ref = instance_properties['image_ref']
        meta = instance_properties['metadata']
        flavor_id = instance_type['flavorid']
        reservation_id = instance_properties['reservation_id']

        files = kwargs['injected_files']
        ipgroup = None  # Not supported in OS API ... yet

        child_zone = zone_info['child_zone']
        child_blob = zone_info['child_blob']
        zone = db.zone_get(context, child_zone)
        url = zone.api_url
        LOG.debug(_("Forwarding instance create call to child zone %(url)s"
                    ". ReservationID=%(reservation_id)s")
                    % locals())
        nova = None
        try:
            nova = novaclient.OpenStack(zone.username, zone.password, None,
                                        url)
            nova.authenticate()
        except novaclient.exceptions.BadRequest, e:
            raise exception.NotAuthorized(_("Bad credentials attempting "
                            "to talk to zone at %(url)s.") % locals())

        nova.servers.create(name, image_ref, flavor_id, ipgroup, meta, files,
                            child_blob, reservation_id=reservation_id)

    def _provision_resource_from_blob(self, context, item, instance_id,
                                          request_spec, kwargs):
        """Create the requested resource locally or in a child zone
           based on what is stored in the zone blob info.

           Attempt to decrypt the blob to see if this request is:
           1. valid, and
           2. intended for this zone or a child zone.

           Note: If we have "blob" that means the request was passed
           into us from a parent zone. If we have "child_blob" that
           means we gathered the info from one of our children.
           It's possible that, when we decrypt the 'blob' field, it
           contains "child_blob" data. In which case we forward the
           request."""

        host_info = None
        if "blob" in item:
            # Request was passed in from above. Is it for us?
            host_info = self._decrypt_blob(item['blob'])
        elif "child_blob" in item:
            # Our immediate child zone provided this info ...
            host_info = item

        if not host_info:
            raise InvalidBlob()

        # Valid data ... is it for us?
        if 'child_zone' in host_info and 'child_blob' in host_info:
            self._ask_child_zone_to_create_instance(context, host_info,
                                                    request_spec, kwargs)
        else:
            self._provision_resource_locally(context, host_info,
                                             instance_id, kwargs)

    def _provision_resource(self, context, item, instance_id, request_spec,
                           kwargs):
        """Create the requested resource in this Zone or a child zone."""
        if "hostname" in item:
            self._provision_resource_locally(context, item, instance_id,
                            kwargs)
            return

        self._provision_resource_from_blob(context, item, instance_id,
                                               request_spec, kwargs)

    def schedule_run_instance(self, context, instance_id, request_spec,
                              *args, **kwargs):
        """This method is called from nova.compute.api to provision
        an instance. However we need to look at the parameters being
        passed in to see if this is a request to:
        1. Create a Build Plan and then provision, or
        2. Use the Build Plan information in the request parameters
           to simply create the instance (either in this zone or
           a child zone).
        """

        # TODO(sandy): We'll have to look for richer specs at some point.

        blob = request_spec.get('blob')
        if blob:
            self._provision_resource(context, request_spec, instance_id,
                                    request_spec, kwargs)
            return None

        num_instances = request_spec['num_instances']
        LOG.debug(_("Attemping to build %d instance%s") %
            (num_instances, "" if num_instances == 1 else "s"))

        # Create build plan and provision ...
        build_plan = self.select(context, request_spec)
        if not build_plan:
            raise driver.NoValidHost(_('No hosts were available'))

        for num in xrange(num_instances):
            if not build_plan:
                break
            item = build_plan.pop(0)
            self._provision_resource(context, item, instance_id,
                    request_spec, kwargs)

        # Returning None short-circuits the routing to Compute (since
        # we've already done it here)
        return None

    def select(self, context, request_spec, *args, **kwargs):
        """Select returns a list of weights and zone/host information
        corresponding to the best hosts to service the request. Any
        child zone information has been encrypted so as not to reveal
        anything about the children.
        """
        return self._schedule(context, "compute", request_spec,
                              *args, **kwargs)

    # TODO(sandy): We're only focused on compute instances right now,
    # so we don't implement the default "schedule()" method required
    # of Schedulers.
    def schedule(self, context, topic, request_spec, *args, **kwargs):
        """The schedule() contract requires we return the one
        best-suited host for this request.
        """
        raise driver.NoValidHost(_('No hosts were available'))

    def _schedule(self, context, topic, request_spec, *args, **kwargs):
        """Returns a list of hosts that meet the required specs,
        ordered by their fitness.
        """

        if topic != "compute":
            raise NotImplemented(_("Zone Aware Scheduler only understands "
                                   "Compute nodes (for now)"))

        num_instances = request_spec['num_instances']
        instance_type = request_spec['instance_type']

        weighted = []
        host_list = None

        for i in xrange(num_instances):
            # Filter local hosts based on requirements ...
            #
            # The first pass through here will pass 'None' as the
            # host_list.. which tells the filter to build the full
            # list of hosts.
            # On a 2nd pass, the filter can modify the host_list with
            # any updates it needs to make based on resources that
            # may have been consumed from a previous build..
            host_list = self.filter_hosts(topic, request_spec, host_list)
            if not host_list:
                LOG.warn(_("Ran out of available hosts after weighing "
                        "%d of %d instances") % (i, num_instances))
                break

            # then weigh the selected hosts.
            # weighted = [{weight=weight, hostname=hostname,
            #              capabilities=capabs}, ...]
            weights = self.weigh_hosts(topic, request_spec, host_list)
            weights.sort(key=operator.itemgetter('weight'))
            best_weight = weights[0]
            weighted.append(best_weight)
            self.consume_resources(topic, best_weight['capabilities'],
                    instance_type)

        # Next, tack on the best weights from the child zones ...
        json_spec = json.dumps(request_spec)
        child_results = self._call_zone_method(context, "select",
                specs=json_spec)
        for child_zone, result in child_results:
            for weighting in result:
                # Remember the child_zone so we can get back to
                # it later if needed. This implicitly builds a zone
                # path structure.
                host_dict = {"weight": weighting["weight"],
                             "child_zone": child_zone,
                             "child_blob": weighting["blob"]}
                weighted.append(host_dict)

        weighted.sort(key=operator.itemgetter('weight'))
        return weighted

    def compute_filter(self, hostname, capabilities, request_spec):
        """Return whether or not we can schedule to this compute node.
        Derived classes should override this and return True if the host
        is acceptable for scheduling.
        """
        instance_type = request_spec['instance_type']
        requested_mem = instance_type['memory_mb'] * 1024 * 1024
        return capabilities['host_memory_free'] >= requested_mem

    def filter_hosts(self, topic, request_spec, host_list=None):
        """Return a list of hosts which are acceptable for scheduling.
        Return value should be a list of (hostname, capability_dict)s.
        Derived classes may override this, but may find the
        '<topic>_filter' function more appropriate.
        """

        def _default_filter(self, hostname, capabilities, request_spec):
            """Default filter function if there's no <topic>_filter"""
            # NOTE(sirp): The default logic is the equivalent to
            # AllHostsFilter
            return True

        filter_func = getattr(self, '%s_filter' % topic, _default_filter)

        if host_list is None:
            first_run = True
            host_list = self.zone_manager.service_states.iteritems()
        else:
            first_run = False

        filtered_hosts = []
        for host, services in host_list:
            if first_run:
                if topic not in services:
                    continue
                services = services['topic']
            if filter_func(host, services, request_spec):
                filtered_hosts.append((host, services))
        return filtered_hosts

    def weigh_hosts(self, topic, request_spec, hosts):
        """Derived classes may override this to provide more sophisticated
        scheduling objectives
        """
        # NOTE(sirp): The default logic is the same as the NoopCostFunction
        return [dict(weight=1, hostname=hostname, capabilities=capabilities)
                for hostname, capabilities in hosts]

    def compute_consume(self, capabilities, instance_type):
        """Consume compute resources for selected host"""

        requested_mem = max(instance_type['memory_mb'], 0) * 1024 * 1024
        capabilities['host_memory_free'] -= requested_mem

    def consume_resources(self, topic, capabilities, instance_type):
        """Consume resources for a specific host.  'host' is a tuple
        of the hostname and the services"""

        consume_func = getattr(self, '%s_consume' % topic, None)
        if not consume_func:
            return
        consume_func(capabilities, instance_type)
<commit_msg>same typo i made before!<commit_after># Copyright (c) 2011 Openstack, LLC.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

"""
The Zone Aware Scheduler is a base class Scheduler for creating instances
across zones. There are two expansion points to this class for:
1. Assigning Weights to hosts for requested instances
2. Filtering Hosts based on required instance capabilities
"""

import operator
import json

import M2Crypto
import novaclient

from nova import crypto
from nova import db
from nova import exception
from nova import flags
from nova import log as logging
from nova import rpc

from nova.scheduler import api
from nova.scheduler import driver

FLAGS = flags.FLAGS
LOG = logging.getLogger('nova.scheduler.zone_aware_scheduler')

class InvalidBlob(exception.NovaException):
    message = _("Ill-formed or incorrectly routed 'blob' data sent "
                "to instance create request.")

class ZoneAwareScheduler(driver.Scheduler):
    """Base class for creating Zone Aware Schedulers."""

    def _call_zone_method(self, context, method, specs):
        """Call novaclient zone method. Broken out for testing."""
        return api.call_zone_method(context, method, specs=specs)

    def _provision_resource_locally(self, context, item, instance_id, kwargs):
        """Create the requested resource in this Zone."""
        host = item['hostname']
        kwargs['instance_id'] = instance_id
        rpc.cast(context,
                 db.queue_get_for(context, "compute", host),
                 {"method": "run_instance",
                  "args": kwargs})
        LOG.debug(_("Provisioning locally via compute node %(host)s")
                            % locals())

    def _decrypt_blob(self, blob):
        """Returns the decrypted blob or None if invalid. Broken out
        for testing."""
        decryptor = crypto.decryptor(FLAGS.build_plan_encryption_key)
        try:
            json_entry = decryptor(blob)
            return json.dumps(entry)
        except M2Crypto.EVP.EVPError:
            pass
        return None

    def _ask_child_zone_to_create_instance(self, context, zone_info,
                                           request_spec, kwargs):
        """Once we have determined that the request should go to one
        of our children, we need to fabricate a new POST /servers/
        call with the same parameters that were passed into us.

        Note that we have to reverse engineer from our args to get back the
        image, flavor, ipgroup, etc. since the original call could have
        come in from EC2 (which doesn't use these things)."""

        instance_type = request_spec['instance_type']
        instance_properties = request_spec['instance_properties']

        name = instance_properties['display_name']
        image_ref = instance_properties['image_ref']
        meta = instance_properties['metadata']
        flavor_id = instance_type['flavorid']
        reservation_id = instance_properties['reservation_id']

        files = kwargs['injected_files']
        ipgroup = None  # Not supported in OS API ... yet

        child_zone = zone_info['child_zone']
        child_blob = zone_info['child_blob']
        zone = db.zone_get(context, child_zone)
        url = zone.api_url
        LOG.debug(_("Forwarding instance create call to child zone %(url)s"
                    ". ReservationID=%(reservation_id)s")
                    % locals())
        nova = None
        try:
            nova = novaclient.OpenStack(zone.username, zone.password, None,
                                        url)
            nova.authenticate()
        except novaclient.exceptions.BadRequest, e:
            raise exception.NotAuthorized(_("Bad credentials attempting "
                            "to talk to zone at %(url)s.") % locals())

        nova.servers.create(name, image_ref, flavor_id, ipgroup, meta, files,
                            child_blob, reservation_id=reservation_id)

    def _provision_resource_from_blob(self, context, item, instance_id,
                                          request_spec, kwargs):
        """Create the requested resource locally or in a child zone
           based on what is stored in the zone blob info.

           Attempt to decrypt the blob to see if this request is:
           1. valid, and
           2. intended for this zone or a child zone.

           Note: If we have "blob" that means the request was passed
           into us from a parent zone. If we have "child_blob" that
           means we gathered the info from one of our children.
           It's possible that, when we decrypt the 'blob' field, it
           contains "child_blob" data. In which case we forward the
           request."""

        host_info = None
        if "blob" in item:
            # Request was passed in from above. Is it for us?
            host_info = self._decrypt_blob(item['blob'])
        elif "child_blob" in item:
            # Our immediate child zone provided this info ...
            host_info = item

        if not host_info:
            raise InvalidBlob()

        # Valid data ... is it for us?
        if 'child_zone' in host_info and 'child_blob' in host_info:
            self._ask_child_zone_to_create_instance(context, host_info,
                                                    request_spec, kwargs)
        else:
            self._provision_resource_locally(context, host_info,
                                             instance_id, kwargs)

    def _provision_resource(self, context, item, instance_id, request_spec,
                           kwargs):
        """Create the requested resource in this Zone or a child zone."""
        if "hostname" in item:
            self._provision_resource_locally(context, item, instance_id,
                            kwargs)
            return

        self._provision_resource_from_blob(context, item, instance_id,
                                               request_spec, kwargs)

    def schedule_run_instance(self, context, instance_id, request_spec,
                              *args, **kwargs):
        """This method is called from nova.compute.api to provision
        an instance. However we need to look at the parameters being
        passed in to see if this is a request to:
        1. Create a Build Plan and then provision, or
        2. Use the Build Plan information in the request parameters
           to simply create the instance (either in this zone or
           a child zone).
        """

        # TODO(sandy): We'll have to look for richer specs at some point.

        blob = request_spec.get('blob')
        if blob:
            self._provision_resource(context, request_spec, instance_id,
                                    request_spec, kwargs)
            return None

        num_instances = request_spec['num_instances']
        LOG.debug(_("Attemping to build %d instance%s") %
            (num_instances, "" if num_instances == 1 else "s"))

        # Create build plan and provision ...
        build_plan = self.select(context, request_spec)
        if not build_plan:
            raise driver.NoValidHost(_('No hosts were available'))

        for num in xrange(num_instances):
            if not build_plan:
                break
            item = build_plan.pop(0)
            self._provision_resource(context, item, instance_id,
                    request_spec, kwargs)

        # Returning None short-circuits the routing to Compute (since
        # we've already done it here)
        return None

    def select(self, context, request_spec, *args, **kwargs):
        """Select returns a list of weights and zone/host information
        corresponding to the best hosts to service the request. Any
        child zone information has been encrypted so as not to reveal
        anything about the children.
        """
        return self._schedule(context, "compute", request_spec,
                              *args, **kwargs)

    # TODO(sandy): We're only focused on compute instances right now,
    # so we don't implement the default "schedule()" method required
    # of Schedulers.
    def schedule(self, context, topic, request_spec, *args, **kwargs):
        """The schedule() contract requires we return the one
        best-suited host for this request.
        """
        raise driver.NoValidHost(_('No hosts were available'))

    def _schedule(self, context, topic, request_spec, *args, **kwargs):
        """Returns a list of hosts that meet the required specs,
        ordered by their fitness.
        """

        if topic != "compute":
            raise NotImplemented(_("Zone Aware Scheduler only understands "
                                   "Compute nodes (for now)"))

        num_instances = request_spec['num_instances']
        instance_type = request_spec['instance_type']

        weighted = []
        host_list = None

        for i in xrange(num_instances):
            # Filter local hosts based on requirements ...
            #
            # The first pass through here will pass 'None' as the
            # host_list.. which tells the filter to build the full
            # list of hosts.
            # On a 2nd pass, the filter can modify the host_list with
            # any updates it needs to make based on resources that
            # may have been consumed from a previous build..
            host_list = self.filter_hosts(topic, request_spec, host_list)
            if not host_list:
                LOG.warn(_("Ran out of available hosts after weighing "
                        "%d of %d instances") % (i, num_instances))
                break

            # then weigh the selected hosts.
            # weighted = [{weight=weight, hostname=hostname,
            #              capabilities=capabs}, ...]
            weights = self.weigh_hosts(topic, request_spec, host_list)
            weights.sort(key=operator.itemgetter('weight'))
            best_weight = weights[0]
            weighted.append(best_weight)
            self.consume_resources(topic, best_weight['capabilities'],
                    instance_type)

        # Next, tack on the best weights from the child zones ...
        json_spec = json.dumps(request_spec)
        child_results = self._call_zone_method(context, "select",
                specs=json_spec)
        for child_zone, result in child_results:
            for weighting in result:
                # Remember the child_zone so we can get back to
                # it later if needed. This implicitly builds a zone
                # path structure.
                host_dict = {"weight": weighting["weight"],
                             "child_zone": child_zone,
                             "child_blob": weighting["blob"]}
                weighted.append(host_dict)

        weighted.sort(key=operator.itemgetter('weight'))
        return weighted

    def compute_filter(self, hostname, capabilities, request_spec):
        """Return whether or not we can schedule to this compute node.
        Derived classes should override this and return True if the host
        is acceptable for scheduling.
        """
        instance_type = request_spec['instance_type']
        requested_mem = instance_type['memory_mb'] * 1024 * 1024
        return capabilities['host_memory_free'] >= requested_mem

    def filter_hosts(self, topic, request_spec, host_list=None):
        """Return a list of hosts which are acceptable for scheduling.
        Return value should be a list of (hostname, capability_dict)s.
        Derived classes may override this, but may find the
        '<topic>_filter' function more appropriate.
        """

        def _default_filter(self, hostname, capabilities, request_spec):
            """Default filter function if there's no <topic>_filter"""
            # NOTE(sirp): The default logic is the equivalent to
            # AllHostsFilter
            return True

        filter_func = getattr(self, '%s_filter' % topic, _default_filter)

        if host_list is None:
            first_run = True
            host_list = self.zone_manager.service_states.iteritems()
        else:
            first_run = False

        filtered_hosts = []
        for host, services in host_list:
            if first_run:
                if topic not in services:
                    continue
                services = services[topic]
            if filter_func(host, services, request_spec):
                filtered_hosts.append((host, services))
        return filtered_hosts

    def weigh_hosts(self, topic, request_spec, hosts):
        """Derived classes may override this to provide more sophisticated
        scheduling objectives
        """
        # NOTE(sirp): The default logic is the same as the NoopCostFunction
        return [dict(weight=1, hostname=hostname, capabilities=capabilities)
                for hostname, capabilities in hosts]

    def compute_consume(self, capabilities, instance_type):
        """Consume compute resources for selected host"""

        requested_mem = max(instance_type['memory_mb'], 0) * 1024 * 1024
        capabilities['host_memory_free'] -= requested_mem

    def consume_resources(self, topic, capabilities, instance_type):
        """Consume resources for a specific host.  'host' is a tuple
        of the hostname and the services"""

        consume_func = getattr(self, '%s_consume' % topic, None)
        if not consume_func:
            return
        consume_func(capabilities, instance_type)
<|endoftext|>


<commit_before>from django.conf import settings

from rest_framework import viewsets, response, permissions

from metaci.api.views.testmethod_perf import BuildFlowFilterSet, DEFAULTS

from metaci.api.views.testmethod_perf import TestMethodPerfFilterSet
from metaci.api.views.testresults import TestMethodResultFilterSet

class TestMethodPerfUIApiView(viewsets.ViewSet):
    permission_classes = (permissions.AllowAny,)

    defaults = [
        (item, value)
        for (item, value) in DEFAULTS.__dict__.items()
        if not item.startswith("_")
    ]

    def list(self, request, format=None):
        buildflow_filters = self.collect_filter_defs(BuildFlowFilterSet, [])
        testmethodperf_filters = self.collect_filter_defs(
            TestMethodPerfFilterSet, buildflow_filters.keys()
        )
        testmethodresult_filters = self.collect_filter_defs(
            TestMethodResultFilterSet, buildflow_filters.keys()
        )

        json = {
            "buildflow_filters": buildflow_filters.values(),
            "testmethod_perf": {
                "includable_fields": testmethodperf_filters["include_fields"][
                    "choices"
                ],
                "filters": testmethodperf_filters.values(),
                "defaults": {**dict(self.defaults)},
            },
            "testmethod_results": {
                "includable_fields": testmethodresult_filters["include_fields"][
                    "choices"
                ],
                "filters": testmethodresult_filters.values(),
                "defaults": {**dict(self.defaults)},
            },
            "debug": settings.DEBUG,
        }

        return response.Response(json)

    def collect_filter_defs(self, filterSet, exclusion_names):
        json_filter_defs = {}
        defaults = dict(self.defaults)
        for name, filter in filterSet.get_filters().items():
            if name in exclusion_names:
                continue
            obj = {
                "name": name,
                "label": filter._label,
                "field_type": filter.field_class.__name__,
                "field_module": filter.field_class.__module__,
                "lookup_expr": filter.lookup_expr,
            }
            choices = filter.extra.get("choices")
            if choices:
                obj["choices"] = choices

            initial = filter.extra.get("initial")
            if initial:
                obj["initial"] = initial
            elif name in defaults:
                obj["initial"] = defaults[name]

            if "ui" in filter.extra:
                for key, value in filter.extra["ui"].items():
                    obj[key] = value

            json_filter_defs[name] = obj

        return json_filter_defs
<commit_msg>Minor cleanups.<commit_after>from django.conf import settings

from rest_framework import viewsets, response, permissions

from metaci.api.views.testmethod_perf import (
    BuildFlowFilterSet,
    TestMethodPerfFilterSet,
    DEFAULTS,
)
from metaci.api.views.testresults import TestMethodResultFilterSet

class TestMethodPerfUIApiView(viewsets.ViewSet):
    permission_classes = (permissions.AllowAny,)

    defaults = [
        (item, value)
        for (item, value) in DEFAULTS.__dict__.items()
        if not item.startswith("_")
    ]

    def list(self, request, format=None):
        buildflow_filters = self.collect_filter_defs(BuildFlowFilterSet, [])
        testmethodperf_filters = self.collect_filter_defs(
            TestMethodPerfFilterSet, buildflow_filters.keys()
        )
        testmethodresult_filters = self.collect_filter_defs(
            TestMethodResultFilterSet, buildflow_filters.keys()
        )

        json = {
            "buildflow_filters": buildflow_filters.values(),
            "testmethod_perf": {
                "includable_fields": testmethodperf_filters["include_fields"][
                    "choices"
                ],
                "filters": testmethodperf_filters.values(),
                "defaults": {**dict(self.defaults)},
            },
            "testmethod_results": {
                "includable_fields": testmethodresult_filters["include_fields"][
                    "choices"
                ],
                "filters": testmethodresult_filters.values(),
                "defaults": {**dict(self.defaults)},
            },
            "debug": settings.DEBUG,
        }

        return response.Response(json)

    def collect_filter_defs(self, filterSet, exclusion_names):
        json_filter_defs = {}
        defaults = dict(self.defaults)
        for name, filter in filterSet.get_filters().items():
            if name in exclusion_names:
                continue
            obj = {
                "name": name,
                "label": filter._label,
                "field_type": filter.field_class.__name__,
                "field_module": filter.field_class.__module__,
                "lookup_expr": filter.lookup_expr,
            }
            choices = filter.extra.get("choices")
            if choices:
                obj["choices"] = choices

            initial = filter.extra.get("initial")
            if initial:
                obj["initial"] = initial
            elif name in defaults:
                obj["initial"] = defaults[name]

            json_filter_defs[name] = obj

        return json_filter_defs
<|endoftext|>


<commit_before># Copyright 2013-2014, Simon Kennedy, code@sffjunkie.co.uk
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Part of 'hiss' the Python notification library

import pytest
import asyncio
from hiss.target import Target
from hiss.notifier import Notifier
from hiss.handler.xbmc import XBMCHandler

HOST = '127.0.0.1'
#HOST = '10.84.23.66'

asyncio.log.logger.setLevel(asyncio.log.logging.INFO)

@pytest.fixture
def notifier():
    n = Notifier('Notifier Test', 'application/x-vnd.sffjunkie.hiss')
    n.add_notification('New', 'New email received.')
    n.add_notification('Old', 'Old as an old thing.')
    return n

def test_notification(notifier):
    loop = asyncio.get_event_loop()

    @asyncio.coroutine
    def coro():
        h = XBMCHandler(loop=loop)

        t = Target('xbmc://%s' % HOST)

        notification = notifier.create_notification(name='New',
                                             title="A brave new world",
                                             text="Say hello to Prism",
                                             icon='info')
        response = yield from h.notify(notification, t)
        if response['status'] == 'ERROR':
            assert response['reason'] == 'All hosts are unreachable.'
        else:
            assert response['status_code'] == 0

    c = coro()
    loop.run_until_complete(c)

if __name__ == '__main__':
    test_notification()
    <commit_msg>* Renamed test functions * Removed __main__ test and function calls<commit_after># Copyright 2013-2014, Simon Kennedy, code@sffjunkie.co.uk
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Part of 'hiss' the Python notification library

import pytest
import asyncio
from hiss.target import Target
from hiss.notifier import Notifier
from hiss.handler.xbmc import XBMCHandler

HOST = '127.0.0.1'
#HOST = '10.84.23.66'

asyncio.log.logger.setLevel(asyncio.log.logging.INFO)

@pytest.fixture
def notifier():
    n = Notifier('Notifier Test', 'application/x-vnd.sffjunkie.hiss')
    n.add_notification('New', 'New email received.')
    n.add_notification('Old', 'Old as an old thing.')
    return n

def test_XBMC_Notification(notifier):
    loop = asyncio.get_event_loop()

    @asyncio.coroutine
    def coro():
        h = XBMCHandler(loop=loop)

        t = Target('xbmc://%s' % HOST)

        notification = notifier.create_notification(name='New',
                                             title="A brave new world",
                                             text="Say hello to Prism",
                                             icon='info')
        response = yield from h.notify(notification, t)
        if response['status'] == 'ERROR':
            assert response['reason'] == 'All hosts are unreachable.'
        else:
            assert response['status_code'] == 0

    c = coro()
    loop.run_until_complete(c)
<|endoftext|>


<commit_before>from pyspark import keyword_only
from pyspark.ml.util import JavaMLWritable
from pyspark.ml.wrapper import JavaTransformer, JavaEstimator
from pyspark.ml.param.shared import Param, Params, TypeConverters
from pyspark.ml.pipeline import Pipeline, PipelineModel, Estimator, Transformer
from sparknlp.common import ParamsGettersSetters, AnnotatorProperties
from sparknlp.util import AnnotatorJavaMLReadable
import sparknlp.internal as _internal

class AnnotatorTransformer(JavaTransformer, AnnotatorJavaMLReadable, JavaMLWritable, ParamsGettersSetters):
    @keyword_only
    def __init__(self, classname):
        super(AnnotatorTransformer, self).__init__()
        kwargs = self._input_kwargs
        if 'classname' in kwargs:
            kwargs.pop('classname')
        self.setParams(**kwargs)
        self.__class__._java_class_name = classname
        self._java_obj = self._new_java_obj(classname, self.uid)

class JavaRecursiveEstimator(JavaEstimator):

    def _fit_java(self, dataset, pipeline=None):
        """
        Fits a Java model to the input dataset.
        :param dataset: input dataset, which is an instance of
                        :py:class:`pyspark.sql.DataFrame`
        :param params: additional params (overwriting embedded values)
        :return: fitted Java model
        """
        self._transfer_params_to_java()
        if pipeline:
            return self._java_obj.recursiveFit(dataset._jdf, pipeline._to_java())
        else:
            return self._java_obj.fit(dataset._jdf)

    def _fit(self, dataset, pipeline=None):
        java_model = self._fit_java(dataset, pipeline)
        model = self._create_model(java_model)
        return self._copyValues(model)

    def fit(self, dataset, params=None, pipeline=None):
        """
        Fits a model to the input dataset with optional parameters.
        :param dataset: input dataset, which is an instance of :py:class:`pyspark.sql.DataFrame`
        :param params: an optional param map that overrides embedded params. If a list/tuple of
                       param maps is given, this calls fit on each param map and returns a list of
                       models.
        :returns: fitted model(s)
        """
        if params is None:
            params = dict()
        if isinstance(params, (list, tuple)):
            models = [None] * len(params)
            for index, model in self.fitMultiple(dataset, params):
                models[index] = model
            return models
        elif isinstance(params, dict):
            if params:
                return self.copy(params)._fit(dataset, pipeline=pipeline)
            else:
                return self._fit(dataset, pipeline=pipeline)
        else:
            raise ValueError("Params must be either a param map or a list/tuple of param maps, "
                             "but got %s." % type(params))

class Annotation:
    def __init__(self, annotator_type, begin, end, result, metadata, embeddings=[], sentence_embeddings=[]):
        self.annotator_type = annotator_type
        self.begin = begin
        self.end = end
        self.result = result
        self.metadata = metadata
        self.embeddings = embeddings
        self.sentence_embeddings = sentence_embeddings

class LightPipeline:
    def __init__(self, pipelineModel, parse_embeddings=False):
        self.pipeline_model = pipelineModel
        self._lightPipeline = _internal._LightPipeline(pipelineModel, parse_embeddings).apply()

    @staticmethod
    def _annotation_from_java(java_annotations):
        annotations = []
        for annotation in java_annotations:
            annotations.append(Annotation(annotation.annotatorType(),
                                          annotation.begin(),
                                          annotation.end(),
                                          annotation.result(),
                                          dict(annotation.metadata()),
                                          annotation.embeddings,
                                          annotation.sentence_embeddings
                                          )
                               )
        return annotations

    def fullAnnotate(self, target):
        result = []
        if type(target) is str:
            target = [target]
        for row in self._lightPipeline.fullAnnotateJava(target):
            kas = {}
            for atype, annotations in row.items():
                kas[atype] = self._annotation_from_java(annotations)
            result.append(kas)
        return result

    def annotate(self, target):

        def reformat(annotations):
            return {k: list(v) for k, v in annotations.items()}

        annotations = self._lightPipeline.annotateJava(target)

        if type(target) is str:
            result = reformat(annotations)
        elif type(target) is list:
            result = list(map(lambda a: reformat(a), list(annotations)))
        else:
            raise TypeError("target for annotation may be 'str' or 'list'")

        return result

    def transform(self, dataframe):
        return self.pipeline_model.transform(dataframe)

    def setIgnoreUnsupported(self, value):
        self._lightPipeline.setIgnoreUnsupported(value)
        return self

    def getIgnoreUnsupported(self):
        return self._lightPipeline.getIgnoreUnsupported()

class RecursivePipeline(Pipeline, JavaEstimator):
    @keyword_only
    def __init__(self, *args, **kwargs):
        super(RecursivePipeline, self).__init__(*args, **kwargs)
        self._java_obj = self._new_java_obj("com.johnsnowlabs.nlp.RecursivePipeline", self.uid)
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    def _fit(self, dataset):
        stages = self.getStages()
        for stage in stages:
            if not (isinstance(stage, Estimator) or isinstance(stage, Transformer)):
                raise TypeError(
                    "Cannot recognize a pipeline stage of type %s." % type(stage))
        indexOfLastEstimator = -1
        for i, stage in enumerate(stages):
            if isinstance(stage, Estimator):
                indexOfLastEstimator = i
        transformers = []
        for i, stage in enumerate(stages):
            if isinstance(stage, Transformer):
                transformers.append(stage)
                dataset = stage.transform(dataset)
            elif isinstance(stage, JavaRecursiveEstimator):
                model = stage.fit(dataset, pipeline=PipelineModel(transformers))
                transformers.append(model)
                if i < indexOfLastEstimator:
                    dataset = model.transform(dataset)
            else:
                model = stage.fit(dataset)
                transformers.append(model)
                if i < indexOfLastEstimator:
                    dataset = model.transform(dataset)
            if i <= indexOfLastEstimator:
                pass
            else:
                transformers.append(stage)
        return PipelineModel(transformers)

class DocumentAssembler(AnnotatorTransformer):

    inputCol = Param(Params._dummy(), "inputCol", "input column name", typeConverter=TypeConverters.toString)
    outputCol = Param(Params._dummy(), "outputCol", "output column name", typeConverter=TypeConverters.toString)
    idCol = Param(Params._dummy(), "idCol", "column for setting an id to such string in row", typeConverter=TypeConverters.toString)
    metadataCol = Param(Params._dummy(), "metadataCol", "String to String map column to use as metadata", typeConverter=TypeConverters.toString)
    calculationsCol = Param(Params._dummy(), "calculationsCol", "String to Float vector map column to use as embeddigns and other representations", typeConverter=TypeConverters.toString)
    cleanupMode = Param(Params._dummy(), "cleanupMode", "possible values: disabled, inplace, inplace_full, shrink, shrink_full, each, each_full, delete_full", typeConverter=TypeConverters.toString)
    name = 'DocumentAssembler'

    @keyword_only
    def __init__(self):
        super(DocumentAssembler, self).__init__(classname="com.johnsnowlabs.nlp.DocumentAssembler")
        self._setDefault(outputCol="document", cleanupMode='disabled')

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    def setInputCol(self, value):
        return self._set(inputCol=value)

    def setOutputCol(self, value):
        return self._set(outputCol=value)

    def setIdCol(self, value):
        return self._set(idCol=value)

    def setMetadataCol(self, value):
        return self._set(metadataCol=value)

    def setCalculationsCol(self, value):
        return self._set(metadataCol=value)

    def setCleanupMode(self, value):
        if value.strip().lower() not in ['disabled', 'inplace', 'inplace_full', 'shrink', 'shrink_full', 'each', 'each_full', 'delete_full']:
            raise Exception("Cleanup mode possible values: disabled, inplace, inplace_full, shrink, shrink_full, each, each_full, delete_full")
        return self._set(cleanupMode=value)

class TokenAssembler(AnnotatorTransformer, AnnotatorProperties):

    name = "TokenAssembler"

    @keyword_only
    def __init__(self):
        super(TokenAssembler, self).__init__(classname="com.johnsnowlabs.nlp.TokenAssembler")

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

class Doc2Chunk(AnnotatorTransformer, AnnotatorProperties):

    chunkCol = Param(Params._dummy(), "chunkCol", "column that contains string. Must be part of DOCUMENT", typeConverter=TypeConverters.toString)
    startCol = Param(Params._dummy(), "startCol", "column that has a reference of where chunk begins", typeConverter=TypeConverters.toString)
    startColByTokenIndex = Param(Params._dummy(), "startColByTokenIndex", "whether start col is by whitespace tokens", typeConverter=TypeConverters.toBoolean)
    isArray = Param(Params._dummy(), "isArray", "whether the chunkCol is an array of strings", typeConverter=TypeConverters.toBoolean)
    failOnMissing = Param(Params._dummy(), "failOnMissing", "whether to fail the job if a chunk is not found within document. return empty otherwise", typeConverter=TypeConverters.toBoolean)
    lowerCase = Param(Params._dummy(), "lowerCase", "whether to lower case for matching case", typeConverter=TypeConverters.toBoolean)
    name = "Doc2Chunk"

    @keyword_only
    def __init__(self):
        super(Doc2Chunk, self).__init__(classname="com.johnsnowlabs.nlp.Doc2Chunk")
        self._setDefault(
            isArray=False
        )

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    def setChunkCol(self, value):
        return self._set(chunkCol=value)

    def setIsArray(self, value):
        return self._set(isArray=value)

    def setStartCol(self, value):
        return self._set(startCol=value)

    def setStartColByTokenIndex(self, value):
        return self._set(startColByTokenIndex=value)

    def setFailOnMissing(self, value):
        return self._set(failOnMissing=value)

    def setLowerCase(self, value):
        return self._set(lowerCase=value)

class Chunk2Doc(AnnotatorTransformer, AnnotatorProperties):

    name = "Chunk2Doc"

    @keyword_only
    def __init__(self):
        super(Chunk2Doc, self).__init__(classname="com.johnsnowlabs.nlp.Chunk2Doc")

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

class Finisher(AnnotatorTransformer):

    inputCols = Param(Params._dummy(), "inputCols", "input annotations", typeConverter=TypeConverters.toListString)
    outputCols = Param(Params._dummy(), "outputCols", "output finished annotation cols", typeConverter=TypeConverters.toListString)
    valueSplitSymbol = Param(Params._dummy(), "valueSplitSymbol", "character separating annotations", typeConverter=TypeConverters.toString)
    annotationSplitSymbol = Param(Params._dummy(), "annotationSplitSymbol", "character separating annotations", typeConverter=TypeConverters.toString)
    cleanAnnotations = Param(Params._dummy(), "cleanAnnotations", "whether to remove annotation columns", typeConverter=TypeConverters.toBoolean)
    includeMetadata = Param(Params._dummy(), "includeMetadata", "annotation metadata format", typeConverter=TypeConverters.toBoolean)
    outputAsArray = Param(Params._dummy(), "outputAsArray", "finisher generates an Array with the results instead of string", typeConverter=TypeConverters.toBoolean)
    parseEmbeddingsVectors = Param(Params._dummy(), "parseEmbeddingsVectors", "whether to include embeddings vectors in the process", typeConverter=TypeConverters.toBoolean)

    name = "Finisher"

    @keyword_only
    def __init__(self):
        super(Finisher, self).__init__(classname="com.johnsnowlabs.nlp.Finisher")
        self._setDefault(
            cleanAnnotations=True,
            includeMetadata=False,
            outputAsArray=True,
            parseEmbeddingsVectors=False
        )

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    def setInputCols(self, *value):
        if len(value) == 1 and type(value[0]) == list:
            return self._set(inputCols=value[0])
        else:
            return self._set(inputCols=list(value))

    def setOutputCols(self, *value):
        if len(value) == 1 and type(value[0]) == list:
            return self._set(outputCols=value[0])
        else:
            return self._set(outputCols=list(value))

    def setValueSplitSymbol(self, value):
        return self._set(valueSplitSymbol=value)

    def setAnnotationSplitSymbol(self, value):
        return self._set(annotationSplitSymbol=value)

    def setCleanAnnotations(self, value):
        return self._set(cleanAnnotations=value)

    def setIncludeMetadata(self, value):
        return self._set(includeMetadata=value)

    def setOutputAsArray(self, value):
        return self._set(outputAsArray=value)
<commit_msg>Added missing setter<commit_after>from pyspark import keyword_only
from pyspark.ml.util import JavaMLWritable
from pyspark.ml.wrapper import JavaTransformer, JavaEstimator
from pyspark.ml.param.shared import Param, Params, TypeConverters
from pyspark.ml.pipeline import Pipeline, PipelineModel, Estimator, Transformer
from sparknlp.common import ParamsGettersSetters, AnnotatorProperties
from sparknlp.util import AnnotatorJavaMLReadable
import sparknlp.internal as _internal

class AnnotatorTransformer(JavaTransformer, AnnotatorJavaMLReadable, JavaMLWritable, ParamsGettersSetters):
    @keyword_only
    def __init__(self, classname):
        super(AnnotatorTransformer, self).__init__()
        kwargs = self._input_kwargs
        if 'classname' in kwargs:
            kwargs.pop('classname')
        self.setParams(**kwargs)
        self.__class__._java_class_name = classname
        self._java_obj = self._new_java_obj(classname, self.uid)

class JavaRecursiveEstimator(JavaEstimator):

    def _fit_java(self, dataset, pipeline=None):
        """
        Fits a Java model to the input dataset.
        :param dataset: input dataset, which is an instance of
                        :py:class:`pyspark.sql.DataFrame`
        :param params: additional params (overwriting embedded values)
        :return: fitted Java model
        """
        self._transfer_params_to_java()
        if pipeline:
            return self._java_obj.recursiveFit(dataset._jdf, pipeline._to_java())
        else:
            return self._java_obj.fit(dataset._jdf)

    def _fit(self, dataset, pipeline=None):
        java_model = self._fit_java(dataset, pipeline)
        model = self._create_model(java_model)
        return self._copyValues(model)

    def fit(self, dataset, params=None, pipeline=None):
        """
        Fits a model to the input dataset with optional parameters.
        :param dataset: input dataset, which is an instance of :py:class:`pyspark.sql.DataFrame`
        :param params: an optional param map that overrides embedded params. If a list/tuple of
                       param maps is given, this calls fit on each param map and returns a list of
                       models.
        :returns: fitted model(s)
        """
        if params is None:
            params = dict()
        if isinstance(params, (list, tuple)):
            models = [None] * len(params)
            for index, model in self.fitMultiple(dataset, params):
                models[index] = model
            return models
        elif isinstance(params, dict):
            if params:
                return self.copy(params)._fit(dataset, pipeline=pipeline)
            else:
                return self._fit(dataset, pipeline=pipeline)
        else:
            raise ValueError("Params must be either a param map or a list/tuple of param maps, "
                             "but got %s." % type(params))

class Annotation:
    def __init__(self, annotator_type, begin, end, result, metadata, embeddings=[], sentence_embeddings=[]):
        self.annotator_type = annotator_type
        self.begin = begin
        self.end = end
        self.result = result
        self.metadata = metadata
        self.embeddings = embeddings
        self.sentence_embeddings = sentence_embeddings

class LightPipeline:
    def __init__(self, pipelineModel, parse_embeddings=False):
        self.pipeline_model = pipelineModel
        self._lightPipeline = _internal._LightPipeline(pipelineModel, parse_embeddings).apply()

    @staticmethod
    def _annotation_from_java(java_annotations):
        annotations = []
        for annotation in java_annotations:
            annotations.append(Annotation(annotation.annotatorType(),
                                          annotation.begin(),
                                          annotation.end(),
                                          annotation.result(),
                                          dict(annotation.metadata()),
                                          annotation.embeddings,
                                          annotation.sentence_embeddings
                                          )
                               )
        return annotations

    def fullAnnotate(self, target):
        result = []
        if type(target) is str:
            target = [target]
        for row in self._lightPipeline.fullAnnotateJava(target):
            kas = {}
            for atype, annotations in row.items():
                kas[atype] = self._annotation_from_java(annotations)
            result.append(kas)
        return result

    def annotate(self, target):

        def reformat(annotations):
            return {k: list(v) for k, v in annotations.items()}

        annotations = self._lightPipeline.annotateJava(target)

        if type(target) is str:
            result = reformat(annotations)
        elif type(target) is list:
            result = list(map(lambda a: reformat(a), list(annotations)))
        else:
            raise TypeError("target for annotation may be 'str' or 'list'")

        return result

    def transform(self, dataframe):
        return self.pipeline_model.transform(dataframe)

    def setIgnoreUnsupported(self, value):
        self._lightPipeline.setIgnoreUnsupported(value)
        return self

    def getIgnoreUnsupported(self):
        return self._lightPipeline.getIgnoreUnsupported()

class RecursivePipeline(Pipeline, JavaEstimator):
    @keyword_only
    def __init__(self, *args, **kwargs):
        super(RecursivePipeline, self).__init__(*args, **kwargs)
        self._java_obj = self._new_java_obj("com.johnsnowlabs.nlp.RecursivePipeline", self.uid)
        kwargs = self._input_kwargs
        self.setParams(**kwargs)

    def _fit(self, dataset):
        stages = self.getStages()
        for stage in stages:
            if not (isinstance(stage, Estimator) or isinstance(stage, Transformer)):
                raise TypeError(
                    "Cannot recognize a pipeline stage of type %s." % type(stage))
        indexOfLastEstimator = -1
        for i, stage in enumerate(stages):
            if isinstance(stage, Estimator):
                indexOfLastEstimator = i
        transformers = []
        for i, stage in enumerate(stages):
            if isinstance(stage, Transformer):
                transformers.append(stage)
                dataset = stage.transform(dataset)
            elif isinstance(stage, JavaRecursiveEstimator):
                model = stage.fit(dataset, pipeline=PipelineModel(transformers))
                transformers.append(model)
                if i < indexOfLastEstimator:
                    dataset = model.transform(dataset)
            else:
                model = stage.fit(dataset)
                transformers.append(model)
                if i < indexOfLastEstimator:
                    dataset = model.transform(dataset)
            if i <= indexOfLastEstimator:
                pass
            else:
                transformers.append(stage)
        return PipelineModel(transformers)

class DocumentAssembler(AnnotatorTransformer):

    inputCol = Param(Params._dummy(), "inputCol", "input column name", typeConverter=TypeConverters.toString)
    outputCol = Param(Params._dummy(), "outputCol", "output column name", typeConverter=TypeConverters.toString)
    idCol = Param(Params._dummy(), "idCol", "column for setting an id to such string in row", typeConverter=TypeConverters.toString)
    metadataCol = Param(Params._dummy(), "metadataCol", "String to String map column to use as metadata", typeConverter=TypeConverters.toString)
    calculationsCol = Param(Params._dummy(), "calculationsCol", "String to Float vector map column to use as embeddigns and other representations", typeConverter=TypeConverters.toString)
    cleanupMode = Param(Params._dummy(), "cleanupMode", "possible values: disabled, inplace, inplace_full, shrink, shrink_full, each, each_full, delete_full", typeConverter=TypeConverters.toString)
    name = 'DocumentAssembler'

    @keyword_only
    def __init__(self):
        super(DocumentAssembler, self).__init__(classname="com.johnsnowlabs.nlp.DocumentAssembler")
        self._setDefault(outputCol="document", cleanupMode='disabled')

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    def setInputCol(self, value):
        return self._set(inputCol=value)

    def setOutputCol(self, value):
        return self._set(outputCol=value)

    def setIdCol(self, value):
        return self._set(idCol=value)

    def setMetadataCol(self, value):
        return self._set(metadataCol=value)

    def setCalculationsCol(self, value):
        return self._set(metadataCol=value)

    def setCleanupMode(self, value):
        if value.strip().lower() not in ['disabled', 'inplace', 'inplace_full', 'shrink', 'shrink_full', 'each', 'each_full', 'delete_full']:
            raise Exception("Cleanup mode possible values: disabled, inplace, inplace_full, shrink, shrink_full, each, each_full, delete_full")
        return self._set(cleanupMode=value)

class TokenAssembler(AnnotatorTransformer, AnnotatorProperties):

    name = "TokenAssembler"

    @keyword_only
    def __init__(self):
        super(TokenAssembler, self).__init__(classname="com.johnsnowlabs.nlp.TokenAssembler")

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

class Doc2Chunk(AnnotatorTransformer, AnnotatorProperties):

    chunkCol = Param(Params._dummy(), "chunkCol", "column that contains string. Must be part of DOCUMENT", typeConverter=TypeConverters.toString)
    startCol = Param(Params._dummy(), "startCol", "column that has a reference of where chunk begins", typeConverter=TypeConverters.toString)
    startColByTokenIndex = Param(Params._dummy(), "startColByTokenIndex", "whether start col is by whitespace tokens", typeConverter=TypeConverters.toBoolean)
    isArray = Param(Params._dummy(), "isArray", "whether the chunkCol is an array of strings", typeConverter=TypeConverters.toBoolean)
    failOnMissing = Param(Params._dummy(), "failOnMissing", "whether to fail the job if a chunk is not found within document. return empty otherwise", typeConverter=TypeConverters.toBoolean)
    lowerCase = Param(Params._dummy(), "lowerCase", "whether to lower case for matching case", typeConverter=TypeConverters.toBoolean)
    name = "Doc2Chunk"

    @keyword_only
    def __init__(self):
        super(Doc2Chunk, self).__init__(classname="com.johnsnowlabs.nlp.Doc2Chunk")
        self._setDefault(
            isArray=False
        )

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    def setChunkCol(self, value):
        return self._set(chunkCol=value)

    def setIsArray(self, value):
        return self._set(isArray=value)

    def setStartCol(self, value):
        return self._set(startCol=value)

    def setStartColByTokenIndex(self, value):
        return self._set(startColByTokenIndex=value)

    def setFailOnMissing(self, value):
        return self._set(failOnMissing=value)

    def setLowerCase(self, value):
        return self._set(lowerCase=value)

class Chunk2Doc(AnnotatorTransformer, AnnotatorProperties):

    name = "Chunk2Doc"

    @keyword_only
    def __init__(self):
        super(Chunk2Doc, self).__init__(classname="com.johnsnowlabs.nlp.Chunk2Doc")

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

class Finisher(AnnotatorTransformer):

    inputCols = Param(Params._dummy(), "inputCols", "input annotations", typeConverter=TypeConverters.toListString)
    outputCols = Param(Params._dummy(), "outputCols", "output finished annotation cols", typeConverter=TypeConverters.toListString)
    valueSplitSymbol = Param(Params._dummy(), "valueSplitSymbol", "character separating annotations", typeConverter=TypeConverters.toString)
    annotationSplitSymbol = Param(Params._dummy(), "annotationSplitSymbol", "character separating annotations", typeConverter=TypeConverters.toString)
    cleanAnnotations = Param(Params._dummy(), "cleanAnnotations", "whether to remove annotation columns", typeConverter=TypeConverters.toBoolean)
    includeMetadata = Param(Params._dummy(), "includeMetadata", "annotation metadata format", typeConverter=TypeConverters.toBoolean)
    outputAsArray = Param(Params._dummy(), "outputAsArray", "finisher generates an Array with the results instead of string", typeConverter=TypeConverters.toBoolean)
    parseEmbeddingsVectors = Param(Params._dummy(), "parseEmbeddingsVectors", "whether to include embeddings vectors in the process", typeConverter=TypeConverters.toBoolean)

    name = "Finisher"

    @keyword_only
    def __init__(self):
        super(Finisher, self).__init__(classname="com.johnsnowlabs.nlp.Finisher")
        self._setDefault(
            cleanAnnotations=True,
            includeMetadata=False,
            outputAsArray=True,
            parseEmbeddingsVectors=False
        )

    @keyword_only
    def setParams(self):
        kwargs = self._input_kwargs
        return self._set(**kwargs)

    def setInputCols(self, *value):
        if len(value) == 1 and type(value[0]) == list:
            return self._set(inputCols=value[0])
        else:
            return self._set(inputCols=list(value))

    def setOutputCols(self, *value):
        if len(value) == 1 and type(value[0]) == list:
            return self._set(outputCols=value[0])
        else:
            return self._set(outputCols=list(value))

    def setValueSplitSymbol(self, value):
        return self._set(valueSplitSymbol=value)

    def setAnnotationSplitSymbol(self, value):
        return self._set(annotationSplitSymbol=value)

    def setCleanAnnotations(self, value):
        return self._set(cleanAnnotations=value)

    def setIncludeMetadata(self, value):
        return self._set(includeMetadata=value)

    def setOutputAsArray(self, value):
        return self._set(outputAsArray=value)

    def setParseEmbeddingsVectors(self, value):
        return self._set(parseEmbeddingsVectors=value)
<|endoftext|>