<commit_before>
from __future__ import absolute_import, with_statement

import functools
import logging
import threading
import time

import pyuv

from collections import deque
from twisted.internet.base import _SignalReactorMixin
from twisted.internet.posixbase import PosixReactorBase
from twisted.internet.interfaces import IReactorFDSet, IDelayedCall, IReactorTime
from twisted.python import failure, log
from twisted.internet import error
from zope.interface import implements

class UVWaker(object):
    def __init__(self, reactor):
        self._async = pyuv.Async(reactor._loop, self._cb)
        self._async.unref()
    def _cb(self, handle):
        pass
    def wakeUp(self):
        self._async.send()

class UVDelayedCall(object):
    implements(IDelayedCall)

    def __init__(self, reactor, seconds, f, *args, **kw):
        self._reactor = reactor
        self._func = functools.partial(f, *args, **kw)
        self._time = self._reactor.seconds() + seconds
        self._timer = pyuv.Timer(self._reactor._loop)
        self._timer.start(self._called, self._time-self._reactor.seconds(), 0.0)
        self._active = True

    def _called(self, handle):
        self._active = False
        self._timer.stop()
        self._reactor._removeDelayedCall(self)
        try:
            self._func()
        except Exception:
            logging.error("_called caught exception", exc_info=True)

    def getTime(self):
        return self._time

    def cancel(self):
        self._active = False
        self._timer.stop()
        self._reactor._removeDelayedCall(self)

    def delay(self, seconds):
        self._timer.stop()
        self._time += seconds
        self._timer.start(self._called, self._time-self._reactor.seconds(), 0.0)

    def reset(self, seconds):
        self._timer.stop()
        self._time = self._reactor.seconds() + seconds
        self._timer.start(self._called, self._time-self._reactor.seconds(), 0.0)

    def active(self):
        return self._active

class UVReactor(PosixReactorBase):
    implements(IReactorTime, IReactorFDSet)

    def __init__(self):
        self._loop = pyuv.Loop()
        self._signal_watcher = pyuv.Signal(self._loop)
        self._async_handle = pyuv.Async(self._loop, self._async_cb)
        self._async_handle_lock = threading.Lock()
        self._async_callbacks = deque()
        self._readers = {}  # map of reader objects to fd
        self._writers = {}  # map of writer objects to fd
        self._fds = {}      # map of fd to a (reader, writer) tuple
        self._delayedCalls = {}
        self._poll_handles = {}
        PosixReactorBase.__init__(self)

    def _get_loop_handles(self):
        handles = set()
        def cb(handle):
            if not handle.closed:
                handles.add(handle)
        self._loop.walk(cb)
        return handles

    # IReactorTime
    def seconds(self):
        return time.time()

    def callLater(self, seconds, f, *args, **kw):
        dc = UVDelayedCall(self, seconds, f, *args, **kw)
        self._delayedCalls[dc] = True
        return dc

    def getDelayedCalls(self):
        return [x for x in self._delayedCalls if x._active]

    def _removeDelayedCall(self, dc):
        if dc in self._delayedCalls:
            del self._delayedCalls[dc]

    def _async_cb(self, handle):
        with self._async_handle_lock:
            callbacks = self._async_callbacks
            self._async_callbacks = deque()
        while callbacks:
            cb = callbacks.popleft()
            try:
                cb()
            except Exception:
                log.err()

    # IReactorThreads

    def callFromThread(self, f, *args, **kw):
        """See `twisted.internet.interfaces.IReactorThreads.callFromThread`"""
        assert callable(f), "%s is not callable" % f
        cb = functools.partial(f, *args, **kw)
        with self._async_handle_lock:
            self._async_callbacks.append(cb)
        self._async_handle.send()

    def _handleSignals(self):
        """Bypass installing the child waker"""
        _SignalReactorMixin._handleSignals(self)

    # IReactorProcess

    def spawnProcess(self, processProtocol, executable, args=(), env={}, path=None, uid=None, gid=None, usePTY=0, childFDs=None):
        raise NotImplementedError("spawnProcess")

    def installWaker(self):
        self.waker = UVWaker(self)

    def wakeUp(self):
        if self.waker:
            self.waker.wakeUp()

    def _invoke_callback(self, handle, events, poll_error):
        fd = handle.fd
        reader, writer = self._fds[fd]
        if reader:
            err = None
            if reader.fileno() == -1:
                err = error.ConnectionLost()
            elif events & pyuv.UV_READABLE:
                err = log.callWithLogger(reader, reader.doRead)
            if err is None and poll_error is not None:
                err = error.ConnectionLost()
            if err is not None:
                self.removeReader(reader)
                reader.readConnectionLost(failure.Failure(err))
        if writer:
            err = None
            if writer.fileno() == -1:
                err = error.ConnectionLost()
            elif events & pyuv.UV_WRITABLE:
                err = log.callWithLogger(writer, writer.doWrite)
            if err is None and poll_error is not None:
                err = error.ConnectionLost()
            if err is not None:
                self.removeWriter(writer)
                writer.writeConnectionLost(failure.Failure(err))

    # IReactorFDSet

    def addReader(self, reader):
        """Add a FileDescriptor for notification of data available to read."""
        if reader in self._readers:
            # Don't add the reader if it's already there
            return
        fd = reader.fileno()
        self._readers[reader] = fd
        if fd in self._fds:
            _, writer = self._fds[fd]
            self._fds[fd] = (reader, writer)
            if writer:
                # We already registered this fd for write events,
                # update it for read events as well.
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_READABLE | pyuv.UV_WRITABLE, self._invoke_callback)
        else:
            self._fds[fd] = (reader, None)
            poll_handle = pyuv.Poll(self._loop, fd)
            poll_handle.start(pyuv.UV_READABLE, self._invoke_callback)
            poll_handle.fd = fd
            self._poll_handles[fd] = poll_handle

    def addWriter(self, writer):
        """Add a FileDescriptor for notification of data available to write."""
        if writer in self._writers:
            return
        fd = writer.fileno()
        self._writers[writer] = fd
        if fd in self._fds:
            reader, _ = self._fds[fd]
            self._fds[fd] = (reader, writer)
            if reader:
                # We already registered this fd for read events,
                # update it for write events as well.
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_READABLE | pyuv.UV_WRITABLE, self._invoke_callback)
        else:
            self._fds[fd] = (None, writer)
            poll_handle = pyuv.Poll(self._loop, fd)
            poll_handle.start(pyuv.UV_WRITABLE, self._invoke_callback)
            poll_handle.fd = fd
            self._poll_handles[fd] = poll_handle

    def removeReader(self, reader):
        """Remove a Selectable for notification of data available to read."""
        if reader in self._readers:
            fd = self._readers.pop(reader)
            _, writer = self._fds[fd]
            if writer:
                # We have a writer so we need to update the IOLoop for
                # write events only.
                self._fds[fd] = (None, writer)
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_WRITABLE, self._invoke_callback)
            else:
                # Since we have no writer registered, we remove the
                # entry from _fds and unregister the handler from the
                # IOLoop
                del self._fds[fd]
                del self._poll_handles[fd]

    def removeWriter(self, writer):
        """Remove a Selectable for notification of data available to write."""
        if writer in self._writers:
            fd = self._writers.pop(writer)
            reader, _ = self._fds[fd]
            if reader:
                # We have a reader so we need to update the IOLoop for
                # read events only.
                self._fds[fd] = (reader, None)
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_READABLE, self._invoke_callback)
            else:
                # Since we have no reader registered, we remove the
                # entry from the _fds and unregister the handler from
                # the IOLoop.
                del self._fds[fd]
                del self._poll_handles[fd]

    def removeAll(self):
        return self._removeAll(self._readers, self._writers)

    def getReaders(self):
        return self._readers.keys()

    def getWriters(self):
        return self._writers.keys()

    def stop(self):
        self._async_handle.close()
        for handle in self._get_loop_handles():
            if hasattr(handle, 'stop') and callable(handle.stop):
                handle.stop()
        PosixReactorBase.stop(self)

    def crash(self):
        # No need to stop handles, it's done by stop
        PosixReactorBase.crash(self)

    def runUntilCurrent(self):
        raise NotImplementedError("runUntilCurrent")

    def startRunning(self, installSignalHandlers=True):
        self._signal_watcher.start()
        _SignalReactorMixin.startRunning(self, installSignalHandlers)

    def doIteration(self, delay):
        raise NotImplementedError("doIteration")

    def mainLoop(self):
        while self._started:
            try:
                while self._started:
                    if self._stopped:
                        self.fireSystemEvent("shutdown")
                        break
                    self._loop.run()
            except:
                log.msg("Unexpected error in main loop.")
                log.err()
            else:
                for handle in self._get_loop_handles():
                    handle.close()
                # Run the loop so the close callbacks are fired and memory is freed
                # It will not block because all handles are closed
                self._loop.run()
                log.msg('Main loop terminated.')

def install():
    reactor = UVReactor()
    from twisted.internet.main import installReactor
    installReactor(reactor)

<commit_msg>Removed unneded function<commit_after>
from __future__ import absolute_import, with_statement

import functools
import logging
import threading
import time

import pyuv

from collections import deque
from twisted.internet.base import _SignalReactorMixin
from twisted.internet.posixbase import PosixReactorBase
from twisted.internet.interfaces import IReactorFDSet, IDelayedCall, IReactorTime
from twisted.python import failure, log
from twisted.internet import error
from zope.interface import implements

class UVWaker(object):
    def __init__(self, reactor):
        self._async = pyuv.Async(reactor._loop, lambda x: None)
        self._async.unref()
    def wakeUp(self):
        self._async.send()

class UVDelayedCall(object):
    implements(IDelayedCall)

    def __init__(self, reactor, seconds, f, *args, **kw):
        self._reactor = reactor
        self._func = functools.partial(f, *args, **kw)
        self._time = self._reactor.seconds() + seconds
        self._timer = pyuv.Timer(self._reactor._loop)
        self._timer.start(self._called, self._time-self._reactor.seconds(), 0.0)
        self._active = True

    def _called(self, handle):
        self._active = False
        self._timer.stop()
        self._reactor._removeDelayedCall(self)
        try:
            self._func()
        except Exception:
            logging.error("_called caught exception", exc_info=True)

    def getTime(self):
        return self._time

    def cancel(self):
        self._active = False
        self._timer.stop()
        self._reactor._removeDelayedCall(self)

    def delay(self, seconds):
        self._timer.stop()
        self._time += seconds
        self._timer.start(self._called, self._time-self._reactor.seconds(), 0.0)

    def reset(self, seconds):
        self._timer.stop()
        self._time = self._reactor.seconds() + seconds
        self._timer.start(self._called, self._time-self._reactor.seconds(), 0.0)

    def active(self):
        return self._active

class UVReactor(PosixReactorBase):
    implements(IReactorTime, IReactorFDSet)

    def __init__(self):
        self._loop = pyuv.Loop()
        self._signal_watcher = pyuv.Signal(self._loop)
        self._async_handle = pyuv.Async(self._loop, self._async_cb)
        self._async_handle_lock = threading.Lock()
        self._async_callbacks = deque()
        self._readers = {}  # map of reader objects to fd
        self._writers = {}  # map of writer objects to fd
        self._fds = {}      # map of fd to a (reader, writer) tuple
        self._delayedCalls = {}
        self._poll_handles = {}
        PosixReactorBase.__init__(self)

    def _get_loop_handles(self):
        handles = set()
        def cb(handle):
            if not handle.closed:
                handles.add(handle)
        self._loop.walk(cb)
        return handles

    # IReactorTime
    def seconds(self):
        return time.time()

    def callLater(self, seconds, f, *args, **kw):
        dc = UVDelayedCall(self, seconds, f, *args, **kw)
        self._delayedCalls[dc] = True
        return dc

    def getDelayedCalls(self):
        return [x for x in self._delayedCalls if x._active]

    def _removeDelayedCall(self, dc):
        if dc in self._delayedCalls:
            del self._delayedCalls[dc]

    def _async_cb(self, handle):
        with self._async_handle_lock:
            callbacks = self._async_callbacks
            self._async_callbacks = deque()
        while callbacks:
            cb = callbacks.popleft()
            try:
                cb()
            except Exception:
                log.err()

    # IReactorThreads

    def callFromThread(self, f, *args, **kw):
        """See `twisted.internet.interfaces.IReactorThreads.callFromThread`"""
        assert callable(f), "%s is not callable" % f
        cb = functools.partial(f, *args, **kw)
        with self._async_handle_lock:
            self._async_callbacks.append(cb)
        self._async_handle.send()

    def _handleSignals(self):
        """Bypass installing the child waker"""
        _SignalReactorMixin._handleSignals(self)

    # IReactorProcess

    def spawnProcess(self, processProtocol, executable, args=(), env={}, path=None, uid=None, gid=None, usePTY=0, childFDs=None):
        raise NotImplementedError("spawnProcess")

    def installWaker(self):
        self.waker = UVWaker(self)

    def wakeUp(self):
        if self.waker:
            self.waker.wakeUp()

    def _invoke_callback(self, handle, events, poll_error):
        fd = handle.fd
        reader, writer = self._fds[fd]
        if reader:
            err = None
            if reader.fileno() == -1:
                err = error.ConnectionLost()
            elif events & pyuv.UV_READABLE:
                err = log.callWithLogger(reader, reader.doRead)
            if err is None and poll_error is not None:
                err = error.ConnectionLost()
            if err is not None:
                self.removeReader(reader)
                reader.readConnectionLost(failure.Failure(err))
        if writer:
            err = None
            if writer.fileno() == -1:
                err = error.ConnectionLost()
            elif events & pyuv.UV_WRITABLE:
                err = log.callWithLogger(writer, writer.doWrite)
            if err is None and poll_error is not None:
                err = error.ConnectionLost()
            if err is not None:
                self.removeWriter(writer)
                writer.writeConnectionLost(failure.Failure(err))

    # IReactorFDSet

    def addReader(self, reader):
        """Add a FileDescriptor for notification of data available to read."""
        if reader in self._readers:
            # Don't add the reader if it's already there
            return
        fd = reader.fileno()
        self._readers[reader] = fd
        if fd in self._fds:
            _, writer = self._fds[fd]
            self._fds[fd] = (reader, writer)
            if writer:
                # We already registered this fd for write events,
                # update it for read events as well.
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_READABLE | pyuv.UV_WRITABLE, self._invoke_callback)
        else:
            self._fds[fd] = (reader, None)
            poll_handle = pyuv.Poll(self._loop, fd)
            poll_handle.start(pyuv.UV_READABLE, self._invoke_callback)
            poll_handle.fd = fd
            self._poll_handles[fd] = poll_handle

    def addWriter(self, writer):
        """Add a FileDescriptor for notification of data available to write."""
        if writer in self._writers:
            return
        fd = writer.fileno()
        self._writers[writer] = fd
        if fd in self._fds:
            reader, _ = self._fds[fd]
            self._fds[fd] = (reader, writer)
            if reader:
                # We already registered this fd for read events,
                # update it for write events as well.
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_READABLE | pyuv.UV_WRITABLE, self._invoke_callback)
        else:
            self._fds[fd] = (None, writer)
            poll_handle = pyuv.Poll(self._loop, fd)
            poll_handle.start(pyuv.UV_WRITABLE, self._invoke_callback)
            poll_handle.fd = fd
            self._poll_handles[fd] = poll_handle

    def removeReader(self, reader):
        """Remove a Selectable for notification of data available to read."""
        if reader in self._readers:
            fd = self._readers.pop(reader)
            _, writer = self._fds[fd]
            if writer:
                # We have a writer so we need to update the IOLoop for
                # write events only.
                self._fds[fd] = (None, writer)
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_WRITABLE, self._invoke_callback)
            else:
                # Since we have no writer registered, we remove the
                # entry from _fds and unregister the handler from the
                # IOLoop
                del self._fds[fd]
                del self._poll_handles[fd]

    def removeWriter(self, writer):
        """Remove a Selectable for notification of data available to write."""
        if writer in self._writers:
            fd = self._writers.pop(writer)
            reader, _ = self._fds[fd]
            if reader:
                # We have a reader so we need to update the IOLoop for
                # read events only.
                self._fds[fd] = (reader, None)
                poll_handle = self._poll_handles[fd]
                poll_handle.start(pyuv.UV_READABLE, self._invoke_callback)
            else:
                # Since we have no reader registered, we remove the
                # entry from the _fds and unregister the handler from
                # the IOLoop.
                del self._fds[fd]
                del self._poll_handles[fd]

    def removeAll(self):
        return self._removeAll(self._readers, self._writers)

    def getReaders(self):
        return self._readers.keys()

    def getWriters(self):
        return self._writers.keys()

    def stop(self):
        self._async_handle.close()
        for handle in self._get_loop_handles():
            if hasattr(handle, 'stop') and callable(handle.stop):
                handle.stop()
        PosixReactorBase.stop(self)

    def crash(self):
        # No need to stop handles, it's done by stop
        PosixReactorBase.crash(self)

    def runUntilCurrent(self):
        raise NotImplementedError("runUntilCurrent")

    def startRunning(self, installSignalHandlers=True):
        self._signal_watcher.start()
        _SignalReactorMixin.startRunning(self, installSignalHandlers)

    def doIteration(self, delay):
        raise NotImplementedError("doIteration")

    def mainLoop(self):
        while self._started:
            try:
                while self._started:
                    if self._stopped:
                        self.fireSystemEvent("shutdown")
                        break
                    self._loop.run()
            except:
                log.msg("Unexpected error in main loop.")
                log.err()
            else:
                for handle in self._get_loop_handles():
                    handle.close()
                # Run the loop so the close callbacks are fired and memory is freed
                # It will not block because all handles are closed
                self._loop.run()
                log.msg('Main loop terminated.')

def install():
    reactor = UVReactor()
    from twisted.internet.main import installReactor
    installReactor(reactor)

<|endoftext|>


<commit_before>"""Automated installation on RPM systems with the yum package manager.
"""
from fabric.api import *
from fabric.contrib.files import *

from cloudbio.package.shared import _yaml_to_packages

def _yum_packages(to_install):
    """Install rpm packages available via yum.
    """
    pkg_config = os.path.join(env.config_dir, "packages-yum.yaml")
    with settings(warn_only=True):
        sudo("yum check-update")
    sudo("yum -y upgrade")
    # Retrieve packages to get and install each of them
    (packages, _) = _yaml_to_packages(pkg_config, to_install)
    # At this point allow the Flavor to rewrite the package list
    packages = env.flavor.rewrite_config_items("packages", packages)
    for package in packages:
        sudo("yum -y install %s" % package)

def _setup_yum_bashrc():
    return
    """Fix the user bashrc to update compilers.
    """
    # XXX:  this is not very flexible, and malicious :-)  It causes our builds to fail in Scientific Linux
    to_include = ["export CC=gcc44", "export CXX=g++44", "export FC=gfortran44",
                  "export PKG_CONFIG_PATH=${PKG_CONFIG_PATH}:/usr/lib/pkgconfig"]
    fname = run("ls %s" % env.shell_config)
    for line in to_include:
        if not contains(fname, line.split("=")[0]):
            append(fname, line)

def _setup_yum_sources():
    """Add additional useful yum repositories.
    """
    repos = ["http://download.fedora.redhat.com/pub/epel/5/x86_64/epel-release-5-4.noarch.rpm"]
    for repo in repos:
        with settings(warn_only=True):
            sudo("rpm -Uvh %s" % repo)
<commit_msg>Add cloudera repository<commit_after>"""Automated installation on RPM systems with the yum package manager.
"""
from fabric.api import *
from fabric.contrib.files import *

from cloudbio.package.shared import _yaml_to_packages

def _yum_packages(to_install):
    """Install rpm packages available via yum.
    """
    pkg_config = os.path.join(env.config_dir, "packages-yum.yaml")
    with settings(warn_only=True):
        sudo("yum check-update")
    sudo("yum -y upgrade")
    # Retrieve packages to get and install each of them
    (packages, _) = _yaml_to_packages(pkg_config, to_install)
    # At this point allow the Flavor to rewrite the package list
    packages = env.flavor.rewrite_config_items("packages", packages)
    for package in packages:
        sudo("yum -y install %s" % package)

def _setup_yum_bashrc():
    return
    """Fix the user bashrc to update compilers.
    """
    # XXX:  this is not very flexible, and malicious :-)  It causes our builds to fail in Scientific Linux
    to_include = ["export CC=gcc44", "export CXX=g++44", "export FC=gfortran44",
                  "export PKG_CONFIG_PATH=${PKG_CONFIG_PATH}:/usr/lib/pkgconfig"]
    fname = run("ls %s" % env.shell_config)
    for line in to_include:
        if not contains(fname, line.split("=")[0]):
            append(fname, line)

def _setup_yum_sources():
    """Add additional useful yum repositories.
    """
    repos = [
      "http://download.fedora.redhat.com/pub/epel/5/x86_64/epel-release-5-4.noarch.rpm",
      "http://archive.cloudera.com/redhat/6/x86_64/cdh/cdh3-repository-1.0-1.noarch.rpm"
    ]
    for repo in repos:
        with settings(warn_only=True):
            sudo("rpm -Uvh %s" % repo)
<|endoftext|>


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

"""PeriodicActions execute small actions periodically in the training loop."""

import abc
import collections
import concurrent.futures
import contextlib
import queue
import time
from typing import Optional

from clu import metric_writers
from clu import platform
from clu import profiler

import jax

MetricWriter = metric_writers.MetricWriter

@jax.jit
def _squareit(x):
  """Minimalistic function for use in _wait_jax_async_dispatch()."""
  return x**2

def _wait_jax_async_dispatch():
  """Creates a simple JAX program and waits for its completion.

  Since JAX operations are put in a queue and dispatched one after the other,
  all previously enqueued computations will be finished after a call to this
  function.
  """
  _squareit(0.).block_until_ready()

class PeriodicAction(abc.ABC):
  """Abstract base class for perodic actions."""

  def __init__(self,
               *,
               every_steps: Optional[int] = None,
               every_secs: Optional[float] = None):
    """Creates an action that triggers periodically.

    Args:
      every_steps: If the current step is divisible by `every_steps`, then an
        action is triggered.
      every_secs: If no action has triggered for specified `every_secs`, then
        an action is triggered. Note that the previous action might have been
        triggered by `every_steps` or by `every_secs`.
    """
    self._every_steps = every_steps
    self._every_secs = every_secs
    self._previous_step = None
    self._previous_time = None
    self._last_step = None

  def _apply_condition(self, step: int, t: float):
    if self._every_steps is not None and step % self._every_steps == 0:
      return True
    if (self._every_secs is not None and
        t - self._previous_time > self._every_secs):
      return True
    return False

  def __call__(self, step: int, t: Optional[float] = None) -> bool:
    """Method to call the hook after every training step.

    Args:
      step: Current step.
      t: Optional timestamp. Will use `time.time()` if not specified.

    Returns:
      True if the action triggered, False otherwise. Note that the first
      invocation never triggers.
    """
    if t is None:
      t = time.time()
    if self._previous_step is None:
      self._previous_step = step
      self._previous_time = t
      self._last_step = step
      return False

    if self._every_steps is not None:
      if step - self._last_step != 1:
        raise ValueError("EveryNHook must be called after every step (once).")
    self._last_step = step

    if self._apply_condition(step, t):
      self._apply(step, t)
      self._previous_step = step
      self._previous_time = t
      return True
    return False

  @abc.abstractmethod
  def _apply(self, step: int, t: float):
    pass

class ReportProgress(PeriodicAction):
  """This hook will set the progress note on the work unit."""

  def __init__(self,
               *,
               num_train_steps: int,
               writer: Optional[MetricWriter] = None,
               every_steps: Optional[int] = None,
               every_secs: Optional[float] = 60.0):
    """Creates a new ReportProgress hook.

    Warning: The progress and the reported steps_per_sec are estimates. We
    ignore the asynchronous dispatch for JAX and other operations in the
    training loop (e.g. evaluation).

    Args:
      num_train_steps: The total number of training steps for training.
      writer: Optional MetricWriter to report steps_per_sec measurement. This is
        an estimate for precise values use Xprof.
      every_steps: How often to report the progress in number of training steps.
      every_secs: How often to report progress as time interval.
    """
    super().__init__(every_steps=every_steps, every_secs=every_secs)
    self._num_train_steps = num_train_steps
    self._writer = writer
    self._waiting_for_part = collections.defaultdict(queue.Queue)
    self._time_per_part = collections.defaultdict(float)
    self._t0 = time.time()
    # Using max_worker=1 guarantees that the calls to _wait_jax_async_dispatch()
    # happen sequentially.
    self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)

  def _apply_condition(self, step: int, t: float):
    # Always trigger at last step.
    if step == self._num_train_steps:
      return True
    return super()._apply_condition(step, t)

  def _apply(self, step: int, t: float):
    steps_per_sec = (step - self._previous_step) / (t - self._previous_time)
    eta_seconds = (self._num_train_steps - step) / steps_per_sec
    message = (f"{100 * step / self._num_train_steps:.1f}% @{step}, "
               f"{steps_per_sec:.1f} steps/s, ETA: {eta_seconds / 60:.0f} min")
    if self._time_per_part:
      total = time.time() - self._t0
      message += " ({:.0f} min : {})".format(total / 60, ", ".join(
          f"{100 * dt / total:.1f}% {name}"
          for name, dt in sorted(self._time_per_part.items())))
    # This should be relatively cheap so we can do it in the same main thread.
    platform.work_unit().set_notes(message)
    if self._writer is not None:
      self._writer.write_scalars(step, {"steps_per_sec": steps_per_sec})

  @contextlib.contextmanager
  def timed(self, name: str, wait_jax_async_dispatch: bool = True):
    # pylint: disable=g-doc-return-or-yield
    """Measures time spent in a named part of the training loop.

    The reported progress will break down the total time into the different
    parts spent inside blocks.

    Example:

      report_progress = hooks.ReportProgress()
      for step, batch in enumerate(train_iter):
        params = train_step(params, batch)
        report_progress(step + 1)
        if (step + 1) % eval_every_steps == 0:
          with report_progress.timed("eval"):
            evaluate()

    The above example would result in the progress being reported as something
    like "10% @1000 ... (5 min : 10% eval)" - assuming that evaluation takes 10%
    of the entire time in this case.

    Args:
      name: Name of the part to be measured.
      wait_jax_async_dispatch: When set to `True`, JAX async dispatch queue will
        be emptied by creating a new computation and waiting for its completion.
        This makes sure that previous computations (e.g. the last train step)
        have actually finished. The same is done before the time is measured.
        Note that this wait happens in a different thread that is only used for
        measuring start/stop time of timed parts. In other words, the measured
        timings reflect the start/stop of the JAX computations within the
        measured part: the timer is started when the last computation before the
        block has finished, and the timer is stopped when the last computation
        from within the block has finished. Note that due to JAX execution these
        operations asynchronously, the measured time might overlap with non-JAX
        computations outside the measured block.
        When set to `False`, then the measured time is of the Python statements
        within the block.
        If there are no expensive JAX computations enqueued in JAX's async
        dispatch queue, then both measurements are identical.
    """
    # pylint: enable=g-doc-return-or-yield
    def start_measurement():
      if wait_jax_async_dispatch:
        _wait_jax_async_dispatch()
      self._waiting_for_part[name].put(time.time())
    self._executor.submit(start_measurement)

    yield

    def stop_measurement():
      if wait_jax_async_dispatch:
        _wait_jax_async_dispatch()
      dt = time.time() - self._waiting_for_part[name].get()
      self._time_per_part[name] += dt
    self._executor.submit(stop_measurement)

class Profile(PeriodicAction):
  """This hook collects a profile every time it triggers."""

  def __init__(self,
               *,
               num_profile_steps: int = 5,
               first_profile: int = 10,
               every_steps: Optional[int] = None,
               every_secs: Optional[float] = 3600.0,
               logdir: Optional[str] = None):
    """Initializes a new periodic profiler action.

    Args:
      num_profile_steps: Over how many steps the profile should be taken.
      first_profile: First step at which a profile is started.
      every_steps: See `PeriodicAction.__init__()`.
      every_secs: See `PeriodicAction.__init__()`.
      logdir: Where the profile should be stored (required for
        `tf.profiler.experimental`).
    """
    super().__init__(every_steps=every_steps, every_secs=every_secs)
    self._num_profile_steps = num_profile_steps
    self._first_profile = first_profile
    self._session_running = False
    self._logdir = logdir

  def _apply_condition(self, step: int, t: float) -> bool:
    if self._session_running:
      if step >= self._previous_step + self._num_profile_steps:
        self._end_session()
      return False
    if step == self._first_profile:
      return True
    return super()._apply_condition(step, t)

  def _apply(self, step: int, t: float):
    del step, t  # Unused.
    self._start_session()

  def _start_session(self):
    self._session_running = True
    profiler.start(logdir=self._logdir)

  def _end_session(self):
    url = profiler.stop()
    if url is not None:
      platform.work_unit().create_artifact(
          platform.ArtifactType.URL,
          url,
          description=f"[{self._previous_step}] Profile")
    self._session_running = False
<commit_msg>Fix type errors exposed by upcoming JAX change.<commit_after># Copyright 2020 The CLU Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""PeriodicActions execute small actions periodically in the training loop."""

import abc
import collections
import concurrent.futures
import contextlib
import queue
import time
from typing import Optional

from clu import metric_writers
from clu import platform
from clu import profiler

import jax
import jax.numpy as jnp

MetricWriter = metric_writers.MetricWriter

@jax.jit
def _squareit(x):
  """Minimalistic function for use in _wait_jax_async_dispatch()."""
  return x**2

def _wait_jax_async_dispatch():
  """Creates a simple JAX program and waits for its completion.

  Since JAX operations are put in a queue and dispatched one after the other,
  all previously enqueued computations will be finished after a call to this
  function.
  """
  _squareit(jnp.array(0.)).block_until_ready()

class PeriodicAction(abc.ABC):
  """Abstract base class for perodic actions."""

  def __init__(self,
               *,
               every_steps: Optional[int] = None,
               every_secs: Optional[float] = None):
    """Creates an action that triggers periodically.

    Args:
      every_steps: If the current step is divisible by `every_steps`, then an
        action is triggered.
      every_secs: If no action has triggered for specified `every_secs`, then
        an action is triggered. Note that the previous action might have been
        triggered by `every_steps` or by `every_secs`.
    """
    self._every_steps = every_steps
    self._every_secs = every_secs
    self._previous_step = None
    self._previous_time = None
    self._last_step = None

  def _apply_condition(self, step: int, t: float):
    if self._every_steps is not None and step % self._every_steps == 0:
      return True
    if (self._every_secs is not None and
        t - self._previous_time > self._every_secs):
      return True
    return False

  def __call__(self, step: int, t: Optional[float] = None) -> bool:
    """Method to call the hook after every training step.

    Args:
      step: Current step.
      t: Optional timestamp. Will use `time.time()` if not specified.

    Returns:
      True if the action triggered, False otherwise. Note that the first
      invocation never triggers.
    """
    if t is None:
      t = time.time()
    if self._previous_step is None:
      self._previous_step = step
      self._previous_time = t
      self._last_step = step
      return False

    if self._every_steps is not None:
      if step - self._last_step != 1:
        raise ValueError("EveryNHook must be called after every step (once).")
    self._last_step = step

    if self._apply_condition(step, t):
      self._apply(step, t)
      self._previous_step = step
      self._previous_time = t
      return True
    return False

  @abc.abstractmethod
  def _apply(self, step: int, t: float):
    pass

class ReportProgress(PeriodicAction):
  """This hook will set the progress note on the work unit."""

  def __init__(self,
               *,
               num_train_steps: int,
               writer: Optional[MetricWriter] = None,
               every_steps: Optional[int] = None,
               every_secs: Optional[float] = 60.0):
    """Creates a new ReportProgress hook.

    Warning: The progress and the reported steps_per_sec are estimates. We
    ignore the asynchronous dispatch for JAX and other operations in the
    training loop (e.g. evaluation).

    Args:
      num_train_steps: The total number of training steps for training.
      writer: Optional MetricWriter to report steps_per_sec measurement. This is
        an estimate for precise values use Xprof.
      every_steps: How often to report the progress in number of training steps.
      every_secs: How often to report progress as time interval.
    """
    super().__init__(every_steps=every_steps, every_secs=every_secs)
    self._num_train_steps = num_train_steps
    self._writer = writer
    self._waiting_for_part = collections.defaultdict(queue.Queue)
    self._time_per_part = collections.defaultdict(float)
    self._t0 = time.time()
    # Using max_worker=1 guarantees that the calls to _wait_jax_async_dispatch()
    # happen sequentially.
    self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1)

  def _apply_condition(self, step: int, t: float):
    # Always trigger at last step.
    if step == self._num_train_steps:
      return True
    return super()._apply_condition(step, t)

  def _apply(self, step: int, t: float):
    steps_per_sec = (step - self._previous_step) / (t - self._previous_time)
    eta_seconds = (self._num_train_steps - step) / steps_per_sec
    message = (f"{100 * step / self._num_train_steps:.1f}% @{step}, "
               f"{steps_per_sec:.1f} steps/s, ETA: {eta_seconds / 60:.0f} min")
    if self._time_per_part:
      total = time.time() - self._t0
      message += " ({:.0f} min : {})".format(total / 60, ", ".join(
          f"{100 * dt / total:.1f}% {name}"
          for name, dt in sorted(self._time_per_part.items())))
    # This should be relatively cheap so we can do it in the same main thread.
    platform.work_unit().set_notes(message)
    if self._writer is not None:
      self._writer.write_scalars(step, {"steps_per_sec": steps_per_sec})

  @contextlib.contextmanager
  def timed(self, name: str, wait_jax_async_dispatch: bool = True):
    # pylint: disable=g-doc-return-or-yield
    """Measures time spent in a named part of the training loop.

    The reported progress will break down the total time into the different
    parts spent inside blocks.

    Example:

      report_progress = hooks.ReportProgress()
      for step, batch in enumerate(train_iter):
        params = train_step(params, batch)
        report_progress(step + 1)
        if (step + 1) % eval_every_steps == 0:
          with report_progress.timed("eval"):
            evaluate()

    The above example would result in the progress being reported as something
    like "10% @1000 ... (5 min : 10% eval)" - assuming that evaluation takes 10%
    of the entire time in this case.

    Args:
      name: Name of the part to be measured.
      wait_jax_async_dispatch: When set to `True`, JAX async dispatch queue will
        be emptied by creating a new computation and waiting for its completion.
        This makes sure that previous computations (e.g. the last train step)
        have actually finished. The same is done before the time is measured.
        Note that this wait happens in a different thread that is only used for
        measuring start/stop time of timed parts. In other words, the measured
        timings reflect the start/stop of the JAX computations within the
        measured part: the timer is started when the last computation before the
        block has finished, and the timer is stopped when the last computation
        from within the block has finished. Note that due to JAX execution these
        operations asynchronously, the measured time might overlap with non-JAX
        computations outside the measured block.
        When set to `False`, then the measured time is of the Python statements
        within the block.
        If there are no expensive JAX computations enqueued in JAX's async
        dispatch queue, then both measurements are identical.
    """
    # pylint: enable=g-doc-return-or-yield
    def start_measurement():
      if wait_jax_async_dispatch:
        _wait_jax_async_dispatch()
      self._waiting_for_part[name].put(time.time())
    self._executor.submit(start_measurement)

    yield

    def stop_measurement():
      if wait_jax_async_dispatch:
        _wait_jax_async_dispatch()
      dt = time.time() - self._waiting_for_part[name].get()
      self._time_per_part[name] += dt
    self._executor.submit(stop_measurement)

class Profile(PeriodicAction):
  """This hook collects a profile every time it triggers."""

  def __init__(self,
               *,
               num_profile_steps: int = 5,
               first_profile: int = 10,
               every_steps: Optional[int] = None,
               every_secs: Optional[float] = 3600.0,
               logdir: Optional[str] = None):
    """Initializes a new periodic profiler action.

    Args:
      num_profile_steps: Over how many steps the profile should be taken.
      first_profile: First step at which a profile is started.
      every_steps: See `PeriodicAction.__init__()`.
      every_secs: See `PeriodicAction.__init__()`.
      logdir: Where the profile should be stored (required for
        `tf.profiler.experimental`).
    """
    super().__init__(every_steps=every_steps, every_secs=every_secs)
    self._num_profile_steps = num_profile_steps
    self._first_profile = first_profile
    self._session_running = False
    self._logdir = logdir

  def _apply_condition(self, step: int, t: float) -> bool:
    if self._session_running:
      if step >= self._previous_step + self._num_profile_steps:
        self._end_session()
      return False
    if step == self._first_profile:
      return True
    return super()._apply_condition(step, t)

  def _apply(self, step: int, t: float):
    del step, t  # Unused.
    self._start_session()

  def _start_session(self):
    self._session_running = True
    profiler.start(logdir=self._logdir)

  def _end_session(self):
    url = profiler.stop()
    if url is not None:
      platform.work_unit().create_artifact(
          platform.ArtifactType.URL,
          url,
          description=f"[{self._previous_step}] Profile")
    self._session_running = False
<|endoftext|>


<commit_before>import io
import os
import shutil
import subprocess
import tempfile

# An example of how to use LibreOffice and ImageMagick to convert WMF images to
# PNGs.
#
# libreoffice_wmf_conversion uses LibreOffice to convert the image to a PNG.
# This normally creates an image with a large amount of padding, so
# imagemagick_trim can be used to trim the image.
#
# The image can be then be converted using a normal image handler, such as
# mammoth.images.data_uri.
#
# Example usage:
#
# def convert_image(image):
#     image = libreoffice_wmf_conversion(image, post_process=imagemagick_trim)
#     return mammoth.images.data_uri(image)
#
#
# fileobj = open("document.docx", "rb")
# result = mammoth.convert_to_html(fileobj, convert_image=convert_image)

_wmf_extensions = {
    "image/x-wmf": ".wmf",
    "image/x-emf": ".emf",
}

def libreoffice_wmf_conversion(image, post_process=None):
    if post_process is None:
        post_process = lambda x: x

    wmf_extension = _wmf_extensions.get(image.content_type)
    if wmf_extension is None:
        return image
    else:
        temporary_directory = tempfile.mkdtemp()
        try:
            input_path = os.path.join(temporary_directory, "image" + wmf_extension)
            with io.open(input_path, "wb") as input_fileobj:
                with image.open() as image_fileobj:
                    shutil.copyfileobj(image_fileobj, input_fileobj)

            output_path = os.path.join(temporary_directory, "image.png")
            subprocess.check_call([
                "libreoffice",
                "--headless",
                "--convert-to",
                "png",
                input_path,
                "--outdir",
                temporary_directory,
            ])

            with io.open(output_path, "rb") as output_fileobj:
                output = output_fileobj.read()

            def open_image():
                return io.BytesIO(output)

            return post_process(image.copy(
                content_type="image/png",
                open=open_image,
            ))
        finally:
            shutil.rmtree(temporary_directory)

def imagemagick_trim(image):
    command = ["convert", "-", "-trim", "-"]
    process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    try:
        with image.open() as image_fileobj:
            shutil.copyfileobj(image_fileobj, process.stdin)
        output, err_output = process.communicate()
    except:
        process.kill()
        process.wait()
        raise

    return_code = process.poll()
    if return_code:
        raise subprocess.CalledProcessError(return_code, command)
    else:
        def open_image():
            return io.BytesIO(output)

        return image.copy(open=open_image)

<commit_msg>Use "with" to open file in example<commit_after>import io
import os
import shutil
import subprocess
import tempfile

# An example of how to use LibreOffice and ImageMagick to convert WMF images to
# PNGs.
#
# libreoffice_wmf_conversion uses LibreOffice to convert the image to a PNG.
# This normally creates an image with a large amount of padding, so
# imagemagick_trim can be used to trim the image.
#
# The image can be then be converted using a normal image handler, such as
# mammoth.images.data_uri.
#
# Example usage:
#
# def convert_image(image):
#     image = libreoffice_wmf_conversion(image, post_process=imagemagick_trim)
#     return mammoth.images.data_uri(image)
#
# with open("document.docx", "rb") as fileobj:
#     result = mammoth.convert_to_html(fileobj, convert_image=convert_image)

_wmf_extensions = {
    "image/x-wmf": ".wmf",
    "image/x-emf": ".emf",
}

def libreoffice_wmf_conversion(image, post_process=None):
    if post_process is None:
        post_process = lambda x: x

    wmf_extension = _wmf_extensions.get(image.content_type)
    if wmf_extension is None:
        return image
    else:
        temporary_directory = tempfile.mkdtemp()
        try:
            input_path = os.path.join(temporary_directory, "image" + wmf_extension)
            with io.open(input_path, "wb") as input_fileobj:
                with image.open() as image_fileobj:
                    shutil.copyfileobj(image_fileobj, input_fileobj)

            output_path = os.path.join(temporary_directory, "image.png")
            subprocess.check_call([
                "libreoffice",
                "--headless",
                "--convert-to",
                "png",
                input_path,
                "--outdir",
                temporary_directory,
            ])

            with io.open(output_path, "rb") as output_fileobj:
                output = output_fileobj.read()

            def open_image():
                return io.BytesIO(output)

            return post_process(image.copy(
                content_type="image/png",
                open=open_image,
            ))
        finally:
            shutil.rmtree(temporary_directory)

def imagemagick_trim(image):
    command = ["convert", "-", "-trim", "-"]
    process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE)
    try:
        with image.open() as image_fileobj:
            shutil.copyfileobj(image_fileobj, process.stdin)
        output, err_output = process.communicate()
    except:
        process.kill()
        process.wait()
        raise

    return_code = process.poll()
    if return_code:
        raise subprocess.CalledProcessError(return_code, command)
    else:
        def open_image():
            return io.BytesIO(output)

        return image.copy(open=open_image)

<|endoftext|>


<commit_before>import functools
from .result import DeferredResult

class Pipeline(object):
    __slots__ = ['_client', '_stack', '_callbacks']

    def __init__(self, redis_client):
        self._client = redis_client
        self._stack = []
        self._callbacks = []

    def __getattr__(self, item):
        f = getattr(self._client, item)
        if not callable(f):
            return f

        @functools.wraps(f)
        def inner(*args, **kwargs):
            f(*args, **kwargs)
            ref = DeferredResult()
            self._stack.append(ref)
            return ref

        return inner

    def execute(self, raise_on_error=True):
        stack = self._stack
        callbacks = self._callbacks
        self._stack = []
        self._callbacks = []
        res = self._client.execute(raise_on_error=raise_on_error)
        for i, v in enumerate(res):
            stack[i].set(v)
        for cb in callbacks:
            cb()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._client.__exit__(exc_type, exc_val, exc_tb)
        self._stack = []
        self._callbacks = []

    def reset(self):
        self._client.reset()
        self._stack = []
        self._callbacks = []

    def on_execute(self, callback):
        self._callbacks.append(callback)

class NestedPipeline(object):
    __slots__ = ['_pipe', '_stack', '_callbacks']

    def __init__(self, pipe):
        self._pipe = pipe
        self._stack = []
        self._callbacks = []

    def __getattr__(self, item):
        f = getattr(self._pipe, item)
        if not callable(f):
            return f

        @functools.wraps(f)
        def inner(*args, **kwargs):
            ref = DeferredResult()
            self._stack.append((f, args, kwargs, ref))
            return ref

        return inner

    def execute(self, raise_on_error=True):
        stack = self._stack
        callbacks = self._callbacks
        self._stack = []
        self._callbacks = []

        def build(res, ref):
            def cb():
                ref.result = res.result

            self._pipe.on_execute(cb)

        for f, args, kwargs, ref in stack:
            build(f(*args, **kwargs), ref)

        for cb in callbacks:
            self._pipe.on_execute(cb)

    def on_execute(self, callback):
        self._callbacks.append(callback)

    def reset(self):
        self._stack = []
        self._callbacks = []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.reset()
<commit_msg>improvements to nested pipeline and reset logic<commit_after>import functools
from .result import DeferredResult

class Pipeline(object):
    """
    Wrapper for redispy pipeline object.
    It returns a reference that contains a result
    once the pipeline executes.
    This allows us to be able to pipeline
    lots of calls within nested functions
    and not have to wait for the execute call.
    """
    __slots__ = ['_pipe', '_stack', '_callbacks']

    def __init__(self, pipe):
        """
        pass in the redispy client pipeline object.
        :param pipe: redis.StrictRedis.pipeline() or redis.Redis.pipeline()
        """
        self._pipe = pipe
        self._stack = []
        self._callbacks = []

    def __getattr__(self, item):
        """
        magic method to intercept all calls bound for the internal
        client object and return a deferred result reference object.
        :param item: function name
        :return:
        """
        # get the original attribute from the client.
        # this could be a class method, or an instance attribute.
        f = getattr(self._pipe, item)

        # if it is just an attribute, return it.
        if not callable(f):
            return f

        # build a decorator for the internal method.
        # when the function is called, we create a DeferredResult object
        # and return it, and invoke the internal method against the
        # redispy pipeline object.
        # keep track of the ref object we return so we can put data into
        # the result on execute.
        @functools.wraps(f)
        def inner(*args, **kwargs):
            f(*args, **kwargs)
            ref = DeferredResult()
            self._stack.append(ref)
            return ref

        return inner

    def execute(self, raise_on_error=True):
        """
        Invoke the redispy pipeline.execute() method and take all the values
        returned in sequential order of commands and map them to the
        DeferredResult objects we returned when each command was queued inside
        the pipeline.
        Also invoke all the callback functions queued up.
        :param raise_on_error: boolean
        :return: None
        """
        stack = self._stack
        callbacks = self._callbacks
        self._stack = []
        self._callbacks = []
        res = self._pipe.execute(raise_on_error=raise_on_error)
        for i, v in enumerate(res):
            stack[i].set(v)
        for cb in callbacks:
            cb()

    def __enter__(self):
        """
        magic method to allow us to use in context like this:

            with Pipeline(redis.StrictRedis().pipeline()) as pipe:
                ref = pipe.set('foo', 'bar')
                pipe.execute()

        we are overriding the behavior in redispy.
        :return: Pipeline instance
        """
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """
        context manager cleanup method.
        :param exc_type:
        :param exc_val:
        :param exc_tb:
        :return:
        """
        self._pipe.__exit__(exc_type, exc_val, exc_tb)
        self._stack = []
        self._callbacks = []

    def reset(self):
        """
        cleanup method. get rid of the stack and callbacks.
        :return:
        """
        self._pipe.reset()
        self._stack = []
        self._callbacks = []

    def on_execute(self, callback):
        """

        :param callback:
        :return:
        """
        self._callbacks.append(callback)

class NestedPipeline(object):
    __slots__ = ['_pipe', '_stack', '_callbacks']

    def __init__(self, pipe):
        self._pipe = pipe
        self._stack = []
        self._callbacks = []

    def __getattr__(self, item):
        f = getattr(self._pipe, item)
        if not callable(f):
            return f

        @functools.wraps(f)
        def inner(*args, **kwargs):
            ref = DeferredResult()
            self._stack.append((f, args, kwargs, ref))
            return ref

        return inner

    def execute(self, raise_on_error=True):
        stack = self._stack
        callbacks = self._callbacks
        self._stack = []
        self._callbacks = []

        def build(res, ref):
            def cb():
                ref.result = res.result

            self._pipe.on_execute(cb)

        for f, args, kwargs, ref in stack:
            build(f(*args, **kwargs), ref)

        for cb in callbacks:
            self._pipe.on_execute(cb)

    def on_execute(self, callback):
        self._callbacks.append(callback)

    def reset(self):
        self._stack = []
        self._callbacks = []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.reset()
<|endoftext|>


<commit_before>from __future__ import print_function
from os.path import isfile
from time import sleep
from yaml import load
from reflector.util import *

KEY_TITLE = {'xml': 'title', 'json': 'Id'}
KEY_CONTENT = 'content'
KEY_SRC = 'src'
KEY_REL = 'rel'
KEY_HREF = 'href'
VALUE_NEXT = {'xml': 'next', 'json': '__next'}

class Config(object):
    def __init__(self):
        with open('config/config.yaml') as data_file:
            c = load(data_file)
            self.remote_url = c['remote']['url'].rstrip('/')
            self.update_feed = c['remote']['update_feed']
            self.remote_json_api = c['remote']['json_api']
            self.local_url = c['local']['url'].rstrip('/')
            self.local_json_api = c['local']['json_api']
            self.local_api_key = c['local']['api_key']
            self.package_storage_path = c['local']['package_storage_path'].rstrip('/').rstrip('\\')
            self.hash_verify_downloads = c['hash']['verify_downloads']
            self.hash_verify_uploaded = c['hash']['verify_uploaded']
            self.dotnet_path = c['local']['dotnet_path']
            if not self.dotnet_path or not isfile(self.dotnet_path):
                raise EnvironmentError('DotNot CLI executable is not configured or the path specified does not exist!')

class Mirror(object):
    def __init__(self,
                 remote_url,
                 update_feed,
                 remote_json_api,
                 local_url,
                 local_json_api,
                 package_storage_path,
                 local_api_key,
                 dotnet_path,
                 verify_downloads=True,
                 verify_uploaded=True
                 ):
        """
        :param remote_url:
        :param update_feed:
        :param remote_json_api:
        :param local_url:
        :param local_json_api:
        :param package_storage_path:
        :param local_api_key:
        :param dotnet_path:
        :param verify_downloads:
        :param verify_uploaded:
        """
        self.remote_api_url = '/'.join([remote_url, 'api/v2'])
        self.update_feed = update_feed
        self.remote_json_api = remote_json_api
        self.remote_packages_url = '/'.join([self.remote_api_url, 'Packages'])
        self.local_api_url = '/'.join([local_url, 'api/v2'])
        self.local_json_api = local_json_api
        self.local_api_upload_url = '/'.join([self.local_api_url, 'package'])
        self.local_packages_url = '/'.join([self.local_api_url, 'Packages'])
        self.package_storage_path = package_storage_path
        self.local_api_key = local_api_key
        self.dotnet_path = dotnet_path
        self.verify_downloads = verify_downloads
        self.verify_uploaded = verify_uploaded

    def __sync(self, content_url, save_to, package_name, version,
               source_hash=None,
               source_hash_method=None,
               dl_reties=0,
               up_retries=0,
               force_dl=False,
               force_up=False
               ):
        """
        :param content_url:
        :param save_to:
        :param package_name:
        :param version:
        :param source_hash:
        :param source_hash_method:
        :param dl_reties:
        :param up_retries:
        :param force_dl:
        :param force_up:
        :return:
        """
        # Is the package already uploaded? Pull it from the target API
        pull_request = pull_package(package_name, version, self.local_packages_url, self.local_json_api)

        if not pull_request:
            print('Unable to sync package.')
            return False

        # What did the target api return
        if pull_request.status_code == 404 or pull_request.status_code == 200 or force_dl or not isfile(save_to):
            # Download the file if we are forcing or it was not already uploaded or cached
            if pull_request.status_code == 404 or not isfile(save_to) or force_dl:
                if not download_file(content_url, save_to):
                    print('Package does not exists after download!?!')
                    return False

            # Did we get a source hash when this was called
            if source_hash is not None and self.verify_downloads:
                # Verify the cached package hash
                hash_verified = verify_hash(save_to, source_hash, hash_method=source_hash_method)
                # If the hash is not verified and we have retired less than 3 times
                if not hash_verified and dl_reties < 3:
                    # Count a retry
                    dl_reties += 1
                    print('Cache hash does not match source hash... retying download...')
                    # Run another sync
                    return self.__sync(
                        content_url,
                        save_to,
                        package_name,
                        version,
                        source_hash=source_hash,
                        source_hash_method=source_hash_method,
                        dl_reties=dl_reties,
                        up_retries=up_retries,
                        force_dl=True,
                        force_up=force_up
                    )
                elif not hash_verified and dl_reties >= 3:
                    # Reached max retries
                    print('Retried to download the package 3 times. Skipping :( ')
                    return False
            else:
                print('Skipping cache hash verification...')
        else:
            # API Error
            print(''.join(['API error! Code: ', str(pull_request.status_code)]))
            return False

        # Made it here? Cache hash either verified or skipped verification
        if pull_request.status_code == 404 or pull_request.status_code == 200 or force_up:
            if pull_request.status_code == 404 or force_up:
                # Send the package up using the dotnet binary
                return_code = push_package_dotnet(
                    save_to,
                    self.local_api_upload_url,
                    self.local_api_key,
                    self.dotnet_path
                )

                if return_code is not 0:
                    print('Push failed, retying with native library.')
                    print('Uploading package...')
                    # If the dotnet binary does not return 0 try to use a python library
                    push_response = push_package_native(
                        save_to,
                        self.local_api_upload_url,
                        self.local_api_key
                    )
                    print(''.join(['Response code: ', str(push_response.status_code)]))
                    if push_response.status_code is not 200:
                        print('Upload failed... :-(')
                        print(''.join(['Response message: ', str(push_response.content)]))

        else:
            # API Error
            # This should never happen. It should get caught above
            print(''.join(['API error! Code: ', str(pull_request.status_code)]))
            return False

        # If we have a source hash Start verifying it
        if source_hash is not None and self.verify_uploaded:
            use_target_json = self.local_json_api
            # Pull the package after uploading it
            pull_request = pull_package(package_name, version, self.local_packages_url, use_target_json)
            # Did we find the package
            if pull_request.status_code == 200:
                # Get the hash
                if use_target_json:
                    # If we are using a json api get it this way
                    target_hash = pull_request.json()['d']['PackageHash']
                else:
                    # If we use the XML api get it this way
                    target_hash = pull_request.objectified.properties.PackageHash.text

                # Does the source hash match the target repo hash?
                if hashes_match(target_hash, source_hash):
                    print('Package synced and verified!')
                    return True
                else:
                    print('Package synced but checksum do not match!')
                    return False

            elif up_retries <= 3:
                up_retries += 1
                print(''.join(['API error! Code: ', str(pull_request.status_code)]))
                print('Package not synced retrying...')
                return self.__sync(
                    content_url,
                    save_to,
                    package_name,
                    version,
                    source_hash=source_hash,
                    source_hash_method=source_hash_method,
                    dl_reties=dl_reties,
                    up_retries=up_retries,
                    force_dl=False,
                    force_up=True
                )

            elif up_retries >= 3:
                print('Max upload retries reached. Skipping :-( ')
                print(''.join(['API error! Code: ', str(pull_request.status_code)]))
                return False

            else:
                # API Error
                print(''.join(['API error! Code: ', str(pull_request.status_code)]))
                return False
        else:
            print('Package synced!')
            return True

    def sync_package(self, package):
        """
        :param package:
        :return:
        """
        # Extract the info that we need from the package entry
        # Dict keys vary depending if the page was pulled in XML or JSON
        use_remote_json = self.remote_json_api
        package_name = str(package[KEY_TITLE['json']]) if use_remote_json else str(package.title.text)
        version = str(package['Version']) if use_remote_json else str(package.properties.Version.text)
        metadata = package['__metadata'] if use_remote_json else {}
        content_url = metadata['media_src'] if use_remote_json else package.content['src']
        package_n_v = '.'.join([package_name, version])
        save_to = os.path.join(self.package_storage_path, '.'.join([package_n_v, 'nupkg']))
        if use_remote_json:
            remote_hash = str(package['PackageHash'])
            remote_hash_method = str(package['PackageHashAlgorithm']).lower()
        else:
            remote_hash = str(package.properties.PackageHash.text)
            remote_hash_method = str(package.properties.PackageHashAlgorithm.text).lower()

        # Begin package sync
        print('')
        print(''.join(['########## ', package_n_v, ' ##########']))

        sync = self.__sync(content_url, save_to, package_name, version, source_hash=remote_hash,
                           source_hash_method=remote_hash_method)

        print('Done!')
        return sync

    def delta_sync(self):
        url = self.update_feed
        previous_delta = read_delta()
        new_delta = None

        if previous_delta is not None:
            print(' '.join(['Syncing packages since:', previous_delta]))
            previous_delta = utc_to_epoch(previous_delta)
        else:
            print('No previous delta syncs. Syncing all updates!')
            previous_delta = first_epoch()

        # Grab the update feed
        response = pull_updates(url)
        # Did the request go well?
        if response.status_code == 200:
            # Get the page
            page = response.objectified
            # Get all items
            items = page.find_all('item')
            # Loop over the items
            for item in items:
                # Get when this was updated
                updated = utc_to_epoch(item.updated.text)
                # Get the new delta on the first package
                if new_delta is None:
                    new_delta = item.updated.text
                # determine if it has been updated since the last run
                if updated > previous_delta:
                    # Grab the package info
                    parts = str(item.origLink.text).split('/')
                    version = parts[-1]
                    title = parts[-2]
                    # Sync the package
                    pull_response = pull_package(title, version, self.remote_packages_url, self.remote_json_api)
                    if pull_response.status_code == 200:
                        package = pull_response.json() if self.remote_json_api else pull_response.objectified
                        self.sync_package(package)
                    else:
                        print('Received bad http code from remote API when pulling package. Response Code: ' + str(
                            pull_response.status_code))

            # If the new delta is still not set. Set it to the previous one
            if new_delta is None:
                new_delta = epoch_to_utc(previous_delta)
            # write epoch to the delta file
            store_delta(new_delta)

        else:
            print('Received bad http code from remote API. Response Code: ' + str(response.status_code))
            return False

        return True

    def sync_packages(self):
        """
        :return:
        """
        done = False
        cool_down_counter = 250
        url = self.remote_packages_url
        use_remote_json = self.remote_json_api
        while not done:
            print('Pulling packages from: ' + url)
            # pull packages from the remote api
            response = pull_packages(url, json=use_remote_json)
            if response:
                # was the response good?
                if response.status_code == 200:
                    if use_remote_json:
                        # Handle JSON pages
                        page = response.json()
                        # data object
                        data = page['d']
                        # Grab the results
                        results = data['results'] if 'results' in data else []
                        # For each result
                        for package in results:
                            # sync it!
                            self.sync_package(package)
                        # If we have a next key continue to the next page
                        if VALUE_NEXT['json'] in data:
                            # Set the url
                            url = data[VALUE_NEXT['json']]
                        else:
                            # Break out
                            done = True
                    else:
                        # Handle XML pages
                        page = response.objectified
                        entries = page.find_all('entry')
                        # Whats the size of the entry list
                        if len(entries) > 0:
                            for package in entries:
                                # sync it!
                                self.sync_package(package)

                        links = page.find_all('link')
                        # Get the last link on the page
                        link = links[0] if 0 > (len(links) - 1) else links[(len(links) - 1)]
                        # If the last link is the next link set it's url as the target url for the next iteration
                        if link[KEY_REL] == VALUE_NEXT['xml']:
                            url = str(link['href'])
                            print(' ')
                            cool_down_counter -= 1
                            if cool_down_counter == 1:
                                # Cool down for 30 seconds every 250 pages...
                                print('Cooling down for 30 seconds...')
                                print(' ')
                                sleep(30)
                                cool_down_counter = 250
                        else:
                            print(' ')
                            print('Done!')
                            # Break out
                            done = True
                else:
                    print(
                        'Received bad http code from remote API. Sleeping for 10 and trying again. Response Code: ' + str(
                            response.status_code))
                    sleep(10)
            else:
                print('Timed out when pulling package list... sleeping for 25 then trying again.')
                sleep(25)
        return True
<commit_msg>added cooldown time to avoid spamming local and remote API<commit_after>from __future__ import print_function
from os.path import isfile
from time import sleep
from yaml import load
from reflector.util import *

KEY_TITLE = {'xml': 'title', 'json': 'Id'}
KEY_CONTENT = 'content'
KEY_SRC = 'src'
KEY_REL = 'rel'
KEY_HREF = 'href'
VALUE_NEXT = {'xml': 'next', 'json': '__next'}

class Config(object):
    def __init__(self):
        with open('config/config.yaml') as data_file:
            c = load(data_file)
            self.remote_url = c['remote']['url'].rstrip('/')
            self.update_feed = c['remote']['update_feed']
            self.remote_json_api = c['remote']['json_api']
            self.local_url = c['local']['url'].rstrip('/')
            self.local_json_api = c['local']['json_api']
            self.local_api_key = c['local']['api_key']
            self.package_storage_path = c['local']['package_storage_path'].rstrip('/').rstrip('\\')
            self.hash_verify_downloads = c['hash']['verify_downloads']
            self.hash_verify_uploaded = c['hash']['verify_uploaded']
            self.dotnet_path = c['local']['dotnet_path']
            if not self.dotnet_path or not isfile(self.dotnet_path):
                raise EnvironmentError('DotNot CLI executable is not configured or the path specified does not exist!')

class Mirror(object):
    def __init__(self,
                 remote_url,
                 update_feed,
                 remote_json_api,
                 local_url,
                 local_json_api,
                 package_storage_path,
                 local_api_key,
                 dotnet_path,
                 verify_downloads=True,
                 verify_uploaded=True
                 ):
        """
        :param remote_url:
        :param update_feed:
        :param remote_json_api:
        :param local_url:
        :param local_json_api:
        :param package_storage_path:
        :param local_api_key:
        :param dotnet_path:
        :param verify_downloads:
        :param verify_uploaded:
        """
        self.remote_api_url = '/'.join([remote_url, 'api/v2'])
        self.update_feed = update_feed
        self.remote_json_api = remote_json_api
        self.remote_packages_url = '/'.join([self.remote_api_url, 'Packages'])
        self.local_api_url = '/'.join([local_url, 'api/v2'])
        self.local_json_api = local_json_api
        self.local_api_upload_url = '/'.join([self.local_api_url, 'package'])
        self.local_packages_url = '/'.join([self.local_api_url, 'Packages'])
        self.package_storage_path = package_storage_path
        self.local_api_key = local_api_key
        self.dotnet_path = dotnet_path
        self.verify_downloads = verify_downloads
        self.verify_uploaded = verify_uploaded

    def __sync(self, content_url, save_to, package_name, version,
               source_hash=None,
               source_hash_method=None,
               dl_reties=0,
               up_retries=0,
               force_dl=False,
               force_up=False
               ):
        """
        :param content_url:
        :param save_to:
        :param package_name:
        :param version:
        :param source_hash:
        :param source_hash_method:
        :param dl_reties:
        :param up_retries:
        :param force_dl:
        :param force_up:
        :return:
        """
        # Is the package already uploaded? Pull it from the target API
        pull_request = pull_package(package_name, version, self.local_packages_url, self.local_json_api)

        if not pull_request:
            print('Unable to sync package.')
            return False

        # What did the target api return
        if pull_request.status_code == 404 or pull_request.status_code == 200 or force_dl or not isfile(save_to):
            # Download the file if we are forcing or it was not already uploaded or cached
            if pull_request.status_code == 404 or not isfile(save_to) or force_dl:
                if not download_file(content_url, save_to):
                    print('Package does not exists after download!?!')
                    return False

            # Did we get a source hash when this was called
            if source_hash is not None and self.verify_downloads:
                # Verify the cached package hash
                hash_verified = verify_hash(save_to, source_hash, hash_method=source_hash_method)
                # If the hash is not verified and we have retired less than 3 times
                if not hash_verified and dl_reties < 3:
                    # Count a retry
                    dl_reties += 1
                    print('Cache hash does not match source hash... retying download...')
                    # Run another sync
                    return self.__sync(
                        content_url,
                        save_to,
                        package_name,
                        version,
                        source_hash=source_hash,
                        source_hash_method=source_hash_method,
                        dl_reties=dl_reties,
                        up_retries=up_retries,
                        force_dl=True,
                        force_up=force_up
                    )
                elif not hash_verified and dl_reties >= 3:
                    # Reached max retries
                    print('Retried to download the package 3 times. Skipping :( ')
                    return False
            else:
                print('Skipping cache hash verification...')
        else:
            # API Error
            print(''.join(['API error! Code: ', str(pull_request.status_code)]))
            return False

        # Made it here? Cache hash either verified or skipped verification
        if pull_request.status_code == 404 or pull_request.status_code == 200 or force_up:
            if pull_request.status_code == 404 or force_up:
                # Send the package up using the dotnet binary
                return_code = push_package_dotnet(
                    save_to,
                    self.local_api_upload_url,
                    self.local_api_key,
                    self.dotnet_path
                )

                if return_code is not 0:
                    print('Push failed, retying with native library.')
                    print('Uploading package...')
                    # If the dotnet binary does not return 0 try to use a python library
                    push_response = push_package_native(
                        save_to,
                        self.local_api_upload_url,
                        self.local_api_key
                    )
                    print(''.join(['Response code: ', str(push_response.status_code)]))
                    if push_response.status_code is not 200:
                        print('Upload failed... :-(')
                        print(''.join(['Response message: ', str(push_response.content)]))

        else:
            # API Error
            # This should never happen. It should get caught above
            print(''.join(['API error! Code: ', str(pull_request.status_code)]))
            return False

        # If we have a source hash Start verifying it
        if source_hash is not None and self.verify_uploaded:
            use_target_json = self.local_json_api
            # Pull the package after uploading it
            pull_request = pull_package(package_name, version, self.local_packages_url, use_target_json)
            # Did we find the package
            if pull_request.status_code == 200:
                # Get the hash
                if use_target_json:
                    # If we are using a json api get it this way
                    target_hash = pull_request.json()['d']['PackageHash']
                else:
                    # If we use the XML api get it this way
                    target_hash = pull_request.objectified.properties.PackageHash.text

                # Does the source hash match the target repo hash?
                if hashes_match(target_hash, source_hash):
                    print('Package synced and verified!')
                    return True
                else:
                    print('Package synced but checksum do not match!')
                    return False

            elif up_retries <= 3:
                up_retries += 1
                print(''.join(['API error! Code: ', str(pull_request.status_code)]))
                print('Package not synced retrying...')
                return self.__sync(
                    content_url,
                    save_to,
                    package_name,
                    version,
                    source_hash=source_hash,
                    source_hash_method=source_hash_method,
                    dl_reties=dl_reties,
                    up_retries=up_retries,
                    force_dl=False,
                    force_up=True
                )

            elif up_retries >= 3:
                print('Max upload retries reached. Skipping :-( ')
                print(''.join(['API error! Code: ', str(pull_request.status_code)]))
                return False

            else:
                # API Error
                print(''.join(['API error! Code: ', str(pull_request.status_code)]))
                return False
        else:
            print('Package synced!')
            return True

    def sync_package(self, package):
        """
        :param package:
        :return:
        """
        # Extract the info that we need from the package entry
        # Dict keys vary depending if the page was pulled in XML or JSON
        use_remote_json = self.remote_json_api
        package_name = str(package[KEY_TITLE['json']]) if use_remote_json else str(package.title.text)
        version = str(package['Version']) if use_remote_json else str(package.properties.Version.text)
        metadata = package['__metadata'] if use_remote_json else {}
        content_url = metadata['media_src'] if use_remote_json else package.content['src']
        package_n_v = '.'.join([package_name, version])
        save_to = os.path.join(self.package_storage_path, '.'.join([package_n_v, 'nupkg']))
        if use_remote_json:
            remote_hash = str(package['PackageHash'])
            remote_hash_method = str(package['PackageHashAlgorithm']).lower()
        else:
            remote_hash = str(package.properties.PackageHash.text)
            remote_hash_method = str(package.properties.PackageHashAlgorithm.text).lower()

        # Begin package sync
        print('')
        print(''.join(['########## ', package_n_v, ' ##########']))

        sync = self.__sync(content_url, save_to, package_name, version, source_hash=remote_hash,
                           source_hash_method=remote_hash_method)

        print('Done!')
        return sync

    def delta_sync(self):
        url = self.update_feed
        previous_delta = read_delta()
        new_delta = None

        if previous_delta is not None:
            print(' '.join(['Syncing packages since:', previous_delta]))
            previous_delta = utc_to_epoch(previous_delta)
        else:
            print('No previous delta syncs. Syncing all updates!')
            previous_delta = first_epoch()

        # Grab the update feed
        response = pull_updates(url)
        # Did the request go well?
        if response.status_code == 200:
            # Get the page
            page = response.objectified
            # Get all items
            items = page.find_all('item')
            # Loop over the items
            for item in items:
                # Get when this was updated
                updated = utc_to_epoch(item.updated.text)
                # Get the new delta on the first package
                if new_delta is None:
                    new_delta = item.updated.text
                # determine if it has been updated since the last run
                if updated > previous_delta:
                    # Grab the package info
                    parts = str(item.origLink.text).split('/')
                    version = parts[-1]
                    title = parts[-2]
                    # Sync the package
                    pull_response = pull_package(title, version, self.remote_packages_url, self.remote_json_api)
                    if pull_response.status_code == 200:
                        package = pull_response.json() if self.remote_json_api else pull_response.objectified
                        self.sync_package(package)
                    else:
                        print('Received bad http code from remote API when pulling package. Response Code: ' + str(
                            pull_response.status_code))

            # If the new delta is still not set. Set it to the previous one
            if new_delta is None:
                new_delta = epoch_to_utc(previous_delta)
            # write epoch to the delta file
            store_delta(new_delta)

        else:
            print('Received bad http code from remote API. Response Code: ' + str(response.status_code))
            return False

        return True

    def sync_packages(self):
        """
        :return:
        """
        done = False
        cool_down_counter = 250
        url = self.remote_packages_url
        use_remote_json = self.remote_json_api
        while not done:
            print('Pulling packages from: ' + url)
            # pull packages from the remote api
            response = pull_packages(url, json=use_remote_json)
            if response:
                # was the response good?
                if response.status_code == 200:
                    if use_remote_json:
                        # Handle JSON pages
                        page = response.json()
                        # data object
                        data = page['d']
                        # Grab the results
                        results = data['results'] if 'results' in data else []
                        # For each result
                        for package in results:
                            # sync it!
                            self.sync_package(package)
                        # If we have a next key continue to the next page
                        if VALUE_NEXT['json'] in data:
                            # Set the url
                            url = data[VALUE_NEXT['json']]
                        else:
                            # Break out
                            done = True
                    else:
                        # Handle XML pages
                        page = response.objectified
                        entries = page.find_all('entry')
                        # Whats the size of the entry list
                        if len(entries) > 0:
                            for package in entries:
                                # sync it!
                                self.sync_package(package)

                        links = page.find_all('link')
                        # Get the last link on the page
                        link = links[0] if 0 > (len(links) - 1) else links[(len(links) - 1)]
                        # If the last link is the next link set it's url as the target url for the next iteration
                        if link[KEY_REL] == VALUE_NEXT['xml']:
                            url = str(link['href'])
                            print(' ')
                            cool_down_counter -= 1
                            if cool_down_counter == 1:
                                # Cool down for 5 seconds every 250 pages...
                                print('Cooling down for 5 seconds...')
                                print(' ')
                                sleep(30)
                                cool_down_counter = 250
                        else:
                            print(' ')
                            print('Done!')
                            # Break out
                            done = True
                else:
                    print(
                        'Received bad http code from remote API. Sleeping for 10 and trying again. Response Code: ' + str(
                            response.status_code))
                    sleep(10)
            else:
                print('Timed out when pulling package list... sleeping for 25 then trying again.')
                sleep(25)
        return True
<|endoftext|>