<commit_before># -*- coding: utf-8 -*-
"""The message module allows for quickly and easily sending emails. For quickly
sending messages consider using the send_message function, however, there is
also the :class:`dyn.mm.message.EMail` class which will give you additional
control over the messages you're sending.
"""
from .errors import DynInvalidArgumentError
from .session import session

__author__ = 'jnappi'

def send_message(from_field, to, subject, cc=None, body=None, html=None,
                 replyto=None, xheaders=None):
    """Create and send an email on the fly. For information on the arguments
    accepted by this function see the documentation for
    :class:`dyn.mm.message.EMail`
    """
    EMail(from_field, to, subject, cc, body, html, replyto, xheaders).send()

class EMail(object):
    """Create an and Send it from one of your approved senders"""
    uri = '/send'

    def __init__(self, from_field, to, subject, cc=None, body=None, html=None,
                 replyto=None, xheaders=None):
        """Create a new :class:`EMail` object

        :param from_field: Sender email address - This can either be an email
            address or a properly formatted from header (example: "From Name"
            <example@email.com>). NOTE: The sender must be one of your account's
            Approved Senders
        :param to: A `list` of Address(es) or a single Address that the email
            will be sent to — This/These can either be an email address or a
            properly formatted from header (example: "To Name"
            <example@email.com>). The To field in the email will contain
            all the addresses when it is sent out and will be sent to all the
            addresses.
        :param subject: The subject of the email being sent
        :param cc: Address(es) to copy the email to - This can either be an
            email address or a properly formatted cc header (example: "cc Name"
            <example@email.com>). For multiple addresses, each address must have
            its own 'cc' field. (example: cc = "example1@email.com", cc =
            "example2@email.com").
        :param body: The plain/text version of the email; this field may be
            encoded in Base64 (recommended), quoted-printable, 8-bit, or 7-bit.
        :param html: The text/html version of the email; this field may be
            encoded in 7-bit, 8-bit, quoted-printable, or base64.
        :param replyto: The email address for the recipient to reply to. If left
            blank, defaults to the from address.
        :param xheaders: Any additional custom X-headers to send in the email -
            Pass the X-header's name as the field name and the X-header's value
            as the value (example: x-demonheader=zoom).
        """
        if body is None and html is None:
            raise DynInvalidArgumentError('body and html', (None, None))
        self.from_field = from_field
        self.to = to
        self.subject = subject
        self.cc = cc
        self.body = body
        self.html = html
        self.replyto = replyto
        self.xheaders = xheaders

    def send(self):
        """Send the content of this :class:`Email` object to the provided list
        of recipients.
        """
        d = self.__dict__
        api_args = {x: d[x] for x in d if d[x] is not None and
                    not hasattr(d[x], '__call__')}
        session().execute(self.uri, 'POST', api_args)
<commit_msg>Added __all__ attribute and EMail subclasses to allow for easier templating of emails while sending.<commit_after># -*- coding: utf-8 -*-
"""The message module allows for quickly and easily sending emails. For quickly
sending messages consider using the send_message function, however, there is
also the :class:`dyn.mm.message.EMail` class which will give you additional
control over the messages you're sending.
"""
from .errors import DynInvalidArgumentError
from .session import session
from ..core import cleared_class_dict

__all__ = ['send_message', 'EMail', 'HTMLEMail', 'TemplateEMail',
           'HTMLTemplateEMail']
__author__ = 'jnappi'

def send_message(from_field, to, subject, cc=None, body=None, html=None,
                 replyto=None, xheaders=None):
    """Create and send an email on the fly. For information on the arguments
    accepted by this function see the documentation for
    :class:`~dyn.mm.message.EMail`
    """
    EMail(from_field, to, subject, cc, body, html, replyto, xheaders).send()

class EMail(object):
    """Create an and Send it from one of your approved senders"""
    uri = '/send'

    def __init__(self, from_field, to, subject, cc=None, body=None, html=None,
                 replyto=None, xheaders=None):
        """Create a new :class:`EMail` object

        :param from_field: Sender email address - This can either be an email
            address or a properly formatted from header (example: "From Name"
            <example@email.com>). NOTE: The sender must be one of your account's
            Approved Senders
        :param to: A `list` of Address(es) or a single Address that the email
            will be sent to — This/These can either be an email address or a
            properly formatted from header (example: "To Name"
            <example@email.com>). The To field in the email will contain
            all the addresses when it is sent out and will be sent to all the
            addresses.
        :param subject: The subject of the email being sent
        :param cc: Address(es) to copy the email to - This can either be an
            email address or a properly formatted cc header (example: "cc Name"
            <example@email.com>). For multiple addresses, each address must have
            its own 'cc' field. (example: cc = "example1@email.com", cc =
            "example2@email.com").
        :param body: The plain/text version of the email; this field may be
            encoded in Base64 (recommended), quoted-printable, 8-bit, or 7-bit.
        :param html: The text/html version of the email; this field may be
            encoded in 7-bit, 8-bit, quoted-printable, or base64.
        :param replyto: The email address for the recipient to reply to. If left
            blank, defaults to the from address.
        :param xheaders: Any additional custom X-headers to send in the email -
            Pass the X-header's name as the field name and the X-header's value
            as the value (example: x-demonheader=zoom).
        """
        self.from_field = from_field
        self.to = to
        self.subject = subject
        self.cc = cc
        self.bodytext = body
        self.htmltext = html
        self.replyto = replyto
        self.xheaders = xheaders

    def send(self, content=None):
        """Send the content of this :class:`Email` object to the provided list
        of recipients.

        :param content: The optional content field can be used to overrwrite, or
            to specify the actual content of the body of the message. Note: If
            *content*, this instance's body, and this instance's html fields are
            all *None*, then an
            :exception:`~dyn.mm.errors.DynInvalidArgumentError` will be raised.
        """
        if content is None and self.body is None and self.html is None:
            raise DynInvalidArgumentError('body and html', (None, None))
        api_args = cleared_class_dict(self.__dict__)
        if content is not None:
            api_args['body'] = content
        session().execute(self.uri, 'POST', api_args)

class HTMLEMail(EMail):
    """:class:`~dyn.mm.message.EMail` subclass with an overridden send method
    for specifying html content on the fly
    """
    def send(self, content=None):
        """Send the content of this :class:`Email` object to the provided list
        of recipients.

        :param content: The optional content field can be used to overrwrite, or
            to specify the actual content of the html of the message. Note: If
            *content*, this instance's body, and this instance's html fields are
            all *None*, then an
            :exception:`~dyn.mm.errors.DynInvalidArgumentError` will be raised.
        """
        if content is None and self.body is None and self.html is None:
            raise DynInvalidArgumentError('body and html', (None, None))
        api_args = cleared_class_dict(self.__dict__)
        if content is not None:
            api_args['html'] = content
        session().execute(self.uri, 'POST', api_args)

class TemplateEMail(EMail):
    """:class:`~dyn.mm.message.EMail` subclass which treats it's body field as
    a template before sending.
    """
    def send(self, formatters=None):
        """Send the content of this :class:`Email` object to the provided list
        of recipients.

        :param formatters: The optional content field can be used to overrwrite, or
            to specify the actual content of the html of the message. Note: If
            *content*, this instance's body, and this instance's html fields are
            all *None*, then an
            :exception:`~dyn.mm.errors.DynInvalidArgumentError` will be raised.
        """
        if formatters is None:
            raise DynInvalidArgumentError('send content', None)

        if self.body is None and self.html is None:
            raise DynInvalidArgumentError('body and html', (None, None))

        for formatter in formatters:
            super(TemplateEMail, self).send(self.body % formatter)

class HTMLTemplateEMail(HTMLEMail):
    """:class:`~dyn.mm.message.EMail` subclass which treats it's body field as
    a template before sending.
    """
    def send(self, formatters=None):
        """Send the content of this :class:`Email` object to the provided list
        of recipients.

        :param formatters: The optional content field can be used to overrwrite, or
            to specify the actual content of the html of the message. Note: If
            *content*, this instance's body, and this instance's html fields are
            all *None*, then an
            :exception:`~dyn.mm.errors.DynInvalidArgumentError` will be raised.
        """
        if formatters is None:
            raise DynInvalidArgumentError('send content', None)

        if self.body is None and self.html is None:
            raise DynInvalidArgumentError('body and html', (None, None))

        for formatter in formatters:
            super(HTMLTemplateEMail, self).send(self.html % formatter)
<|endoftext|>


<commit_before>"""Test AdaNet estimator cluster training support.

Copyright 2019 The AdaNet Authors. 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

    https://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.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import copy
import itertools
import json
import os
import shutil
import socket
import subprocess
import sys
import time

from absl import flags
from absl import logging
from absl.testing import parameterized
from adanet.core.timer import _CountDownTimer
import tensorflow as tf

# Maximum number of characters to log per process.
# NOTE: The full process output is written to disk.
MAX_OUTPUT_CHARS = 15000

# A process. name is a string identifying the process in logs. stderr is a file
# object of the process's stderr.
_ProcessInfo = collections.namedtuple("_ProcessInfo",
                                      ["name", "popen", "stderr"])

def _create_task_process(task_type, task_index, estimator_type,
                         placement_strategy, tf_config, model_dir):
  """Creates a process for a single estimator task.

  Args:
    task_type: 'chief', 'worker' or 'ps'.
    task_index: The index of the task within the cluster.
    estimator_type: The estimator type to train. 'estimator' or 'autoensemble'.
    placement_strategy: The distributed placement strategy.
    tf_config: Dictionary representation of the TF_CONFIG environment variable.
      This method creates a copy as to not mutate the input dict.
    model_dir: The Estimator's model directory.

  Returns:
    A _ProcessInfo namedtuple of the running process. The stderr field of this
      tuple must be closed by the caller once the process ends.
  """

  process_name = "%s_%s" % (task_type, task_index)
  args = ["python", "adanet/core/estimator_distributed_test_runner.py"]
  args.append("--estimator_type={}".format(estimator_type))
  args.append("--placement_strategy={}".format(placement_strategy))
  # Log everything to stderr.
  args.append("--stderrthreshold=info")
  args.append("--model_dir={}".format(model_dir))
  logging.info("Spawning %s process: %s", process_name, " ".join(args))
  stderr_filename = os.path.join(model_dir, "%s_stderr.txt" % process_name)
  logging.info("Logging to %s", model_dir)
  stderr_file = open(stderr_filename, "w+")
  tf_config = copy.deepcopy(tf_config)
  tf_config["task"]["type"] = task_type
  tf_config["task"]["index"] = task_index
  json_tf_config = json.dumps(tf_config)
  env = os.environ.copy()
  # Allow stderr to be viewed before the process ends.
  env["PYTHONUNBUFFERED"] = "1"
  env["TF_CPP_MIN_LOG_LEVEL"] = "0"
  env["TF_CONFIG"] = json_tf_config
  # Change gRPC polling strategy to prevent blocking forever.
  # See https://github.com/tensorflow/tensorflow/issues/17852.
  env["GRPC_POLL_STRATEGY"] = "poll"
  popen = subprocess.Popen(args, stderr=stderr_file, env=env)
  return _ProcessInfo(process_name, popen, stderr_file)

def _pick_unused_port():
  """Returns a free port on localhost."""

  for family in (socket.AF_INET6, socket.AF_INET):
    try:
      sock = socket.socket(family, socket.SOCK_STREAM)
      sock.bind(("", 0))  # Passing port '0' binds to a free port on localhost.
      port = sock.getsockname()[1]
      sock.close()
      return port
    except socket.error:
      continue
  raise socket.error

class EstimatorDistributedTrainingTest(parameterized.TestCase,
                                       tf.test.TestCase):
  """Tests distributed training."""

  def setUp(self):
    super(EstimatorDistributedTrainingTest, self).setUp()
    flags.FLAGS(sys.argv)
    # Setup and cleanup test directory.
    self.test_subdirectory = os.path.join(flags.FLAGS.test_tmpdir, self.id())
    shutil.rmtree(self.test_subdirectory, ignore_errors=True)
    os.makedirs(self.test_subdirectory)

  def _wait_for_processes(self, wait_processes, kill_processes, timeout_secs):
    """Waits until all `wait_processes` finish, then kills `kill_processes`.

    Fails an assert if a process in `wait_processes` finishes unsuccessfully.
    The processes in `kill_processes` are assumed to never finish so they are
    killed.

    Args:
      wait_processes: A list of _ProcessInfo tuples. This function will wait for
        each to finish.
      kill_processes: A list of _ProcessInfo tuples. Each will be killed once
        every process in `wait_processes` is finished.
      timeout_secs: Seconds to wait before timing out and terminating processes.

    Returns:
      A list of strings, each which is a string of the stderr of a wait process.

    Raises:
      Exception: When waiting for tasks to finish times out.
    """

    timer = _CountDownTimer(timeout_secs)
    wait_process_stderrs = [None] * len(wait_processes)
    finished_wait_processes = set()
    while len(finished_wait_processes) < len(wait_processes):
      if timer.secs_remaining() == 0:
        logging.error("Timed out! Outputting logs of unfinished processes:")
        for i, wait_process in enumerate(wait_processes):
          if i in finished_wait_processes:
            continue
          wait_process.stderr.seek(0)
          wait_process_stderrs[i] = wait_process.stderr.read()
          logging.info("stderr for incomplete %s (last %d chars): %s\n",
                       wait_process.name, MAX_OUTPUT_CHARS,
                       wait_process.stderr.read()[-MAX_OUTPUT_CHARS:])
        raise Exception("Timed out waiting for tasks to complete.")
      for i, wait_process in enumerate(wait_processes):
        if i in finished_wait_processes:
          continue
        ret_code = wait_process.popen.poll()
        if ret_code is None:
          continue
        logging.info("%s finished", wait_process.name)
        wait_process.stderr.seek(0)
        wait_process_stderrs[i] = wait_process.stderr.read()
        logging.info("stderr for %s (last %d chars): %s\n", wait_process.name,
                     MAX_OUTPUT_CHARS,
                     wait_process_stderrs[i][-MAX_OUTPUT_CHARS:])
        self.assertEqual(0, ret_code)
        finished_wait_processes.add(i)
      for kill_process in kill_processes:
        ret_code = kill_process.popen.poll()
        # Kill processes should not end until we kill them.
        # If it returns early, note the return code.
        self.assertIsNone(ret_code)
      # Delay between polling loops.
      time.sleep(0.25)
    logging.info("All wait processes finished")
    for i, kill_process in enumerate(kill_processes):
      # Kill each kill process.
      kill_process.popen.kill()
      kill_process.popen.wait()
      kill_process.stderr.seek(0)
      logging.info("stderr for %s (last %d chars): %s\n", kill_process.name,
                   MAX_OUTPUT_CHARS,
                   kill_process.stderr.read()[-MAX_OUTPUT_CHARS:])
    return wait_process_stderrs

  # pylint: disable=g-complex-comprehension
  @parameterized.named_parameters(
      itertools.chain(*[[
          {
              "testcase_name": "{}_one_worker".format(placement),
              "placement_strategy": placement,
              "num_workers": 1,
              "num_ps": 0,
          },
          {
              "testcase_name": "{}_one_worker_one_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 1,
              "num_ps": 1,
          },
          {
              "testcase_name": "{}_two_workers_one_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 2,
              "num_ps": 1,
          },
          {
              "testcase_name": "{}_three_workers_three_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 3,
              "num_ps": 3,
          },
          {
              "testcase_name": "{}_five_workers_three_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 5,
              "num_ps": 3,
          },
          {
              "testcase_name":
                  "autoensemble_{}_five_workers_three_ps".format(placement),
              "estimator":
                  "autoensemble",
              "placement_strategy":
                  placement,
              "num_workers":
                  5,
              "num_ps":
                  3,
          },
          {
              "testcase_name":
                  "autoensemble_trees_multiclass_{}_five_workers_three_ps"
                  .format(placement),
              "estimator":
                  "autoensemble_trees_multiclass",
              "placement_strategy":
                  placement,
              "num_workers":
                  5,
              "num_ps":
                  3,
          },
      ] for placement in ["replication", "round_robin"]]))
  # pylint: enable=g-complex-comprehension
  def test_distributed_training(self,
                                num_workers,
                                num_ps,
                                placement_strategy,
                                estimator="estimator"):
    """Uses multiprocessing to simulate a distributed training environment."""

    # Inspired by `tf.test.create_local_cluster`.
    worker_ports = [_pick_unused_port() for _ in range(num_workers)]
    ps_ports = [_pick_unused_port() for _ in range(num_ps)]
    ws_targets = ["localhost:%s" % port for port in worker_ports]
    ps_targets = ["localhost:%s" % port for port in ps_ports]

    # For details see:
    # https://www.tensorflow.org/api_docs/python/tf/estimator/train_and_evaluate
    tf_config = {
        "cluster": {
            # The chief is always worker 0.
            "chief": [ws_targets[0]],
        },
        "task": {
            "type": "chief",
            "index": 0
        },
    }

    # The chief is already worker 0.
    if len(ws_targets) > 1:
      tf_config["cluster"]["worker"] = ws_targets[1:]
    if ps_targets:
      tf_config["cluster"]["ps"] = ps_targets

    worker_processes = []
    ps_processes = []
    evaluator_processes = []

    model_dir = self.test_subdirectory

    # Chief
    worker_processes.append(
        _create_task_process("chief", 0, estimator, placement_strategy,
                             tf_config, model_dir))
    # Workers
    for i in range(len(ws_targets[1:])):
      worker_processes.append(
          _create_task_process("worker", i, estimator, placement_strategy,
                               tf_config, model_dir))
    # Parameter Servers (PS)
    for i in range(len(ps_targets)):
      ps_processes.append(
          _create_task_process("ps", i, estimator, placement_strategy,
                               tf_config, model_dir))
    # Evaluator
    evaluator_processes.append(
        _create_task_process("evaluator", 0, estimator, placement_strategy,
                             tf_config, model_dir))

    # Run processes.
    try:
      # NOTE: Parameter servers do not shut down on their own.
      self._wait_for_processes(
          worker_processes + evaluator_processes,
          kill_processes=ps_processes,
          timeout_secs=500)
    finally:
      for process in worker_processes + ps_processes + evaluator_processes:
        try:
          process.popen.kill()
        except OSError:
          pass  # It's OK (and expected) if the process already exited.
        process.stderr.close()

if __name__ == "__main__":
  tf.test.main()
<commit_msg>Add/fix logging for estimator_distributed_test.py.<commit_after>"""Test AdaNet estimator cluster training support.

Copyright 2019 The AdaNet Authors. 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

    https://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.
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections
import copy
import itertools
import json
import os
import shutil
import socket
import subprocess
import sys
import time

from absl import flags
from absl import logging
from absl.testing import parameterized
from adanet.core.timer import _CountDownTimer
import tensorflow as tf

# Maximum number of characters to log per process.
# NOTE: The full process output is written to disk.
MAX_OUTPUT_CHARS = 15000

# A process. name is a string identifying the process in logs. stderr is a file
# object of the process's stderr.
_ProcessInfo = collections.namedtuple("_ProcessInfo",
                                      ["name", "popen", "stderr"])

def _create_task_process(task_type, task_index, estimator_type,
                         placement_strategy, tf_config, model_dir):
  """Creates a process for a single estimator task.

  Args:
    task_type: 'chief', 'worker' or 'ps'.
    task_index: The index of the task within the cluster.
    estimator_type: The estimator type to train. 'estimator' or 'autoensemble'.
    placement_strategy: The distributed placement strategy.
    tf_config: Dictionary representation of the TF_CONFIG environment variable.
      This method creates a copy as to not mutate the input dict.
    model_dir: The Estimator's model directory.

  Returns:
    A _ProcessInfo namedtuple of the running process. The stderr field of this
      tuple must be closed by the caller once the process ends.
  """

  process_name = "%s_%s" % (task_type, task_index)
  args = ["python", "adanet/core/estimator_distributed_test_runner.py"]
  args.append("--estimator_type={}".format(estimator_type))
  args.append("--placement_strategy={}".format(placement_strategy))
  # Log everything to stderr.
  args.append("--stderrthreshold=info")
  args.append("--model_dir={}".format(model_dir))
  logging.info("Spawning %s process: %s", process_name, " ".join(args))
  stderr_filename = os.path.join(model_dir, "%s_stderr.txt" % process_name)
  logging.info("Logging to %s", model_dir)
  stderr_file = open(stderr_filename, "w+")
  tf_config = copy.deepcopy(tf_config)
  tf_config["task"]["type"] = task_type
  tf_config["task"]["index"] = task_index
  json_tf_config = json.dumps(tf_config)
  env = os.environ.copy()
  # Allow stderr to be viewed before the process ends.
  env["PYTHONUNBUFFERED"] = "1"
  env["TF_CPP_MIN_LOG_LEVEL"] = "0"
  env["TF_CONFIG"] = json_tf_config
  # Change gRPC polling strategy to prevent blocking forever.
  # See https://github.com/tensorflow/tensorflow/issues/17852.
  env["GRPC_POLL_STRATEGY"] = "poll"
  popen = subprocess.Popen(args, stderr=stderr_file, env=env)
  return _ProcessInfo(process_name, popen, stderr_file)

def _pick_unused_port():
  """Returns a free port on localhost."""

  for family in (socket.AF_INET6, socket.AF_INET):
    try:
      sock = socket.socket(family, socket.SOCK_STREAM)
      sock.bind(("", 0))  # Passing port '0' binds to a free port on localhost.
      port = sock.getsockname()[1]
      sock.close()
      return port
    except socket.error:
      continue
  raise socket.error

class EstimatorDistributedTrainingTest(parameterized.TestCase,
                                       tf.test.TestCase):
  """Tests distributed training."""

  def setUp(self):
    super(EstimatorDistributedTrainingTest, self).setUp()
    flags.FLAGS(sys.argv)
    # Setup and cleanup test directory.
    self.test_subdirectory = os.path.join(flags.FLAGS.test_tmpdir, self.id())
    shutil.rmtree(self.test_subdirectory, ignore_errors=True)
    os.makedirs(self.test_subdirectory)

  def _wait_for_processes(self, wait_processes, kill_processes, timeout_secs):
    """Waits until all `wait_processes` finish, then kills `kill_processes`.

    Fails an assert if a process in `wait_processes` finishes unsuccessfully.
    The processes in `kill_processes` are assumed to never finish so they are
    killed.

    Args:
      wait_processes: A list of _ProcessInfo tuples. This function will wait for
        each to finish.
      kill_processes: A list of _ProcessInfo tuples. Each will be killed once
        every process in `wait_processes` is finished.
      timeout_secs: Seconds to wait before timing out and terminating processes.

    Returns:
      A list of strings, each which is a string of the stderr of a wait process.

    Raises:
      Exception: When waiting for tasks to finish times out.
    """

    timer = _CountDownTimer(timeout_secs)
    wait_process_stderrs = [None] * len(wait_processes)
    finished_wait_processes = set()
    poll_count = {wait_process: 0.0 for wait_process in wait_processes}

    while len(finished_wait_processes) < len(wait_processes):
      if timer.secs_remaining() == 0:
        logging.error("Timed out! Outputting logs of unfinished processes:")
        for i, wait_process in enumerate(wait_processes):
          if i in finished_wait_processes:
            continue
          wait_process.stderr.seek(0)
          wait_process_stderrs[i] = wait_process.stderr.read()
          logging.info("stderr for incomplete %s (last %d chars): %s\n",
                       wait_process.name, MAX_OUTPUT_CHARS,
                       wait_process_stderrs[i][-MAX_OUTPUT_CHARS:])
        raise Exception("Timed out waiting for tasks to complete.")
      for i, wait_process in enumerate(wait_processes):
        if i in finished_wait_processes:
          continue
        ret_code = wait_process.popen.poll()
        if ret_code is None:
          poll_count[wait_process] += 0.25
          if ((poll_count[wait_process] / 10.) -
              int(poll_count[wait_process] / 10.)) == 0:
            logging.info("%d secs has elapsed for %s", poll_count[wait_process],
                         wait_process.name)
          continue
        logging.info("%s finished", wait_process.name)
        wait_process.stderr.seek(0)
        wait_process_stderrs[i] = wait_process.stderr.read()
        logging.info("stderr for %s (last %d chars): %s\n", wait_process.name,
                     MAX_OUTPUT_CHARS,
                     wait_process_stderrs[i][-MAX_OUTPUT_CHARS:])
        self.assertEqual(0, ret_code)
        finished_wait_processes.add(i)
      for kill_process in kill_processes:
        ret_code = kill_process.popen.poll()
        # Kill processes should not end until we kill them.
        # If it returns early, note the return code.
        self.assertIsNone(ret_code)
      # Delay between polling loops.
      time.sleep(0.25)
    logging.info("All wait processes finished")
    for i, kill_process in enumerate(kill_processes):
      # Kill each kill process.
      kill_process.popen.kill()
      kill_process.popen.wait()
      kill_process.stderr.seek(0)
      logging.info("stderr for %s (last %d chars): %s\n", kill_process.name,
                   MAX_OUTPUT_CHARS,
                   kill_process.stderr.read()[-MAX_OUTPUT_CHARS:])
    return wait_process_stderrs

  # pylint: disable=g-complex-comprehension
  @parameterized.named_parameters(
      itertools.chain(*[[
          {
              "testcase_name": "{}_one_worker".format(placement),
              "placement_strategy": placement,
              "num_workers": 1,
              "num_ps": 0,
          },
          {
              "testcase_name": "{}_one_worker_one_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 1,
              "num_ps": 1,
          },
          {
              "testcase_name": "{}_two_workers_one_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 2,
              "num_ps": 1,
          },
          {
              "testcase_name": "{}_three_workers_three_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 3,
              "num_ps": 3,
          },
          {
              "testcase_name": "{}_five_workers_three_ps".format(placement),
              "placement_strategy": placement,
              "num_workers": 5,
              "num_ps": 3,
          },
          {
              "testcase_name":
                  "autoensemble_{}_five_workers_three_ps".format(placement),
              "estimator":
                  "autoensemble",
              "placement_strategy":
                  placement,
              "num_workers":
                  5,
              "num_ps":
                  3,
          },
          {
              "testcase_name":
                  "autoensemble_trees_multiclass_{}_five_workers_three_ps"
                  .format(placement),
              "estimator":
                  "autoensemble_trees_multiclass",
              "placement_strategy":
                  placement,
              "num_workers":
                  5,
              "num_ps":
                  3,
          },
      ] for placement in ["replication", "round_robin"]]))
  # pylint: enable=g-complex-comprehension
  def test_distributed_training(self,
                                num_workers,
                                num_ps,
                                placement_strategy,
                                estimator="estimator"):
    """Uses multiprocessing to simulate a distributed training environment."""

    # Inspired by `tf.test.create_local_cluster`.
    worker_ports = [_pick_unused_port() for _ in range(num_workers)]
    ps_ports = [_pick_unused_port() for _ in range(num_ps)]
    ws_targets = ["localhost:%s" % port for port in worker_ports]
    ps_targets = ["localhost:%s" % port for port in ps_ports]

    # For details see:
    # https://www.tensorflow.org/api_docs/python/tf/estimator/train_and_evaluate
    tf_config = {
        "cluster": {
            # The chief is always worker 0.
            "chief": [ws_targets[0]],
        },
        "task": {
            "type": "chief",
            "index": 0
        },
    }

    # The chief is already worker 0.
    if len(ws_targets) > 1:
      tf_config["cluster"]["worker"] = ws_targets[1:]
    if ps_targets:
      tf_config["cluster"]["ps"] = ps_targets

    worker_processes = []
    ps_processes = []
    evaluator_processes = []

    model_dir = self.test_subdirectory

    # Chief
    worker_processes.append(
        _create_task_process("chief", 0, estimator, placement_strategy,
                             tf_config, model_dir))
    # Workers
    for i in range(len(ws_targets[1:])):
      worker_processes.append(
          _create_task_process("worker", i, estimator, placement_strategy,
                               tf_config, model_dir))
    # Parameter Servers (PS)
    for i in range(len(ps_targets)):
      ps_processes.append(
          _create_task_process("ps", i, estimator, placement_strategy,
                               tf_config, model_dir))
    # Evaluator
    evaluator_processes.append(
        _create_task_process("evaluator", 0, estimator, placement_strategy,
                             tf_config, model_dir))

    # Run processes.
    try:
      # NOTE: Parameter servers do not shut down on their own.
      self._wait_for_processes(
          worker_processes + evaluator_processes,
          kill_processes=ps_processes,
          timeout_secs=500)
    finally:
      for process in worker_processes + ps_processes + evaluator_processes:
        try:
          process.popen.kill()
        except OSError:
          pass  # It's OK (and expected) if the process already exited.
        process.stderr.close()

if __name__ == "__main__":
  tf.test.main()
<|endoftext|>


<commit_before># Copyright (c) 2016-2017, Neil Booth
#
# All rights reserved.
#
# See the file "LICENCE" for information about the copyright
# and warranty status of this software.

'''Class for handling asynchronous connections to a blockchain
daemon.'''

import asyncio
import json
import time
from calendar import timegm
from struct import pack
from time import strptime

import aiohttp

from electrumx.lib.util import int_to_varint, hex_to_bytes, class_logger
from electrumx.lib.hash import hex_str_to_hash
from aiorpcx import JSONRPC

class DaemonError(Exception):
    '''Raised when the daemon returns an error in its results.'''

class Daemon(object):
    '''Handles connections to a daemon at the given URL.'''

    WARMING_UP = -28
    RPC_MISC_ERROR = -1

    class DaemonWarmingUpError(Exception):
        '''Raised when the daemon returns an error in its results.'''

    def __init__(self, env):
        self.logger = class_logger(__name__, self.__class__.__name__)
        self.coin = env.coin
        self.set_urls(env.coin.daemon_urls(env.daemon_url))
        self._height = None
        self._mempool_hashes = set()
        self.mempool_refresh_event = asyncio.Event()
        # Limit concurrent RPC calls to this number.
        # See DEFAULT_HTTP_WORKQUEUE in bitcoind, which is typically 16
        self.workqueue_semaphore = asyncio.Semaphore(value=10)
        self.down = False
        self.last_error_time = 0
        self.req_id = 0
        # assignment of asyncio.TimeoutError are essentially ignored
        if aiohttp.__version__.startswith('1.'):
            self.ClientHttpProcessingError = aiohttp.ClientHttpProcessingError
            self.ClientPayloadError = asyncio.TimeoutError
        else:
            self.ClientHttpProcessingError = asyncio.TimeoutError
            self.ClientPayloadError = aiohttp.ClientPayloadError
        self._available_rpcs = {}  # caches results for _is_rpc_available()

    def next_req_id(self):
        '''Retrns the next request ID.'''
        self.req_id += 1
        return self.req_id

    def set_urls(self, urls):
        '''Set the URLS to the given list, and switch to the first one.'''
        if not urls:
            raise DaemonError('no daemon URLs provided')
        self.urls = urls
        self.url_index = 0
        for n, url in enumerate(urls):
            self.logger.info('daemon #{:d} at {}{}'
                             .format(n + 1, self.logged_url(url),
                                     '' if n else ' (current)'))

    def url(self):
        '''Returns the current daemon URL.'''
        return self.urls[self.url_index]

    def failover(self):
        '''Call to fail-over to the next daemon URL.

        Returns False if there is only one, otherwise True.
        '''
        if len(self.urls) > 1:
            self.url_index = (self.url_index + 1) % len(self.urls)
            self.logger.info('failing over to {}'.format(self.logged_url()))
            return True
        return False

    def client_session(self):
        '''An aiohttp client session.'''
        return aiohttp.ClientSession()

    async def _send_data(self, data):
        async with self.workqueue_semaphore:
            async with self.client_session() as session:
                async with session.post(self.url(), data=data) as resp:
                    # If bitcoind can't find a tx, for some reason
                    # it returns 500 but fills out the JSON.
                    # Should still return 200 IMO.
                    if resp.status in (200, 404, 500):
                        return await resp.json()
                    return (resp.status, resp.reason)

    async def _send(self, payload, processor):
        '''Send a payload to be converted to JSON.

        Handles temporary connection issues.  Daemon reponse errors
        are raise through DaemonError.
        '''
        def log_error(error):
            self.down = True
            now = time.time()
            prior_time = self.last_error_time
            if now - prior_time > 60:
                self.last_error_time = now
                if prior_time and self.failover():
                    secs = 0
                else:
                    self.logger.error('{}  Retrying occasionally...'
                                      .format(error))

        data = json.dumps(payload)
        secs = 1
        max_secs = 4
        while True:
            try:
                result = await self._send_data(data)
                if not isinstance(result, tuple):
                    result = processor(result)
                    if self.down:
                        self.down = False
                        self.last_error_time = 0
                        self.logger.info('connection restored')
                    return result
                log_error('HTTP error code {:d}: {}'
                          .format(result[0], result[1]))
            except asyncio.TimeoutError:
                log_error('timeout error.')
            except aiohttp.ServerDisconnectedError:
                log_error('disconnected.')
            except self.ClientHttpProcessingError:
                log_error('HTTP error.')
            except self.ClientPayloadError:
                log_error('payload encoding error.')
            except aiohttp.ClientConnectionError:
                log_error('connection problem - is your daemon running?')
            except self.DaemonWarmingUpError:
                log_error('starting up checking blocks.')
            except (asyncio.CancelledError, DaemonError):
                raise
            except Exception as e:
                self.logger.exception(f'uncaught exception: {e}')

            await asyncio.sleep(secs)
            secs = min(max_secs, secs * 2, 1)

    def logged_url(self, url=None):
        '''The host and port part, for logging.'''
        url = url or self.url()
        return url[url.rindex('@') + 1:]

    async def _send_single(self, method, params=None):
        '''Send a single request to the daemon.'''
        def processor(result):
            err = result['error']
            if not err:
                return result['result']
            if err.get('code') == self.WARMING_UP:
                raise self.DaemonWarmingUpError
            raise DaemonError(err)

        payload = {'method': method, 'id': self.next_req_id()}
        if params:
            payload['params'] = params
        return await self._send(payload, processor)

    async def _send_vector(self, method, params_iterable, replace_errs=False):
        '''Send several requests of the same method.

        The result will be an array of the same length as params_iterable.
        If replace_errs is true, any item with an error is returned as None,
        otherwise an exception is raised.'''
        def processor(result):
            errs = [item['error'] for item in result if item['error']]
            if any(err.get('code') == self.WARMING_UP for err in errs):
                raise self.DaemonWarmingUpError
            if not errs or replace_errs:
                return [item['result'] for item in result]
            raise DaemonError(errs)

        payload = [{'method': method, 'params': p, 'id': self.next_req_id()}
                   for p in params_iterable]
        if payload:
            return await self._send(payload, processor)
        return []

    async def _is_rpc_available(self, method):
        '''Return whether given RPC method is available in the daemon.

        Results are cached and the daemon will generally not be queried with
        the same method more than once.'''
        available = self._available_rpcs.get(method, None)
        if available is None:
            try:
                await self._send_single(method)
                available = True
            except DaemonError as e:
                err = e.args[0]
                error_code = err.get("code")
                if error_code == JSONRPC.METHOD_NOT_FOUND:
                    available = False
                elif error_code == self.RPC_MISC_ERROR:
                    # method found but exception was thrown in command handling
                    # probably because we did not provide arguments
                    available = True
                else:
                    self.logger.warning('error (code {:d}: {}) when testing '
                                        'RPC availability of method {}'
                                        .format(error_code, err.get("message"),
                                                method))
                    available = False
            self._available_rpcs[method] = available
        return available

    async def block_hex_hashes(self, first, count):
        '''Return the hex hashes of count block starting at height first.'''
        params_iterable = ((h, ) for h in range(first, first + count))
        return await self._send_vector('getblockhash', params_iterable)

    async def deserialised_block(self, hex_hash):
        '''Return the deserialised block with the given hex hash.'''
        return await self._send_single('getblock', (hex_hash, True))

    async def raw_blocks(self, hex_hashes):
        '''Return the raw binary blocks with the given hex hashes.'''
        params_iterable = ((h, False) for h in hex_hashes)
        blocks = await self._send_vector('getblock', params_iterable)
        # Convert hex string to bytes
        return [hex_to_bytes(block) for block in blocks]

    async def mempool_hashes(self):
        '''Update our record of the daemon's mempool hashes.'''
        return await self._send_single('getrawmempool')

    async def estimatefee(self, params):
        '''Return the fee estimate for the given parameters.'''
        if await self._is_rpc_available('estimatesmartfee'):
            estimate = await self._send_single('estimatesmartfee', params)
            return estimate.get('feerate', -1)
        return await self._send_single('estimatefee', params)

    async def getnetworkinfo(self):
        '''Return the result of the 'getnetworkinfo' RPC call.'''
        return await self._send_single('getnetworkinfo')

    async def relayfee(self):
        '''The minimum fee a low-priority tx must pay in order to be accepted
        to the daemon's memory pool.'''
        network_info = await self.getnetworkinfo()
        return network_info['relayfee']

    async def getrawtransaction(self, hex_hash, verbose=False):
        '''Return the serialized raw transaction with the given hash.'''
        # Cast to int because some coin daemons are old and require it
        return await self._send_single('getrawtransaction',
                                       (hex_hash, int(verbose)))

    async def getrawtransactions(self, hex_hashes, replace_errs=True):
        '''Return the serialized raw transactions with the given hashes.

        Replaces errors with None by default.'''
        params_iterable = ((hex_hash, False) for hex_hash in hex_hashes)
        txs = await self._send_vector('getrawtransaction', params_iterable,
                                      replace_errs=replace_errs)
        # Convert hex strings to bytes
        return [hex_to_bytes(tx) if tx else None for tx in txs]

    async def sendrawtransaction(self, params):
        '''Broadcast a transaction to the network.'''
        return await self._send_single('sendrawtransaction', params)

    async def height(self, mempool=False):
        '''Query the daemon for its current height.'''
        self._height = await self._send_single('getblockcount')
        if mempool:
            self._mempool_hashes = set(await self.mempool_hashes())
            self.mempool_refresh_event.set()
        return self._height

    def cached_mempool_hashes(self):
        '''Return the cached mempool hashes.'''
        return self._mempool_hashes

    def cached_height(self):
        '''Return the cached daemon height.

        If the daemon has not been queried yet this returns None.'''
        return self._height

class DashDaemon(Daemon):

    async def masternode_broadcast(self, params):
        '''Broadcast a transaction to the network.'''
        return await self._send_single('masternodebroadcast', params)

    async def masternode_list(self, params):
        '''Return the masternode status.'''
        return await self._send_single('masternodelist', params)

class FakeEstimateFeeDaemon(Daemon):
    '''Daemon that simulates estimatefee and relayfee RPC calls. Coin that
    wants to use this daemon must define ESTIMATE_FEE & RELAY_FEE'''

    async def estimatefee(self, params):
        '''Return the fee estimate for the given parameters.'''
        return self.coin.ESTIMATE_FEE

    async def relayfee(self):
        '''The minimum fee a low-priority tx must pay in order to be accepted
        to the daemon's memory pool.'''
        return self.coin.RELAY_FEE

class LegacyRPCDaemon(Daemon):
    '''Handles connections to a daemon at the given URL.

    This class is useful for daemons that don't have the new 'getblock'
    RPC call that returns the block in hex, the workaround is to manually
    recreate the block bytes. The recreated block bytes may not be the exact
    as in the underlying blockchain but it is good enough for our indexing
    purposes.'''

    async def raw_blocks(self, hex_hashes):
        '''Return the raw binary blocks with the given hex hashes.'''
        params_iterable = ((h, ) for h in hex_hashes)
        block_info = await self._send_vector('getblock', params_iterable)

        blocks = []
        for i in block_info:
            raw_block = await self.make_raw_block(i)
            blocks.append(raw_block)

        # Convert hex string to bytes
        return blocks

    async def make_raw_header(self, b):
        pbh = b.get('previousblockhash')
        if pbh is None:
            pbh = '0' * 64
        return b''.join([
            pack('<L', b.get('version')),
            hex_str_to_hash(pbh),
            hex_str_to_hash(b.get('merkleroot')),
            pack('<L', self.timestamp_safe(b['time'])),
            pack('<L', int(b.get('bits'), 16)),
            pack('<L', int(b.get('nonce')))
        ])

    async def make_raw_block(self, b):
        '''Construct a raw block'''

        header = await self.make_raw_header(b)

        transactions = []
        if b.get('height') > 0:
            transactions = await self.getrawtransactions(b.get('tx'), False)

        raw_block = header
        num_txs = len(transactions)
        if num_txs > 0:
            raw_block += int_to_varint(num_txs)
            raw_block += b''.join(transactions)
        else:
            raw_block += b'\x00'

        return raw_block

    def timestamp_safe(self, t):
        if isinstance(t, int):
            return t
        return timegm(strptime(t, "%Y-%m-%d %H:%M:%S %Z"))
<commit_msg>Partial reversion of 2cb89814b6c8df2d2e19f0a2192c2c7216964140<commit_after># Copyright (c) 2016-2017, Neil Booth
#
# All rights reserved.
#
# See the file "LICENCE" for information about the copyright
# and warranty status of this software.

'''Class for handling asynchronous connections to a blockchain
daemon.'''

import asyncio
import json
import time
from calendar import timegm
from struct import pack
from time import strptime

import aiohttp

from electrumx.lib.util import int_to_varint, hex_to_bytes, class_logger
from electrumx.lib.hash import hex_str_to_hash
from aiorpcx import JSONRPC

class DaemonError(Exception):
    '''Raised when the daemon returns an error in its results.'''

class Daemon(object):
    '''Handles connections to a daemon at the given URL.'''

    WARMING_UP = -28
    RPC_MISC_ERROR = -1

    class DaemonWarmingUpError(Exception):
        '''Raised when the daemon returns an error in its results.'''

    def __init__(self, env):
        self.logger = class_logger(__name__, self.__class__.__name__)
        self.coin = env.coin
        self.set_urls(env.coin.daemon_urls(env.daemon_url))
        self._height = None
        self._mempool_hashes = set()
        self.mempool_refresh_event = asyncio.Event()
        # Limit concurrent RPC calls to this number.
        # See DEFAULT_HTTP_WORKQUEUE in bitcoind, which is typically 16
        self.workqueue_semaphore = asyncio.Semaphore(value=10)
        self.down = False
        self.last_error_time = 0
        self.req_id = 0
        # assignment of asyncio.TimeoutError are essentially ignored
        if aiohttp.__version__.startswith('1.'):
            self.ClientHttpProcessingError = aiohttp.ClientHttpProcessingError
            self.ClientPayloadError = asyncio.TimeoutError
        else:
            self.ClientHttpProcessingError = asyncio.TimeoutError
            self.ClientPayloadError = aiohttp.ClientPayloadError
        self._available_rpcs = {}  # caches results for _is_rpc_available()

    def next_req_id(self):
        '''Retrns the next request ID.'''
        self.req_id += 1
        return self.req_id

    def set_urls(self, urls):
        '''Set the URLS to the given list, and switch to the first one.'''
        if not urls:
            raise DaemonError('no daemon URLs provided')
        self.urls = urls
        self.url_index = 0
        for n, url in enumerate(urls):
            self.logger.info('daemon #{:d} at {}{}'
                             .format(n + 1, self.logged_url(url),
                                     '' if n else ' (current)'))

    def url(self):
        '''Returns the current daemon URL.'''
        return self.urls[self.url_index]

    def failover(self):
        '''Call to fail-over to the next daemon URL.

        Returns False if there is only one, otherwise True.
        '''
        if len(self.urls) > 1:
            self.url_index = (self.url_index + 1) % len(self.urls)
            self.logger.info('failing over to {}'.format(self.logged_url()))
            return True
        return False

    def client_session(self):
        '''An aiohttp client session.'''
        return aiohttp.ClientSession()

    async def _send_data(self, data):
        async with self.workqueue_semaphore:
            async with self.client_session() as session:
                async with session.post(self.url(), data=data) as resp:
                    # If bitcoind can't find a tx, for some reason
                    # it returns 500 but fills out the JSON.
                    # Should still return 200 IMO.
                    if resp.status in (200, 404, 500):
                        return await resp.json()
                    return (resp.status, resp.reason)

    async def _send(self, payload, processor):
        '''Send a payload to be converted to JSON.

        Handles temporary connection issues.  Daemon reponse errors
        are raise through DaemonError.
        '''
        def log_error(error):
            self.down = True
            now = time.time()
            prior_time = self.last_error_time
            if now - prior_time > 60:
                self.last_error_time = now
                if prior_time and self.failover():
                    secs = 0
                else:
                    self.logger.error('{}  Retrying occasionally...'
                                      .format(error))

        data = json.dumps(payload)
        secs = 1
        max_secs = 4
        while True:
            try:
                result = await self._send_data(data)
                if not isinstance(result, tuple):
                    result = processor(result)
                    if self.down:
                        self.down = False
                        self.last_error_time = 0
                        self.logger.info('connection restored')
                    return result
                log_error('HTTP error code {:d}: {}'
                          .format(result[0], result[1]))
            except asyncio.TimeoutError:
                log_error('timeout error.')
            except aiohttp.ServerDisconnectedError:
                log_error('disconnected.')
            except self.ClientHttpProcessingError:
                log_error('HTTP error.')
            except self.ClientPayloadError:
                log_error('payload encoding error.')
            except aiohttp.ClientConnectionError:
                log_error('connection problem - is your daemon running?')
            except self.DaemonWarmingUpError:
                log_error('starting up checking blocks.')
            except (asyncio.CancelledError, DaemonError):
                raise
            except Exception as e:
                self.logger.exception(f'uncaught exception: {e}')

            await asyncio.sleep(secs)
            secs = min(max_secs, secs * 2, 1)

    def logged_url(self, url=None):
        '''The host and port part, for logging.'''
        url = url or self.url()
        return url[url.rindex('@') + 1:]

    async def _send_single(self, method, params=None):
        '''Send a single request to the daemon.'''
        def processor(result):
            err = result['error']
            if not err:
                return result['result']
            if err.get('code') == self.WARMING_UP:
                raise self.DaemonWarmingUpError
            raise DaemonError(err)

        payload = {'method': method, 'id': self.next_req_id()}
        if params:
            payload['params'] = params
        return await self._send(payload, processor)

    async def _send_vector(self, method, params_iterable, replace_errs=False):
        '''Send several requests of the same method.

        The result will be an array of the same length as params_iterable.
        If replace_errs is true, any item with an error is returned as None,
        otherwise an exception is raised.'''
        def processor(result):
            errs = [item['error'] for item in result if item['error']]
            if any(err.get('code') == self.WARMING_UP for err in errs):
                raise self.DaemonWarmingUpError
            if not errs or replace_errs:
                return [item['result'] for item in result]
            raise DaemonError(errs)

        payload = [{'method': method, 'params': p, 'id': self.next_req_id()}
                   for p in params_iterable]
        if payload:
            return await self._send(payload, processor)
        return []

    async def _is_rpc_available(self, method):
        '''Return whether given RPC method is available in the daemon.

        Results are cached and the daemon will generally not be queried with
        the same method more than once.'''
        available = self._available_rpcs.get(method, None)
        if available is None:
            try:
                await self._send_single(method)
                available = True
            except DaemonError as e:
                err = e.args[0]
                error_code = err.get("code")
                if error_code == JSONRPC.METHOD_NOT_FOUND:
                    available = False
                elif error_code == self.RPC_MISC_ERROR:
                    # method found but exception was thrown in command handling
                    # probably because we did not provide arguments
                    available = True
                else:
                    self.logger.warning('error (code {:d}: {}) when testing '
                                        'RPC availability of method {}'
                                        .format(error_code, err.get("message"),
                                                method))
                    available = False
            self._available_rpcs[method] = available
        return available

    async def block_hex_hashes(self, first, count):
        '''Return the hex hashes of count block starting at height first.'''
        params_iterable = ((h, ) for h in range(first, first + count))
        return await self._send_vector('getblockhash', params_iterable)

    async def deserialised_block(self, hex_hash):
        '''Return the deserialised block with the given hex hash.'''
        return await self._send_single('getblock', (hex_hash, True))

    async def raw_blocks(self, hex_hashes):
        '''Return the raw binary blocks with the given hex hashes.'''
        params_iterable = ((h, False) for h in hex_hashes)
        blocks = await self._send_vector('getblock', params_iterable)
        # Convert hex string to bytes
        return [hex_to_bytes(block) for block in blocks]

    async def mempool_hashes(self):
        '''Update our record of the daemon's mempool hashes.'''
        return await self._send_single('getrawmempool')

    async def estimatefee(self, params):
        '''Return the fee estimate for the given parameters.'''
        if await self._is_rpc_available('estimatesmartfee'):
            estimate = await self._send_single('estimatesmartfee', params)
            return estimate.get('feerate', -1)
        return await self._send_single('estimatefee', params)

    async def getnetworkinfo(self):
        '''Return the result of the 'getnetworkinfo' RPC call.'''
        return await self._send_single('getnetworkinfo')

    async def relayfee(self):
        '''The minimum fee a low-priority tx must pay in order to be accepted
        to the daemon's memory pool.'''
        network_info = await self.getnetworkinfo()
        return network_info['relayfee']

    async def getrawtransaction(self, hex_hash, verbose=False):
        '''Return the serialized raw transaction with the given hash.'''
        # Cast to int because some coin daemons are old and require it
        return await self._send_single('getrawtransaction',
                                       (hex_hash, int(verbose)))

    async def getrawtransactions(self, hex_hashes, replace_errs=True):
        '''Return the serialized raw transactions with the given hashes.

        Replaces errors with None by default.'''
        params_iterable = ((hex_hash, 0) for hex_hash in hex_hashes)
        txs = await self._send_vector('getrawtransaction', params_iterable,
                                      replace_errs=replace_errs)
        # Convert hex strings to bytes
        return [hex_to_bytes(tx) if tx else None for tx in txs]

    async def sendrawtransaction(self, params):
        '''Broadcast a transaction to the network.'''
        return await self._send_single('sendrawtransaction', params)

    async def height(self, mempool=False):
        '''Query the daemon for its current height.'''
        self._height = await self._send_single('getblockcount')
        if mempool:
            self._mempool_hashes = set(await self.mempool_hashes())
            self.mempool_refresh_event.set()
        return self._height

    def cached_mempool_hashes(self):
        '''Return the cached mempool hashes.'''
        return self._mempool_hashes

    def cached_height(self):
        '''Return the cached daemon height.

        If the daemon has not been queried yet this returns None.'''
        return self._height

class DashDaemon(Daemon):

    async def masternode_broadcast(self, params):
        '''Broadcast a transaction to the network.'''
        return await self._send_single('masternodebroadcast', params)

    async def masternode_list(self, params):
        '''Return the masternode status.'''
        return await self._send_single('masternodelist', params)

class FakeEstimateFeeDaemon(Daemon):
    '''Daemon that simulates estimatefee and relayfee RPC calls. Coin that
    wants to use this daemon must define ESTIMATE_FEE & RELAY_FEE'''

    async def estimatefee(self, params):
        '''Return the fee estimate for the given parameters.'''
        return self.coin.ESTIMATE_FEE

    async def relayfee(self):
        '''The minimum fee a low-priority tx must pay in order to be accepted
        to the daemon's memory pool.'''
        return self.coin.RELAY_FEE

class LegacyRPCDaemon(Daemon):
    '''Handles connections to a daemon at the given URL.

    This class is useful for daemons that don't have the new 'getblock'
    RPC call that returns the block in hex, the workaround is to manually
    recreate the block bytes. The recreated block bytes may not be the exact
    as in the underlying blockchain but it is good enough for our indexing
    purposes.'''

    async def raw_blocks(self, hex_hashes):
        '''Return the raw binary blocks with the given hex hashes.'''
        params_iterable = ((h, ) for h in hex_hashes)
        block_info = await self._send_vector('getblock', params_iterable)

        blocks = []
        for i in block_info:
            raw_block = await self.make_raw_block(i)
            blocks.append(raw_block)

        # Convert hex string to bytes
        return blocks

    async def make_raw_header(self, b):
        pbh = b.get('previousblockhash')
        if pbh is None:
            pbh = '0' * 64
        return b''.join([
            pack('<L', b.get('version')),
            hex_str_to_hash(pbh),
            hex_str_to_hash(b.get('merkleroot')),
            pack('<L', self.timestamp_safe(b['time'])),
            pack('<L', int(b.get('bits'), 16)),
            pack('<L', int(b.get('nonce')))
        ])

    async def make_raw_block(self, b):
        '''Construct a raw block'''

        header = await self.make_raw_header(b)

        transactions = []
        if b.get('height') > 0:
            transactions = await self.getrawtransactions(b.get('tx'), False)

        raw_block = header
        num_txs = len(transactions)
        if num_txs > 0:
            raw_block += int_to_varint(num_txs)
            raw_block += b''.join(transactions)
        else:
            raw_block += b'\x00'

        return raw_block

    def timestamp_safe(self, t):
        if isinstance(t, int):
            return t
        return timegm(strptime(t, "%Y-%m-%d %H:%M:%S %Z"))
<|endoftext|>


<commit_before>"""
The command line interface.  Trains a directory of data.
"""
import gc
import sys
import os
import hyperchamber as hc
import tensorflow as tf
from hypergan.gan_component import ValidationException
from .inputs import *
from .viewer import GlobalViewer
from .configuration import Configuration
import hypergan as hg
import time

import os
import shutil
import sys

from hypergan.losses.supervised_loss import SupervisedLoss
from hypergan.multi_component import MultiComponent
from time import sleep

class CLI:
    def __init__(self, gan, args={}):
        self.samples = 0
        self.steps = 0
        self.gan = gan
        if gan is not None:
            self.gan.cli = self

        args = hc.Config(args)
        self.args = args

        crop =  self.args.crop

        self.config_name = self.args.config or 'default'
        self.method = args.method or 'test'
        self.total_steps = args.steps or -1
        self.sample_every = self.args.sample_every or 100

        self.sampler_name = args.sampler
        self.sampler = None
        self.validate()
        if self.args.save_file:
            self.save_file = self.args.save_file
        else:
            default_save_path = os.path.abspath("saves/"+self.config_name)
            self.save_file = default_save_path + "/model.ckpt"
            self.create_path(self.save_file)
        if self.gan is not None:
            self.gan.save_file = self.save_file

        title = "[hypergan] " + self.config_name
        GlobalViewer.enable_menu = self.args.menu
        GlobalViewer.title = title
        GlobalViewer.viewer_size = self.args.viewer_size
        GlobalViewer.enabled = self.args.viewer
        GlobalViewer.zoom = self.args.zoom

    def sample(self, allow_save=True):
        """ Samples to a file.  Useful for visualizing the learning process.

        If allow_save is False then saves will not be created.

        Use with:

             ffmpeg -i samples/grid-%06d.png -vcodec libx264 -crf 22 -threads 0 grid1-7.mp4

        to create a video of the learning process.
        """
        sample_file="samples/%s/%06d.png" % (self.config_name, self.samples)
        self.create_path(sample_file)
        self.lazy_create()
        sample_list = self.sampler.sample(sample_file, allow_save and self.args.save_samples)
        self.samples += 1

        return sample_list

    def validate(self):
        return True

    def lazy_create(self):
        if(self.sampler == None):
            self.sampler = self.gan.sampler_for(self.sampler_name)(self.gan)
            if(self.sampler == None):
                raise ValidationException("No sampler found by the name '"+self.sampler_name+"'")

    def step(self):
        bgan = self.gan
        self.gan.step()
        if bgan.destroy:
            self.sampler=None
            self.gan = self.gan.newgan
            gc.collect()
            refs = gc.get_referrers(bgan)
            d = bgan.trainer._delegate
            bgan.trainer=None
            gc.collect()
            del bgan
            tf.reset_default_graph()

            gc.collect()

        if(self.steps % self.sample_every == 0):
            sample_list = self.sample()

        self.steps+=1

    def create_path(self, filename):
        return os.makedirs(os.path.expanduser(os.path.dirname(filename)), exist_ok=True)

    def build(self):
        return self.gan.build()
    def serve(self, gan):
        return gan_server(self.gan.session, config)

    def sample_forever(self):
        while not self.gan.destroy:
            self.sample()
            GlobalViewer.tick()

    def train(self):
        i=0
        if(self.args.ipython):
            import fcntl
            fd = sys.stdin.fileno()
            fl = fcntl.fcntl(fd, fcntl.F_GETFL)
            fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)

        while((i < self.total_steps or self.total_steps == -1) and not self.gan.destroy):
            i+=1
            start_time = time.time()
            self.step()
            GlobalViewer.tick()

            if (self.args.save_every != None and
                self.args.save_every != -1 and
                self.args.save_every > 0 and
                i % self.args.save_every == 0):
                print(" |= Saving network")
                self.gan.save(self.save_file)
            if self.args.ipython:
                self.check_stdin()
            end_time = time.time()

    def check_stdin(self):
        try:
            input = sys.stdin.read()
            if input[0]=="y":
                return
            from IPython import embed
            # Misc code
            embed()

        except:
            return

    def new(self):
        template = self.args.directory + '.json'
        print("[hypergan] Creating new configuration file '"+template+"' based off of '"+self.config_name+".json'")
        if os.path.isfile(template):
            raise ValidationException("File exists: " + template)
        source_configuration = Configuration.find(self.config_name+".json")
        shutil.copyfile(source_configuration, template)

        return

    def add_supervised_loss(self):
        if self.args.classloss:
            print("[discriminator] Class loss is on.  Semi-supervised learning mode activated.")
            supervised_loss = SupervisedLoss(self.gan, self.gan.config.loss)
            self.gan.loss = MultiComponent(components=[supervised_loss, self.gan.loss], combine='add')
            #EWW
        else:
            print("[discriminator] Class loss is off.  Unsupervised learning mode activated.")

    def run(self):
        if self.method == 'train':
            self.add_supervised_loss() # TODO I think this is broken now(after moving create out)
            self.gan.session.run(tf.global_variables_initializer())

            if not self.gan.load(self.save_file):
                print("Initializing new model")
            else:
                print("Model loaded")
            self.train()
            self.gan.save(self.save_file)
            tf.reset_default_graph()
            self.gan.session.close()
        elif self.method == 'build':
            if not self.gan.load(self.save_file):
                raise "Could not load model: "+ save_file
            else:
                print("Model loaded")
            self.build()
        elif self.method == 'new':
            self.new()
        elif self.method == 'sample':
            self.add_supervised_loss()
            if not self.gan.load(self.save_file):
                print("Initializing new model")
            else:
                print("Model loaded")

            tf.train.start_queue_runners(sess=self.gan.session)
            self.sample_forever()
            tf.reset_default_graph()
            self.gan.session.close()
        elif self.method == 'test':
            print("Hooray!")
            print("Hypergan is installed correctly.  Testing tensorflow for GPU support.")
            with tf.Session() as sess:
                devices = sess.list_devices()

            if not tf.test.gpu_device_name():
                print("Warning: no default GPU device available")
                allgood=False
            else:
                print("Default GPU is available")
                allgood=True
            print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
            print("Current available tensorflow devices:")
            for device in devices:
                print(device)
            if allgood:
                print("Congratulations!  Tensorflow and hypergan both look installed correctly.  If you still experience issues come let us know on discord.")
            else:
                print("There were errors in the test, please see the logs")

<commit_msg>[feature] exception if build cant load<commit_after>"""
The command line interface.  Trains a directory of data.
"""
import gc
import sys
import os
import hyperchamber as hc
import tensorflow as tf
from hypergan.gan_component import ValidationException
from .inputs import *
from .viewer import GlobalViewer
from .configuration import Configuration
import hypergan as hg
import time

import os
import shutil
import sys

from hypergan.losses.supervised_loss import SupervisedLoss
from hypergan.multi_component import MultiComponent
from time import sleep

class CLI:
    def __init__(self, gan, args={}):
        self.samples = 0
        self.steps = 0
        self.gan = gan
        if gan is not None:
            self.gan.cli = self

        args = hc.Config(args)
        self.args = args

        crop =  self.args.crop

        self.config_name = self.args.config or 'default'
        self.method = args.method or 'test'
        self.total_steps = args.steps or -1
        self.sample_every = self.args.sample_every or 100

        self.sampler_name = args.sampler
        self.sampler = None
        self.validate()
        if self.args.save_file:
            self.save_file = self.args.save_file
        else:
            default_save_path = os.path.abspath("saves/"+self.config_name)
            self.save_file = default_save_path + "/model.ckpt"
            self.create_path(self.save_file)
        if self.gan is not None:
            self.gan.save_file = self.save_file

        title = "[hypergan] " + self.config_name
        GlobalViewer.enable_menu = self.args.menu
        GlobalViewer.title = title
        GlobalViewer.viewer_size = self.args.viewer_size
        GlobalViewer.enabled = self.args.viewer
        GlobalViewer.zoom = self.args.zoom

    def sample(self, allow_save=True):
        """ Samples to a file.  Useful for visualizing the learning process.

        If allow_save is False then saves will not be created.

        Use with:

             ffmpeg -i samples/grid-%06d.png -vcodec libx264 -crf 22 -threads 0 grid1-7.mp4

        to create a video of the learning process.
        """
        sample_file="samples/%s/%06d.png" % (self.config_name, self.samples)
        self.create_path(sample_file)
        self.lazy_create()
        sample_list = self.sampler.sample(sample_file, allow_save and self.args.save_samples)
        self.samples += 1

        return sample_list

    def validate(self):
        return True

    def lazy_create(self):
        if(self.sampler == None):
            self.sampler = self.gan.sampler_for(self.sampler_name)(self.gan)
            if(self.sampler == None):
                raise ValidationException("No sampler found by the name '"+self.sampler_name+"'")

    def step(self):
        bgan = self.gan
        self.gan.step()
        if bgan.destroy:
            self.sampler=None
            self.gan = self.gan.newgan
            gc.collect()
            refs = gc.get_referrers(bgan)
            d = bgan.trainer._delegate
            bgan.trainer=None
            gc.collect()
            del bgan
            tf.reset_default_graph()

            gc.collect()

        if(self.steps % self.sample_every == 0):
            sample_list = self.sample()

        self.steps+=1

    def create_path(self, filename):
        return os.makedirs(os.path.expanduser(os.path.dirname(filename)), exist_ok=True)

    def build(self):
        return self.gan.build()
    def serve(self, gan):
        return gan_server(self.gan.session, config)

    def sample_forever(self):
        while not self.gan.destroy:
            self.sample()
            GlobalViewer.tick()

    def train(self):
        i=0
        if(self.args.ipython):
            import fcntl
            fd = sys.stdin.fileno()
            fl = fcntl.fcntl(fd, fcntl.F_GETFL)
            fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)

        while((i < self.total_steps or self.total_steps == -1) and not self.gan.destroy):
            i+=1
            start_time = time.time()
            self.step()
            GlobalViewer.tick()

            if (self.args.save_every != None and
                self.args.save_every != -1 and
                self.args.save_every > 0 and
                i % self.args.save_every == 0):
                print(" |= Saving network")
                self.gan.save(self.save_file)
            if self.args.ipython:
                self.check_stdin()
            end_time = time.time()

    def check_stdin(self):
        try:
            input = sys.stdin.read()
            if input[0]=="y":
                return
            from IPython import embed
            # Misc code
            embed()

        except:
            return

    def new(self):
        template = self.args.directory + '.json'
        print("[hypergan] Creating new configuration file '"+template+"' based off of '"+self.config_name+".json'")
        if os.path.isfile(template):
            raise ValidationException("File exists: " + template)
        source_configuration = Configuration.find(self.config_name+".json")
        shutil.copyfile(source_configuration, template)

        return

    def add_supervised_loss(self):
        if self.args.classloss:
            print("[discriminator] Class loss is on.  Semi-supervised learning mode activated.")
            supervised_loss = SupervisedLoss(self.gan, self.gan.config.loss)
            self.gan.loss = MultiComponent(components=[supervised_loss, self.gan.loss], combine='add')
            #EWW
        else:
            print("[discriminator] Class loss is off.  Unsupervised learning mode activated.")

    def run(self):
        if self.method == 'train':
            self.add_supervised_loss() # TODO I think this is broken now(after moving create out)
            self.gan.session.run(tf.global_variables_initializer())

            if not self.gan.load(self.save_file):
                print("Initializing new model")
            else:
                print("Model loaded")
            self.train()
            self.gan.save(self.save_file)
            tf.reset_default_graph()
            self.gan.session.close()
        elif self.method == 'build':
            if not self.gan.load(self.save_file):
                raise ValidationException("Could not load model: "+ self.save_file)
            else:
                print("Model loaded")
            self.build()
        elif self.method == 'new':
            self.new()
        elif self.method == 'sample':
            self.add_supervised_loss()
            if not self.gan.load(self.save_file):
                print("Initializing new model")
            else:
                print("Model loaded")

            tf.train.start_queue_runners(sess=self.gan.session)
            self.sample_forever()
            tf.reset_default_graph()
            self.gan.session.close()
        elif self.method == 'test':
            print("Hooray!")
            print("Hypergan is installed correctly.  Testing tensorflow for GPU support.")
            with tf.Session() as sess:
                devices = sess.list_devices()

            if not tf.test.gpu_device_name():
                print("Warning: no default GPU device available")
                allgood=False
            else:
                print("Default GPU is available")
                allgood=True
            print('Default GPU Device: {}'.format(tf.test.gpu_device_name()))
            print("Current available tensorflow devices:")
            for device in devices:
                print(device)
            if allgood:
                print("Congratulations!  Tensorflow and hypergan both look installed correctly.  If you still experience issues come let us know on discord.")
            else:
                print("There were errors in the test, please see the logs")

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# -*- test-case-name: logsnarf.test.test_snarf -*-
# pylint: disable=invalid-name

"""Log directory watcher and reader.

This class is responsible for watching a set of directories for updates
to files matching a specific pattern. When those updates happen, it reads
those changes and passes it to the consumer given to it at initialization
time.

Updates are sent to the consumer as complete lines (including newlines).
Infile progress is tracked via a persistent state object, which tracks
the inode and file offset.
"""

import codecs
import logging
import os
import os.path

from zope import interface
from twisted.internet import inotify
from twisted.internet import interfaces
from twisted.python import filepath

class LogSnarf(object):
    """Main logsnarf class.

    implements :twisted:`twisted.internet.interfaces.IPushProducer`
    """
    interface.implements(interfaces.IPushProducer)

    def __init__(self, state_obj, consumer, reactor=None):
        """Initialize a LogSnarf object.

        :param state_obj: current state
        :type state_obj: logsnarf.state.State
        :param consumer: a consumer object
        :type consumer:
            implements(:twisted:`twisted.internet.interfaces.IConsumer`)
        """
        if not reactor:
            from twisted.internet import reactor
        self.consumer = consumer
        self.reactor = reactor
        self._inotifier = inotify.INotify(reactor=self.reactor)
        self._callback = None
        self._state = state_obj
        self._patterns = {}
        self._paused_in_doRead = []
        self.paused = True
        self.log = logging.getLogger(self.__class__.__name__)
        self.consumer.registerProducer(self, True)
        self.setCallback(self.consumer.write)

    def watch(self, path, pattern=None, recursive=True):
        """Add a watch to a given directory.

        :param path: directory to watch for changes in
        :param pattern: regular expression that files much match
        :param recursive: If true, also watch subdirectories.

        :type path: string or :twisted:`twisted.python.filepath.FilePath`
        :type pattern: re.RegexObject or None
        :type recursive: bool
        """
        if not isinstance(path, filepath.FilePath):
            path = filepath.FilePath(path)
        self._patterns[path.path] = pattern
        # noinspection PyUnresolvedReferences
        self.reactor.callWhenRunning(self._do_backlog, path, pattern,
                                     recursive)
        self._inotifier.watch(path, recursive=recursive, autoAdd=True,
                              callbacks=[self._snarfcb],
                              mask=inotify.IN_MODIFY | inotify.IN_DELETE)

    def _do_backlog(self, path, pattern, recursive):
        """Process the backlog.

        Files we know about are processed from last offset, new files
        are processed and added to the state file.
        """
        self.log.info("Processing backlog in %s pattern: %s recursive: %s",
                      path.path, pattern.pattern, recursive)
        if recursive:
            filenames = []
            for dirpath, _, filename in os.walk(path.path):
                filenames.extend([os.path.join(dirpath, f) for f in filename])
        else:
            filenames = [
                os.path.join(path.path, f) for f in os.listdir(path.path)
            ]
        if pattern:
            filenames = filter(pattern.match, filenames)

        self.log.info("Files to process as backlog %s", filenames)
        filenames = map(filepath.FilePath, filenames)

        map(self.doRead, filenames)

    def start(self):
        """Start watching."""
        self._inotifier.startReading()
        self.paused = False

    def setCallback(self, callback):
        """Add a callback function.

        The callback should have a signature of f(line), accepting one line
        (including newline) of text.
        """
        if not callable(callback):
            self.log.error('Attempt to add a non callable callback %r',
                           callback)
            raise TypeError('%r is not callable.', callback)
        if hasattr(callback, 'im_class'):
            name = '%s.%s' % (callback.im_class, callback.__name__)
        else:
            name = callback.__name__
        self.log.debug('Callback set to %s', name)
        self._callback = callback

    def pauseProducing(self):
        """:twisted:`twisted.internet.interfaces.IPushProducer` method

        Pause producing.
        """
        self.paused = True
        self._inotifier.pauseProducing()
        self.log.debug('Paused producing')

    def resumeProducing(self):
        """:twisted:`twisted.internet.interfaces.IPushProducer` method

        Resume producing.
        """
        self.paused = False
        self.log.debug('Resuming production')
        for path in self._paused_in_doRead[:]:
            # We can be paused during one of these calls.
            if self.paused:
                return
            self.log.debug('Resuming read of %s', path.path)
            self._paused_in_doRead.remove(path)
            self.doRead(path)
        self.log.debug('Resuming inotify')
        self._inotifier.resumeProducing()

    def cleanState(self):
        """Remove invalid entries from the state file."""
        for path in self._state:
            if not os.path.exists(path):
                self._state.pop(path, None)

    def checkPattern(self, path):
        """Check a path against our pattern.

        :param twisted.python.filepath.FilePath: path to check
        :returns bool: true if the path matches.
        """
        # Start with most specific (longest) path
        for pth in sorted(self._patterns, reverse=True,
                          key=len):
            if path.path.startswith(pth.path):
                if self._patterns[pth] is None or \
                        self._patterns[pth].search(path.path):
                    return True
        return False

    def doRead(self, path):
        """read from a file.

        :param path: the file to read
        :type path: :twisted:`twisted.python.filepath.FilePath`
        """
        path.restat()
        offset, inode = self._state.get(path.path, [0, path.getInodeNumber()])
        new_inode = path.getInodeNumber()

        if inode != new_inode:
            self.log.warning('%s is a new file, inodes differ. Starting '
                             'from offset 0', path.path)
            offset = 0
            inode = new_inode
        elif path.getsize() < offset:
            self.log.warning('The file at %s is of size %d, shorter than our '
                             'previous offset, which was %d. Starting from '
                             'offset 0', path.path, path.getsize(), offset)
            offset = 0
        self.log.debug('do_read: %s offset: %d', path.path, offset)
        try:
            with path.open() as raw_fp:
                f = codecs.getreader('utf-8')(raw_fp, errors='ignore')
                if offset != 0:
                    f.seek(offset - 1)
                    c = f.read(1)
                    if c != '\n':
                        offset = self._seekNewLine(f, offset)
                        f.seek(offset)
                line = True
                while line:
                    line = f.readline()
                    if line.endswith('\n'):
                        # the reader looks ahead, we need to offset this when
                        # storing, well, offsets
                        offset = f.tell()
                        if f.charbuffer is not None:
                            offset -= len(f.charbuffer)
                        self._callback(line)
                        if self.paused:
                            self.log.debug('Paused by consumer')
                            self._paused_in_doRead.append(path)
                            self._state[path.path] = [offset, inode]
                            return

                self._state[path.path] = [offset, inode]
        except IOError:
            self.log.exception('error while processing %s', path)

    def _seekNewLine(self, fp, offset):
        """Seek backwards to the first newline in fp from offset."""
        fp.seek(offset - 4096)
        buf = fp.read(4096)
        new_offset = offset - 4096 + buf.rfind('\n') + 1
        self.log.debug('Rewinding %r to %d', fp, new_offset)
        fp.seek(new_offset)
        return new_offset

    def _snarfcb(self, _, path, mask):
        """The callback given to :twisted:`twisted.internet.inotify.INotify`"""
        if mask & inotify.IN_DELETE:
            self._state.pop(path.path, None)
            return
        if not mask & inotify.IN_MODIFY:
            return
        if not self.checkPattern(path):
            return
        self.doRead(path)
<commit_msg>Another checkPattern fix, missed a path.path instance.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# -*- test-case-name: logsnarf.test.test_snarf -*-
# pylint: disable=invalid-name

"""Log directory watcher and reader.

This class is responsible for watching a set of directories for updates
to files matching a specific pattern. When those updates happen, it reads
those changes and passes it to the consumer given to it at initialization
time.

Updates are sent to the consumer as complete lines (including newlines).
Infile progress is tracked via a persistent state object, which tracks
the inode and file offset.
"""

import codecs
import logging
import os
import os.path

from zope import interface
from twisted.internet import inotify
from twisted.internet import interfaces
from twisted.python import filepath

class LogSnarf(object):
    """Main logsnarf class.

    implements :twisted:`twisted.internet.interfaces.IPushProducer`
    """
    interface.implements(interfaces.IPushProducer)

    def __init__(self, state_obj, consumer, reactor=None):
        """Initialize a LogSnarf object.

        :param state_obj: current state
        :type state_obj: logsnarf.state.State
        :param consumer: a consumer object
        :type consumer:
            implements(:twisted:`twisted.internet.interfaces.IConsumer`)
        """
        if not reactor:
            from twisted.internet import reactor
        self.consumer = consumer
        self.reactor = reactor
        self._inotifier = inotify.INotify(reactor=self.reactor)
        self._callback = None
        self._state = state_obj
        self._patterns = {}
        self._paused_in_doRead = []
        self.paused = True
        self.log = logging.getLogger(self.__class__.__name__)
        self.consumer.registerProducer(self, True)
        self.setCallback(self.consumer.write)

    def watch(self, path, pattern=None, recursive=True):
        """Add a watch to a given directory.

        :param path: directory to watch for changes in
        :param pattern: regular expression that files much match
        :param recursive: If true, also watch subdirectories.

        :type path: string or :twisted:`twisted.python.filepath.FilePath`
        :type pattern: re.RegexObject or None
        :type recursive: bool
        """
        if not isinstance(path, filepath.FilePath):
            path = filepath.FilePath(path)
        self._patterns[path.path] = pattern
        # noinspection PyUnresolvedReferences
        self.reactor.callWhenRunning(self._do_backlog, path, pattern,
                                     recursive)
        self._inotifier.watch(path, recursive=recursive, autoAdd=True,
                              callbacks=[self._snarfcb],
                              mask=inotify.IN_MODIFY | inotify.IN_DELETE)

    def _do_backlog(self, path, pattern, recursive):
        """Process the backlog.

        Files we know about are processed from last offset, new files
        are processed and added to the state file.
        """
        self.log.info("Processing backlog in %s pattern: %s recursive: %s",
                      path.path, pattern.pattern, recursive)
        if recursive:
            filenames = []
            for dirpath, _, filename in os.walk(path.path):
                filenames.extend([os.path.join(dirpath, f) for f in filename])
        else:
            filenames = [
                os.path.join(path.path, f) for f in os.listdir(path.path)
            ]
        if pattern:
            filenames = filter(pattern.match, filenames)

        self.log.info("Files to process as backlog %s", filenames)
        filenames = map(filepath.FilePath, filenames)

        map(self.doRead, filenames)

    def start(self):
        """Start watching."""
        self._inotifier.startReading()
        self.paused = False

    def setCallback(self, callback):
        """Add a callback function.

        The callback should have a signature of f(line), accepting one line
        (including newline) of text.
        """
        if not callable(callback):
            self.log.error('Attempt to add a non callable callback %r',
                           callback)
            raise TypeError('%r is not callable.', callback)
        if hasattr(callback, 'im_class'):
            name = '%s.%s' % (callback.im_class, callback.__name__)
        else:
            name = callback.__name__
        self.log.debug('Callback set to %s', name)
        self._callback = callback

    def pauseProducing(self):
        """:twisted:`twisted.internet.interfaces.IPushProducer` method

        Pause producing.
        """
        self.paused = True
        self._inotifier.pauseProducing()
        self.log.debug('Paused producing')

    def resumeProducing(self):
        """:twisted:`twisted.internet.interfaces.IPushProducer` method

        Resume producing.
        """
        self.paused = False
        self.log.debug('Resuming production')
        for path in self._paused_in_doRead[:]:
            # We can be paused during one of these calls.
            if self.paused:
                return
            self.log.debug('Resuming read of %s', path.path)
            self._paused_in_doRead.remove(path)
            self.doRead(path)
        self.log.debug('Resuming inotify')
        self._inotifier.resumeProducing()

    def cleanState(self):
        """Remove invalid entries from the state file."""
        for path in self._state:
            if not os.path.exists(path):
                self._state.pop(path, None)

    def checkPattern(self, path):
        """Check a path against our pattern.

        :param twisted.python.filepath.FilePath: path to check
        :returns bool: true if the path matches.
        """
        # Start with most specific (longest) path
        for pth in sorted(self._patterns, reverse=True,
                          key=len):
            if path.path.startswith(pth):
                if self._patterns[pth] is None or \
                        self._patterns[pth].search(path.path):
                    return True
        return False

    def doRead(self, path):
        """read from a file.

        :param path: the file to read
        :type path: :twisted:`twisted.python.filepath.FilePath`
        """
        path.restat()
        offset, inode = self._state.get(path.path, [0, path.getInodeNumber()])
        new_inode = path.getInodeNumber()

        if inode != new_inode:
            self.log.warning('%s is a new file, inodes differ. Starting '
                             'from offset 0', path.path)
            offset = 0
            inode = new_inode
        elif path.getsize() < offset:
            self.log.warning('The file at %s is of size %d, shorter than our '
                             'previous offset, which was %d. Starting from '
                             'offset 0', path.path, path.getsize(), offset)
            offset = 0
        self.log.debug('do_read: %s offset: %d', path.path, offset)
        try:
            with path.open() as raw_fp:
                f = codecs.getreader('utf-8')(raw_fp, errors='ignore')
                if offset != 0:
                    f.seek(offset - 1)
                    c = f.read(1)
                    if c != '\n':
                        offset = self._seekNewLine(f, offset)
                        f.seek(offset)
                line = True
                while line:
                    line = f.readline()
                    if line.endswith('\n'):
                        # the reader looks ahead, we need to offset this when
                        # storing, well, offsets
                        offset = f.tell()
                        if f.charbuffer is not None:
                            offset -= len(f.charbuffer)
                        self._callback(line)
                        if self.paused:
                            self.log.debug('Paused by consumer')
                            self._paused_in_doRead.append(path)
                            self._state[path.path] = [offset, inode]
                            return

                self._state[path.path] = [offset, inode]
        except IOError:
            self.log.exception('error while processing %s', path)

    def _seekNewLine(self, fp, offset):
        """Seek backwards to the first newline in fp from offset."""
        fp.seek(offset - 4096)
        buf = fp.read(4096)
        new_offset = offset - 4096 + buf.rfind('\n') + 1
        self.log.debug('Rewinding %r to %d', fp, new_offset)
        fp.seek(new_offset)
        return new_offset

    def _snarfcb(self, _, path, mask):
        """The callback given to :twisted:`twisted.internet.inotify.INotify`"""
        if mask & inotify.IN_DELETE:
            self._state.pop(path.path, None)
            return
        if not mask & inotify.IN_MODIFY:
            return
        if not self.checkPattern(path):
            return
        self.doRead(path)
<|endoftext|>