<commit_before>import logging
import random
import gevent
import blinker
import time

from itertools import cycle
from collections import defaultdict

from .lookupd import Lookupd
from .nsqd import Nsqd
from .backofftimer import BackoffTimer
from .states import INIT, RUNNING, BACKOFF, THROTTLED, CLOSED

from .errors import (
    NSQException,
    NSQNoConnections,
    NSQRequeueMessage
)

class Reader(object):
    def __init__(
        self,
        topic,
        channel,
        nsqd_tcp_addresses=[],
        lookupd_http_addresses=[],
        name=None,
        async=False,
        max_tries=5,
        max_in_flight=1,
        requeue_delay=0,
        lookupd_poll_interval=60,
        lookupd_poll_jitter=0.3,
        low_ready_idle_timeout=10,
        max_backoff_duration=128,
        **kwargs
    ):
        if not nsqd_tcp_addresses and not lookupd_http_addresses:
            raise ValueError('must specify at least on nsqd or lookupd')

        if isinstance(nsqd_tcp_addresses, basestring):
            self.nsqd_tcp_addresses = [nsqd_tcp_addresses]
        elif isinstance(nsqd_tcp_addresses, (list, tuple)):
            self.nsqd_tcp_addresses = set(nsqd_tcp_addresses)
        else:
            raise TypeError('nsqd_tcp_addresses must be a list, set or tuple')

        if isinstance(lookupd_http_addresses, basestring):
            lookupd_http_addresses = [lookupd_http_addresses]
        elif isinstance(lookupd_http_addresses, (list, tuple)):
            lookupd_http_addresses = list(lookupd_http_addresses)
            random.shuffle(lookupd_http_addresses)
        else:
            msg = 'lookupd_http_addresses must be a list, set or tuple'
            raise TypeError(msg)

        self.lookupds = [Lookupd(a) for a in lookupd_http_addresses]
        self.iterlookupds = cycle(self.lookupds)

        self.topic = topic
        self.channel = channel
        self.async = async
        self.max_tries = max_tries
        self.max_in_flight = max_in_flight
        self.requeue_delay = requeue_delay
        self.lookupd_poll_interval = lookupd_poll_interval
        self.lookupd_poll_jitter = lookupd_poll_jitter
        self.low_ready_idle_timeout = low_ready_idle_timeout
        self.max_backoff_duration = max_backoff_duration
        self.conn_kwargs = kwargs

        if name:
            self.name = name
        else:
            self.name = '{}.{}.{}'.format(__name__, self.topic, self.channel)

        self._need_ready_redistributed = False
        self.last_random_ready = time.time()
        self.state = INIT

        self.logger = logging.getLogger(self.name)
        self.conn_backoffs = defaultdict(self.create_backoff)
        self.backoff = self.create_backoff()

        self.on_response = blinker.Signal()
        self.on_error = blinker.Signal()
        self.on_message = blinker.Signal()
        self.on_finish = blinker.Signal()
        self.on_requeue = blinker.Signal()
        self.on_giving_up = blinker.Signal()
        self.on_exception = blinker.Signal()

        self.conns = set()
        self.pending = set()

        self.workers = []
        self.conn_workers = {}

    def start(self, block=True):
        if not self.state == INIT:
            return

        self.state = RUNNING
        self.query_nsqd()

        if self.lookupds:
            self.query_lookupd()
            self.workers.append(gevent.spawn(self._poll_lookupd))

        self.workers.append(gevent.spawn(self._poll_ready))

        if block:
            self.join()

    def close(self, block=True):
        if not self.is_running:
            return

        self.state = CLOSED

        for worker in self.workers:
            worker.kill(block=block)

        for conn in self.conns:
            conn.close_stream()

        if block:
            self.join()

    def join(self, timeout=None, raise_error=False):
        gevent.joinall(self.workers, timeout, raise_error)
        gevent.joinall(self.conn_workers.values(), timeout, raise_error)

    @property
    def is_running(self):
        return self.state in (RUNNING, BACKOFF, THROTTLED)

    @property
    def is_starved(self):
        for conn in self.conns:
            if conn.in_flight >= max(conn.last_ready * 0.85, 1):
                return True
        return False

    @property
    def connection_max_in_flight(self):
        return max(1, self.max_in_flight / max(1, len(self.conns)))

    @property
    def total_ready_count(self):
        return sum(c.ready_count for c in self.conns)

    def send_ready(self, conn, count):
        if (self.total_ready_count + count) > self.max_in_flight:
            return
        conn.ready(count)

    def query_nsqd(self):
        self.logger.debug('querying nsqd...')
        for address in self.nsqd_tcp_addresses:
            address, port = address.split(':')
            conn = Nsqd(address, int(port), **self.conn_kwargs)
            self.connect_to_nsqd(conn)

    def query_lookupd(self):
        self.logger.debug('querying lookupd...')
        lookupd = self.iterlookupds.next()

        try:
            producers = lookupd.lookup(self.topic)['producers']
            self.logger.debug('found {} producers'.format(len(producers)))

        except Exception as error:
            msg = 'Failed to lookup {} on {} ({})'
            self.logger.warn(msg.format(self.topic, lookupd.address, error))
            return

        for producer in producers:
            conn = Nsqd(
                producer.get('broadcast_address') or producer['address'],
                producer['tcp_port'],
                producer['http_port'],
                **self.conn_kwargs
            )
            self.connect_to_nsqd(conn)

    def create_backoff(self):
        return BackoffTimer(max_interval=self.max_backoff_duration)

    def start_backoff(self):
        self.state = BACKOFF

        for conn in self.conns:
            conn.ready(0)

        interval = self.backoff.get_interval()
        self.logger.info('backing off for {} seconds'.format(interval))
        gevent.sleep(interval)

        self.state = THROTTLED
        if not self.conns:
            return

        conn = self.random_connection()
        self.logger.info('[{}] testing backoff state with RDY 1'.format(conn))
        self.send_ready(conn, 1)

    def complete_backoff(self):
        self.state = RUNNING
        self.logger.info('backoff complete, resuming normal operation')

        count = self.connection_max_in_flight
        for conn in self.conns.values():
            self.send_ready(conn, count)

    def _poll_lookupd(self):
        delay = self.lookupd_poll_interval * self.lookupd_poll_jitter
        gevent.sleep(random.random() * delay)

        while True:
            gevent.sleep(self.lookupd_poll_interval)
            self.query_lookupd()

    def _poll_ready(self):
        while True:
            gevent.sleep(5)
            if not self.need_ready_redistributed:
                continue
            self.redistribute_ready_state()

    @property
    def need_ready_redistributed(self):
        if self.state == BACKOFF:
            return False

        if self._need_ready_redistributed:
            return True

        if len(self.conns) > self.max_in_flight:
            return True

        if self.state == THROTTLED and len(self.conns) > 1:
            return True

    @need_ready_redistributed.setter
    def need_ready_redistributed(self, value):
        self._need_ready_redistributed = value

    def redistribute_ready_state(self):
        self.need_ready_redistributed = False

        # first set RDY 0 to all connections that have not received a message
        # within a configurable timeframe (low_ready_idle_timeout).
        for conn in self.conns:
            if conn.ready_count == 0:
                continue

            if (time.time() - conn.last_message) < self.low_ready_idle_timeout:
                continue

            msg = '[{}] idle connection, giving up RDY count'.format(conn)
            self.logger.info(msg)

            conn.ready(0)

        if self.state == THROTTLED:
            max_in_flight = 1 - self.total_ready
        else:
            max_in_flight = self.max_in_flight - self.total_ready

        if max_in_flight <= 0:
            return

        # randomly walk the list of possible connections and send RDY 1 (up to
        # our calculate "max_in_flight").  We only need to send RDY 1 because in
        # both cases described above your per connection RDY count would never
        # be higher.
        #
        # We also don't attempt to avoid the connections who previously might
        # have had RDY 1 because it would be overly complicated and not actually
        # worth it (ie. given enough redistribution rounds it doesn't matter).
        conns = list(self.conns)
        conns = random.sample(conns, min(max_in_flight, len(self.conn)))

        for conn in conns:
            self.logger.info('[{}] redistributing RDY'.format(conn))
            self.send_ready(conn, 1)

    def random_ready_conn(self, conn):
        # if all connections aren't getting RDY
        # occsionally randomize which connection gets RDY
        if len(self.conns) <= self.max_in_flight:
            return conn

        if (time.time() - self.last_random_ready) < 30:
            return conn

        self.last_random_ready = time.time()
        return random.choice(c for c in self.conns if not c.ready_count)

    def update_ready(self, conn):
        if self.state in (BACKOFF, THROTTLED):
            return

        conn = self.random_ready_conn(conn)
        if conn.ready_count < max(conn.last_ready * 0.25, 2):
            self.send_ready(conn, self.connection_max_in_flight)

    def random_connection(self):
        if not self.conns:
            return None
        return random.choice(list(self.conns))

    def publish(self, topic, message):
        conn = self.random_connection()
        if conn is None:
            raise NSQNoConnections()

        conn.publish(topic, message)

    def connect_to_nsqd(self, conn):
        if not self.is_running:
            return

        if conn in self.conns:
            self.logger.debug('[{}] already connected'.format(conn))
            return

        if conn in self.pending:
            self.logger.debug('[{}] already pending'.format(conn))
            return

        self.logger.debug('[{}] connecting...'.format(conn))

        conn.on_response.connect(self.handle_response)
        conn.on_error.connect(self.handle_error)
        conn.on_message.connect(self.handle_message)
        conn.on_finish.connect(self.handle_finish)
        conn.on_requeue.connect(self.handle_requeue)

        self.pending.add(conn)

        try:
            conn.connect()
            conn.identify()

            if conn.max_ready_count < self.max_in_flight:
                self.logger.warning(' '.join([
                    '[{}] max RDY count {} < reader max in flight {},',
                    'truncation possible'
                ]).format(conn, conn.max_ready_count, self.max_in_flight))

            conn.subscribe(self.topic, self.channel)

            # Send RDY 1 if we're not backing off or we're the first connection
            if self.state != BACKOFF or not self.conns:
                self.send_ready(conn, 1)

        except NSQException as error:
            msg = '[{}] connection failed ({!r})'.format(conn, error)
            self.logger.debug(msg)
            self.handle_connection_failure(conn)
            return

        finally:
            self.pending.remove(conn)

        # Check if we've closed since we started
        if not self.is_running:
            conn.close_stream()
            return

        self.logger.info('[{}] connection successful'.format(conn))
        self.handle_connection_success(conn)

    def _listen(self, conn):
        try:
            conn.listen()
        except NSQException as error:
            msg = '[{}] connection lost ({!r})'.format(conn, error)
            self.logger.warning(msg)

        self.handle_connection_failure(conn)

    def handle_connection_success(self, conn):
        self.conns.add(conn)
        self.conn_workers[conn] = gevent.spawn(self._listen, conn)

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        self.conn_backoffs[conn].success()

    def handle_connection_failure(self, conn):
        self.conns.discard(conn)
        self.conn_workers.pop(conn, None)
        conn.close_stream()

        if conn.ready_count:
            self.need_ready_redistributed = True

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        seconds = self.conn_backoffs[conn].failure().get_interval()
        self.logger.debug('[{}] retrying in {}s'.format(conn, seconds))
        gevent.spawn_later(seconds, self.connect_to_nsqd, conn)

    def handle_response(self, conn, response):
        self.logger.debug('[{}] response: {}'.format(conn, response))
        self.on_response.send(self, conn=conn, response=response)

    def handle_error(self, conn, error):
        self.logger.debug('[{}] error: {}'.format(conn, error))
        self.on_error.send(self, conn=conn, error=error)

    def handle_message(self, conn, message):
        self.logger.debug('[{}] got message: {}'.format(conn, message.id))
        self.update_ready(conn)

        if self.max_tries and message.attempts > self.max_tries:
            msg = "giving up on message '{}' after max tries {}"
            self.logger.warning(msg.format(message.id, self.max_tries))
            self.on_giving_up.send(self, conn, message)
            return message.finish()

        try:
            self.on_message.send(self, conn=conn, message=message)
            if not self.async:
                message.finish()
            return

        except NSQRequeueMessage:
            pass

        except Exception as error:
            msg = '[{}] caught exception while handling message'.format(conn)
            self.logger.exception(msg)
            self.on_exception(self, conn=conn, message=message, error=error)

        message.requeue(self.requeue_delay)

    def handle_finish(self, conn, message_id):
        self.logger.debug('[{}] finished message: {}'.format(conn, message_id))
        self.backoff.success()
        self.handle_backoff()
        self.on_finish.send(self, conn=conn, message_id=message_id)

    def handle_requeue(self, conn, message_id, timeout):
        msg = '[{}] requeued message: {} ({})'
        self.logger.debug(msg.format(conn, message_id, timeout))
        self.backoff.failure()
        self.handle_backoff()
        self.on_requeue.send(
            self,
            conn=conn,
            message_id=message_id,
            timeout=timeout
        )

    def handle_backoff(self):
        if self.state == BACKOFF:
            return

        if self.state == THROTTLED and self.backoff.is_reset():
            return self.complete_backoff()

        if not self.backoff.is_reset():
            return self.start_backoff()
<commit_msg>Fix syntax.<commit_after>import logging
import random
import gevent
import blinker
import time

from itertools import cycle
from collections import defaultdict

from .lookupd import Lookupd
from .nsqd import Nsqd
from .backofftimer import BackoffTimer
from .states import INIT, RUNNING, BACKOFF, THROTTLED, CLOSED

from .errors import (
    NSQException,
    NSQNoConnections,
    NSQRequeueMessage
)

class Reader(object):
    def __init__(
        self,
        topic,
        channel,
        nsqd_tcp_addresses=[],
        lookupd_http_addresses=[],
        name=None,
        async=False,
        max_tries=5,
        max_in_flight=1,
        requeue_delay=0,
        lookupd_poll_interval=60,
        lookupd_poll_jitter=0.3,
        low_ready_idle_timeout=10,
        max_backoff_duration=128,
        **kwargs
    ):
        if not nsqd_tcp_addresses and not lookupd_http_addresses:
            raise ValueError('must specify at least on nsqd or lookupd')

        if isinstance(nsqd_tcp_addresses, basestring):
            self.nsqd_tcp_addresses = [nsqd_tcp_addresses]
        elif isinstance(nsqd_tcp_addresses, (list, tuple)):
            self.nsqd_tcp_addresses = set(nsqd_tcp_addresses)
        else:
            raise TypeError('nsqd_tcp_addresses must be a list, set or tuple')

        if isinstance(lookupd_http_addresses, basestring):
            lookupd_http_addresses = [lookupd_http_addresses]
        elif isinstance(lookupd_http_addresses, (list, tuple)):
            lookupd_http_addresses = list(lookupd_http_addresses)
            random.shuffle(lookupd_http_addresses)
        else:
            msg = 'lookupd_http_addresses must be a list, set or tuple'
            raise TypeError(msg)

        self.lookupds = [Lookupd(a) for a in lookupd_http_addresses]
        self.iterlookupds = cycle(self.lookupds)

        self.topic = topic
        self.channel = channel
        self.async = async
        self.max_tries = max_tries
        self.max_in_flight = max_in_flight
        self.requeue_delay = requeue_delay
        self.lookupd_poll_interval = lookupd_poll_interval
        self.lookupd_poll_jitter = lookupd_poll_jitter
        self.low_ready_idle_timeout = low_ready_idle_timeout
        self.max_backoff_duration = max_backoff_duration
        self.conn_kwargs = kwargs

        if name:
            self.name = name
        else:
            self.name = '{}.{}.{}'.format(__name__, self.topic, self.channel)

        self._need_ready_redistributed = False
        self.last_random_ready = time.time()
        self.state = INIT

        self.logger = logging.getLogger(self.name)
        self.conn_backoffs = defaultdict(self.create_backoff)
        self.backoff = self.create_backoff()

        self.on_response = blinker.Signal()
        self.on_error = blinker.Signal()
        self.on_message = blinker.Signal()
        self.on_finish = blinker.Signal()
        self.on_requeue = blinker.Signal()
        self.on_giving_up = blinker.Signal()
        self.on_exception = blinker.Signal()

        self.conns = set()
        self.pending = set()

        self.workers = []
        self.conn_workers = {}

    def start(self, block=True):
        if self.state != INIT:
            self.logger.debug('{} all ready started'.format(self.name))
            return

        self.logger.debug('starting {}...'.format(self.name))
        self.state = RUNNING
        self.query_nsqd()

        if self.lookupds:
            self.query_lookupd()
            self.workers.append(gevent.spawn(self._poll_lookupd))

        self.workers.append(gevent.spawn(self._poll_ready))

        if block:
            self.join()

    def close(self, block=True):
        if not self.is_running:
            return

        self.state = CLOSED

        for worker in self.workers:
            worker.kill(block=block)

        for conn in self.conns:
            conn.close_stream()

        if block:
            self.join()

    def join(self, timeout=None, raise_error=False):
        gevent.joinall(self.workers, timeout, raise_error)
        gevent.joinall(self.conn_workers.values(), timeout, raise_error)

    @property
    def is_running(self):
        return self.state in (RUNNING, BACKOFF, THROTTLED)

    @property
    def is_starved(self):
        for conn in self.conns:
            if conn.in_flight >= max(conn.last_ready * 0.85, 1):
                return True
        return False

    @property
    def connection_max_in_flight(self):
        return max(1, self.max_in_flight / max(1, len(self.conns)))

    @property
    def total_ready_count(self):
        return sum(c.ready_count for c in self.conns)

    def send_ready(self, conn, count):
        if (self.total_ready_count + count) > self.max_in_flight:
            return
        conn.ready(count)

    def query_nsqd(self):
        self.logger.debug('querying nsqd...')
        for address in self.nsqd_tcp_addresses:
            address, port = address.split(':')
            conn = Nsqd(address, int(port), **self.conn_kwargs)
            self.connect_to_nsqd(conn)

    def query_lookupd(self):
        self.logger.debug('querying lookupd...')
        lookupd = self.iterlookupds.next()

        try:
            producers = lookupd.lookup(self.topic)['producers']
            self.logger.debug('found {} producers'.format(len(producers)))

        except Exception as error:
            msg = 'Failed to lookup {} on {} ({})'
            self.logger.warn(msg.format(self.topic, lookupd.address, error))
            return

        for producer in producers:
            conn = Nsqd(
                producer.get('broadcast_address') or producer['address'],
                producer['tcp_port'],
                producer['http_port'],
                **self.conn_kwargs
            )
            self.connect_to_nsqd(conn)

    def create_backoff(self):
        return BackoffTimer(max_interval=self.max_backoff_duration)

    def start_backoff(self):
        self.state = BACKOFF

        for conn in self.conns:
            conn.ready(0)

        interval = self.backoff.get_interval()
        self.logger.info('backing off for {} seconds'.format(interval))
        gevent.sleep(interval)

        self.state = THROTTLED
        if not self.conns:
            return

        conn = self.random_connection()
        self.logger.info('[{}] testing backoff state with RDY 1'.format(conn))
        self.send_ready(conn, 1)

    def complete_backoff(self):
        self.state = RUNNING
        self.logger.info('backoff complete, resuming normal operation')

        count = self.connection_max_in_flight
        for conn in self.conns.values():
            self.send_ready(conn, count)

    def _poll_lookupd(self):
        delay = self.lookupd_poll_interval * self.lookupd_poll_jitter
        gevent.sleep(random.random() * delay)

        while True:
            gevent.sleep(self.lookupd_poll_interval)
            self.query_lookupd()

    def _poll_ready(self):
        while True:
            gevent.sleep(5)
            if not self.need_ready_redistributed:
                continue
            self.redistribute_ready_state()

    @property
    def need_ready_redistributed(self):
        if self.state == BACKOFF:
            return False

        if self._need_ready_redistributed:
            return True

        if len(self.conns) > self.max_in_flight:
            return True

        if self.state == THROTTLED and len(self.conns) > 1:
            return True

    @need_ready_redistributed.setter
    def need_ready_redistributed(self, value):
        self._need_ready_redistributed = value

    def redistribute_ready_state(self):
        self.need_ready_redistributed = False

        # first set RDY 0 to all connections that have not received a message
        # within a configurable timeframe (low_ready_idle_timeout).
        for conn in self.conns:
            if conn.ready_count == 0:
                continue

            if (time.time() - conn.last_message) < self.low_ready_idle_timeout:
                continue

            msg = '[{}] idle connection, giving up RDY count'.format(conn)
            self.logger.info(msg)

            conn.ready(0)

        if self.state == THROTTLED:
            max_in_flight = 1 - self.total_ready
        else:
            max_in_flight = self.max_in_flight - self.total_ready

        if max_in_flight <= 0:
            return

        # randomly walk the list of possible connections and send RDY 1 (up to
        # our calculate "max_in_flight").  We only need to send RDY 1 because in
        # both cases described above your per connection RDY count would never
        # be higher.
        #
        # We also don't attempt to avoid the connections who previously might
        # have had RDY 1 because it would be overly complicated and not actually
        # worth it (ie. given enough redistribution rounds it doesn't matter).
        conns = list(self.conns)
        conns = random.sample(conns, min(max_in_flight, len(self.conn)))

        for conn in conns:
            self.logger.info('[{}] redistributing RDY'.format(conn))
            self.send_ready(conn, 1)

    def random_ready_conn(self, conn):
        # if all connections aren't getting RDY
        # occsionally randomize which connection gets RDY
        if len(self.conns) <= self.max_in_flight:
            return conn

        if (time.time() - self.last_random_ready) < 30:
            return conn

        self.last_random_ready = time.time()
        return random.choice(c for c in self.conns if not c.ready_count)

    def update_ready(self, conn):
        if self.state in (BACKOFF, THROTTLED):
            return

        conn = self.random_ready_conn(conn)
        if conn.ready_count < max(conn.last_ready * 0.25, 2):
            self.send_ready(conn, self.connection_max_in_flight)

    def random_connection(self):
        if not self.conns:
            return None
        return random.choice(list(self.conns))

    def publish(self, topic, message):
        conn = self.random_connection()
        if conn is None:
            raise NSQNoConnections()

        conn.publish(topic, message)

    def connect_to_nsqd(self, conn):
        if not self.is_running:
            return

        if conn in self.conns:
            self.logger.debug('[{}] already connected'.format(conn))
            return

        if conn in self.pending:
            self.logger.debug('[{}] already pending'.format(conn))
            return

        self.logger.debug('[{}] connecting...'.format(conn))

        conn.on_response.connect(self.handle_response)
        conn.on_error.connect(self.handle_error)
        conn.on_message.connect(self.handle_message)
        conn.on_finish.connect(self.handle_finish)
        conn.on_requeue.connect(self.handle_requeue)

        self.pending.add(conn)

        try:
            conn.connect()
            conn.identify()

            if conn.max_ready_count < self.max_in_flight:
                self.logger.warning(' '.join([
                    '[{}] max RDY count {} < reader max in flight {},',
                    'truncation possible'
                ]).format(conn, conn.max_ready_count, self.max_in_flight))

            conn.subscribe(self.topic, self.channel)

            # Send RDY 1 if we're not backing off or we're the first connection
            if self.state != BACKOFF or not self.conns:
                self.send_ready(conn, 1)

        except NSQException as error:
            msg = '[{}] connection failed ({!r})'.format(conn, error)
            self.logger.debug(msg)
            self.handle_connection_failure(conn)
            return

        finally:
            self.pending.remove(conn)

        # Check if we've closed since we started
        if not self.is_running:
            conn.close_stream()
            return

        self.logger.info('[{}] connection successful'.format(conn))
        self.handle_connection_success(conn)

    def _listen(self, conn):
        try:
            conn.listen()
        except NSQException as error:
            msg = '[{}] connection lost ({!r})'.format(conn, error)
            self.logger.warning(msg)

        self.handle_connection_failure(conn)

    def handle_connection_success(self, conn):
        self.conns.add(conn)
        self.conn_workers[conn] = gevent.spawn(self._listen, conn)

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        self.conn_backoffs[conn].success()

    def handle_connection_failure(self, conn):
        self.conns.discard(conn)
        self.conn_workers.pop(conn, None)
        conn.close_stream()

        if conn.ready_count:
            self.need_ready_redistributed = True

        if str(conn) not in self.nsqd_tcp_addresses:
            return

        seconds = self.conn_backoffs[conn].failure().get_interval()
        self.logger.debug('[{}] retrying in {}s'.format(conn, seconds))
        gevent.spawn_later(seconds, self.connect_to_nsqd, conn)

    def handle_response(self, conn, response):
        self.logger.debug('[{}] response: {}'.format(conn, response))
        self.on_response.send(self, conn=conn, response=response)

    def handle_error(self, conn, error):
        self.logger.debug('[{}] error: {}'.format(conn, error))
        self.on_error.send(self, conn=conn, error=error)

    def handle_message(self, conn, message):
        self.logger.debug('[{}] got message: {}'.format(conn, message.id))
        self.update_ready(conn)

        if self.max_tries and message.attempts > self.max_tries:
            msg = "giving up on message '{}' after max tries {}"
            self.logger.warning(msg.format(message.id, self.max_tries))
            self.on_giving_up.send(self, conn, message)
            return message.finish()

        try:
            self.on_message.send(self, conn=conn, message=message)
            if not self.async:
                message.finish()
            return

        except NSQRequeueMessage:
            pass

        except Exception as error:
            msg = '[{}] caught exception while handling message'.format(conn)
            self.logger.exception(msg)
            self.on_exception(self, conn=conn, message=message, error=error)

        message.requeue(self.requeue_delay)

    def handle_finish(self, conn, message_id):
        self.logger.debug('[{}] finished message: {}'.format(conn, message_id))
        self.backoff.success()
        self.handle_backoff()
        self.on_finish.send(self, conn=conn, message_id=message_id)

    def handle_requeue(self, conn, message_id, timeout):
        msg = '[{}] requeued message: {} ({})'
        self.logger.debug(msg.format(conn, message_id, timeout))
        self.backoff.failure()
        self.handle_backoff()
        self.on_requeue.send(
            self,
            conn=conn,
            message_id=message_id,
            timeout=timeout
        )

    def handle_backoff(self):
        if self.state == BACKOFF:
            return

        if self.state == THROTTLED and self.backoff.is_reset():
            return self.complete_backoff()

        if not self.backoff.is_reset():
            return self.start_backoff()
<|endoftext|>


<commit_before>""" pyservice runner """

import sys
import pyservice

from pyservice.utils import load_process

def run_service():
    # TODO fix issue with correct help presentation when process is not exist
    # TODO incorrect handling when pid file is empty or None

    import argparse

    parser = argparse.ArgumentParser(prog="pyservice", add_help=False)
    parser.add_argument("-v", "--version",
        action="version", version="%(prog)s, v.{}".format(pyservice.__version__))
    parser.add_argument("-h", "--help",
        action="store_true", help="show program's help text and exit")
    parser.add_argument("process", nargs='?', help="""
        process class path to run (modulename.ProcessClass) or
        configuration file path to use (/path/to/config.py)
        """.strip())
    parser.add_argument("action", nargs='?',
        choices="start stop restart reload status".split())

    try:
        args = parser.parse_args()
    except TypeError:
        parser.print_help()
        return

    if args.help:
        parser.print_help()
        return
    try:
        if args.process and args.action in "start stop restart reload status".split():
            if not args.process:
                parser.error("You need to specify a process for {}".format(args.action))
            getattr(ServiceControl(), args.action)(args.process)
        else:
            parser.print_help()
    except RuntimeError, e:
        parser.error(e)

class ServiceControl(object):

    def start(self, process_path):
        ''' start process '''

        print "Starting process with {}...".format(process_path)
        process = load_process(process_path)
        if callable(process):
            pyservice.core.Service(process).start()
        else:
            raise RuntimeError("The process {} is not valid".format(process_path))

    def stop(self, process_path):
        if self._validate(pid):
            print "Stopping process {}...".format(pid)
            os.kill(pid, STOP_SIGNAL)

    def restart(self, process_path):
        self.stop(resolve_pid(target=target))
        self.start(target)

    def status(self, process_path):
        if self._validate(pid):
            print "Process is running as {}.".format(pid)

    def _validate(self, process_path):
        try:
            os.kill(pid, 0)
            return pid
        except (OSError, TypeError):
            print "Process is NOT running."

<commit_msg>refactoring runner.py<commit_after>""" pyservice runner """

import sys
import pyservice

from pyservice.utils import load_process

def run_service():
    # TODO fix issue with correct help presentation when process is not exist
    # TODO incorrect handling when pid file is empty or None

    import argparse

    parser = argparse.ArgumentParser(prog="pyservice", add_help=False)
    parser.add_argument("-v", "--version",
        action="version", version="%(prog)s, v.{}".format(pyservice.__version__))
    parser.add_argument("-h", "--help",
        action="store_true", help="show program's help text and exit")
    parser.add_argument("process", nargs='?', help="""
        process class path to run (modulename.ProcessClass) or
        configuration file path to use (/path/to/config.py)
        """.strip())
    parser.add_argument("action", nargs='?',
        choices="start stop restart reload status".split())

    try:
        args = parser.parse_args()
    except TypeError:
        parser.print_help()
        return

    if args.help:
        parser.print_help()
        return
    try:
        if args.process and args.action in "start stop restart reload status".split():
            if not args.process:
                parser.error("You need to specify a process for {}".format(args.action))
            getattr(ServiceControl(args.process), args.action)()
        else:
            parser.print_help()
    except RuntimeError, e:
        parser.error(e)

class ServiceControl(object):

    def __init__(self, process_path):
        self.process = load_process(process_path)
        if not callable(self.process):
            raise RuntimeError("The process {} is not valid".format(self.process_path))

    def start(self):

        print "Starting process with {}...".format(self.process.__name__)
        pyservice.core.Service(self.process).start()

    def stop(self):

        print "Stopping process {}...".format(self.process.__name__)
        pyservice.core.Service(self.process).stop()

    def restart(self):

        print 'Restarting {} process'.format(self.process.__name__)
        self.stop()
        self.start()
        print 'Process {} was restarted'.format(self.process.__name__)

    def status(self, process_path):

        if self._validate(pid):
            print "Process is running as {}.".format(pid)

    def _validate(self, process_path):
        try:
            os.kill(pid, 0)
            return pid
        except (OSError, TypeError):
            print "Process is NOT running."

<|endoftext|>


<commit_before># =============================================================================
#
# Copyright (c) 2016, Cisco Systems
# All rights reserved.
#
# # Author: Klaudiusz Staniek
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
# =============================================================================

import os

from csmpe.plugins import CSMPlugin
from utils import get_filesystems

class Plugin(CSMPlugin):
    """This plugin captures custom commands and stores in the log directory."""
    name = "Disk Space Check Plugin"
    platforms = {'ASR9K', 'CRS'}
    phases = {'Pre-Add'}
    os = {'XR'}

    def _get_pie_size(self, package_url):
        """
        # admin show install pie-info ftp://terastream:*****@172.20.168.195//home/terastream/4.3.1FCS/asr9k-doc-px.pie-4.3.1
        Contents of pie file '/ftp://terastream:*****@172.20.168.195//home/terastream/4.3.1FCS/asr9k-doc-px.pie-4.3.1':
            Expiry date       : < NOT AVAILABLE
            Uncompressed size : 25792000
            Compressed size   : 5602232
            Size on disk      : 5668352

            asr9k-doc-px-4.3.1
                asr9K-doc-supp-4.3.1
        """
        cmd = "admin show install pie-info " + package_url
        output = self.ctx.send(cmd, timeout=3600)
        if output:
            for line in output.split('\n'):
                if "Compressed" in line:
                    size = long(line.split(":")[1].strip())
                    return size
                if line and line[:6] == "Error:":
                    self.ctx.error(output)

    def run(self):
        try:
            packages = self.ctx.software_packages
        except AttributeError:
            self.ctx.warning("No package list provided. Skipping calculation of required free harddisk space.")
            return

        try:
            server_repository_url = self.ctx.server_repository_url
        except AttributeError:
            self.ctx.warning("No repository path provided.")
            return

        if server_repository_url[:4] != 'tftp' and server_repository_url[:3] != 'ftp':
            self.ctx.info('Skipping as disk space check only works for TFTP/FTP.')
            return

        file_systems = get_filesystems(self.ctx)
        disk0 = file_systems.get('disk0:', None)
        if disk0:
            disk0_free = disk0.get('free', 0)
        else:
            self.ctx.error("No filesystem 'disk0:' on active RP.")
            return

        total_size = 0
        for package in packages:
            if package == "":
                continue

            if 'pie' not in package:
                self.ctx.info("Package: {} cannot be checked, disk space check result will not be accurate.".format(package))
                continue

            package_url = os.path.join(server_repository_url, package)
            size = self._get_pie_size(package_url)
            total_size += size
            self.ctx.info("Package: {} requires {} bytes.".format(package, size))

        self.ctx.info("Total (required/available): {}/{} bytes".format(total_size, disk0_free))
        if disk0_free < total_size:
            self.ctx.error("Not enough space on disk0: to install packages. The install process can't proceed.")
        else:
            self.ctx.info("There is enough space on disk0: to install packages.")
        return True
<commit_msg>Check image uncompressed size<commit_after># =============================================================================
#
# Copyright (c) 2016, Cisco Systems
# All rights reserved.
#
# # Author: Klaudiusz Staniek
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
# Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
# THE POSSIBILITY OF SUCH DAMAGE.
# =============================================================================

import os

from csmpe.plugins import CSMPlugin
from utils import get_filesystems

class Plugin(CSMPlugin):
    """This plugin captures custom commands and stores in the log directory."""
    name = "Disk Space Check Plugin"
    platforms = {'ASR9K', 'CRS'}
    phases = {'Pre-Add'}
    os = {'XR'}

    def _get_pie_size(self, package_url):
        """
        # admin show install pie-info ftp://terastream:*****@172.20.168.195//home/terastream/4.3.1FCS/asr9k-doc-px.pie-4.3.1
        Contents of pie file '/ftp://terastream:*****@172.20.168.195//home/terastream/4.3.1FCS/asr9k-doc-px.pie-4.3.1':
            Expiry date       : < NOT AVAILABLE
            Uncompressed size : 25792000
            Compressed size   : 5602232
            Size on disk      : 5668352

            asr9k-doc-px-4.3.1
                asr9K-doc-supp-4.3.1
        """
        cmd = "admin show install pie-info " + package_url
        output = self.ctx.send(cmd, timeout=3600)
        if output:
            for line in output.split('\n'):
                if "Uncompressed" in line:
                    size = long(line.split(":")[1].strip())
                    return size
                if line and line[:6] == "Error:":
                    self.ctx.error(output)

    def run(self):
        try:
            packages = self.ctx.software_packages
        except AttributeError:
            self.ctx.warning("No package list provided. Skipping calculation of required free harddisk space.")
            return

        try:
            server_repository_url = self.ctx.server_repository_url
        except AttributeError:
            self.ctx.warning("No repository path provided.")
            return

        if server_repository_url[:4] != 'tftp' and server_repository_url[:3] != 'ftp':
            self.ctx.info('Skipping as disk space check only works for TFTP/FTP.')
            return

        file_systems = get_filesystems(self.ctx)
        disk0 = file_systems.get('disk0:', None)
        if disk0:
            disk0_free = disk0.get('free', 0)
        else:
            self.ctx.error("No filesystem 'disk0:' on active RP.")
            return

        total_size = 0
        for package in packages:
            if package == "":
                continue

            if 'pie' not in package:
                self.ctx.info("Package: {} cannot be checked, disk space check result will not be accurate.".format(package))
                continue

            package_url = os.path.join(server_repository_url, package)
            size = self._get_pie_size(package_url)
            total_size += size
            self.ctx.info("Package: {} requires {} bytes.".format(package, size))

        self.ctx.info("Total (required/available): {}/{} bytes".format(total_size, disk0_free))
        if disk0_free < total_size:
            self.ctx.error("Not enough space on disk0: to install packages. The install process can't proceed.")
        else:
            self.ctx.info("There is enough space on disk0: to install packages.")
        return True
<|endoftext|>


<commit_before>from datetime import datetime

from django.conf import settings
from celery.task import periodic_task, task
from celery.utils.log import get_task_logger
from redis.exceptions import LockError
from dimagi.utils.couch.cache.cache_core import get_redis_client
from dimagi.utils.couch.undo import DELETED_SUFFIX

from corehq.apps.repeaters.dbaccessors import iterate_repeat_records
from corehq.apps.repeaters.const import (
    CHECK_REPEATERS_INTERVAL,
    CHECK_REPEATERS_KEY,
)
from corehq.apps.repeaters.models import simple_post_cache_value

logging = get_task_logger(__name__)

@periodic_task(
    run_every=CHECK_REPEATERS_INTERVAL,
    queue=settings.CELERY_PERIODIC_QUEUE,
)
def check_repeaters():
    start = datetime.utcnow()
    cutoff = start + CHECK_REPEATERS_INTERVAL

    redis_client = get_redis_client().client.get_client()

    # Timeout for slightly less than periodic check
    check_repeater_lock = redis_client.lock(
        CHECK_REPEATERS_KEY,
        timeout=CHECK_REPEATERS_INTERVAL.seconds - 10
    )
    if not check_repeater_lock.acquire(blocking=False):
        return

    for record in iterate_repeat_records(start):
        now = datetime.utcnow()
        lock_key = _get_repeat_record_lock_key(record)

        if now > cutoff:
            break

        lock = redis_client.lock(lock_key, timeout=60 * 60 * 48)
        if not lock.acquire(blocking=False):
            continue

        # Short circuit task creation if we already know the URL is failing
        result = simple_post_cache_value(record.url)
        if result:
            record.update_failure(result)
            record.save()
            continue

        process_repeat_record.delay(record)

    try:
        check_repeater_lock.release()
    except LockError:
        # Ignore if already released
        pass

@task(queue=settings.CELERY_REPEAT_RECORD_QUEUE)
def process_repeat_record(repeat_record):
    try:
        if repeat_record.repeater.doc_type.endswith(DELETED_SUFFIX):
            if not repeat_record.doc_type.endswith(DELETED_SUFFIX):
                repeat_record.doc_type += DELETED_SUFFIX
        else:
            repeat_record.fire()
        repeat_record.save()
    except Exception:
        logging.exception('Failed to process repeat record: {}'.format(repeat_record._id))

def _get_repeat_record_lock_key(record):
    """
    Including the rev in the key means that the record will be unlocked for processing
    every time we execute a `save()` call.
    """
    return 'repeat_record_in_progress-{}_{}'.format(record._id, record._rev)
<commit_msg>ensure that failure reason is string<commit_after>from datetime import datetime

from django.conf import settings
from celery.task import periodic_task, task
from celery.utils.log import get_task_logger
from redis.exceptions import LockError
from dimagi.utils.couch.cache.cache_core import get_redis_client
from dimagi.utils.couch.undo import DELETED_SUFFIX

from corehq.apps.repeaters.dbaccessors import iterate_repeat_records
from corehq.apps.repeaters.const import (
    CHECK_REPEATERS_INTERVAL,
    CHECK_REPEATERS_KEY,
)
from corehq.apps.repeaters.models import simple_post_cache_value

logging = get_task_logger(__name__)

@periodic_task(
    run_every=CHECK_REPEATERS_INTERVAL,
    queue=settings.CELERY_PERIODIC_QUEUE,
)
def check_repeaters():
    start = datetime.utcnow()
    cutoff = start + CHECK_REPEATERS_INTERVAL

    redis_client = get_redis_client().client.get_client()

    # Timeout for slightly less than periodic check
    check_repeater_lock = redis_client.lock(
        CHECK_REPEATERS_KEY,
        timeout=CHECK_REPEATERS_INTERVAL.seconds - 10
    )
    if not check_repeater_lock.acquire(blocking=False):
        return

    for record in iterate_repeat_records(start):
        now = datetime.utcnow()
        lock_key = _get_repeat_record_lock_key(record)

        if now > cutoff:
            break

        lock = redis_client.lock(lock_key, timeout=60 * 60 * 48)
        if not lock.acquire(blocking=False):
            continue

        # Short circuit task creation if we already know the URL is failing
        result = simple_post_cache_value(record.url)
        if result:
            record.update_failure(unicode(result))
            record.save()
            continue

        process_repeat_record.delay(record)

    try:
        check_repeater_lock.release()
    except LockError:
        # Ignore if already released
        pass

@task(queue=settings.CELERY_REPEAT_RECORD_QUEUE)
def process_repeat_record(repeat_record):
    try:
        if repeat_record.repeater.doc_type.endswith(DELETED_SUFFIX):
            if not repeat_record.doc_type.endswith(DELETED_SUFFIX):
                repeat_record.doc_type += DELETED_SUFFIX
        else:
            repeat_record.fire()
        repeat_record.save()
    except Exception:
        logging.exception('Failed to process repeat record: {}'.format(repeat_record._id))

def _get_repeat_record_lock_key(record):
    """
    Including the rev in the key means that the record will be unlocked for processing
    every time we execute a `save()` call.
    """
    return 'repeat_record_in_progress-{}_{}'.format(record._id, record._rev)
<|endoftext|>


<commit_before>import os

import numpy as np
import math
import random
import torch
import torch.utils.data
from torch.utils.data.distributed import DistributedSampler
import torchvision.transforms as transforms
import torchvision.datasets as datasets

from torch.utils.data.sampler import Sampler
import torchvision
import pickle
from tqdm import tqdm
import time
import multiprocessing

FINISH_EVENT = "FINISH_EVENT"
class PaddleDataLoader(object):
    def __init__(self, torch_dataset, indices=None, concurrent=24, queue_size=3072, shuffle=True, shuffle_seed=0):
        self.torch_dataset = torch_dataset
        self.data_queue = multiprocessing.Queue(queue_size)
        self.indices = indices
        self.concurrent = concurrent
        self.shuffle = shuffle
        self.shuffle_seed=shuffle_seed

    def _worker_loop(self, dataset, worker_indices, worker_id):
        cnt = 0
        for idx in worker_indices:
            cnt += 1
            img, label = self.torch_dataset[idx]
            img = np.array(img).astype('uint8').transpose((2, 0, 1))
            self.data_queue.put((img, label))
        print("worker: [%d] read [%d] samples. " % (worker_id, cnt))
        self.data_queue.put(FINISH_EVENT)

    def reader(self):
        def _reader_creator():
            worker_processes = []
            total_img = len(self.torch_dataset)
            print("total image: ", total_img)
            if self.shuffle:
                self.indices = [i for i in xrange(total_img)]
                random.seed(self.shuffle_seed)
                random.shuffle(self.indices)
                print("shuffle indices: %s ..." % self.indices[:10])

            imgs_per_worker = int(math.ceil(total_img / self.concurrent))
            for i in xrange(self.concurrent):
                start = i * imgs_per_worker
                end = (i + 1) * imgs_per_worker if i != self.concurrent - 1 else None
                sliced_indices = self.indices[start:end]
                w = multiprocessing.Process(
                    target=self._worker_loop,
                    args=(self.torch_dataset, sliced_indices, i)
                )
                w.daemon = True
                w.start()
                worker_processes.append(w)
            finish_workers = 0
            worker_cnt = len(worker_processes)
            while finish_workers < worker_cnt:
                sample = self.data_queue.get()
                if sample == FINISH_EVENT:
                    finish_workers += 1
                else:
                    yield sample

        return _reader_creator

def train(traindir, sz, min_scale=0.08, shuffle_seed=0):
    train_tfms = [
        transforms.RandomResizedCrop(sz, scale=(min_scale, 1.0)),
        transforms.RandomHorizontalFlip()
    ]
    train_dataset = datasets.ImageFolder(traindir, transforms.Compose(train_tfms))
    return PaddleDataLoader(train_dataset, shuffle_seed=shuffle_seed).reader()

def test(valdir, bs, sz, rect_val=False):
    if rect_val:
        idx_ar_sorted = sort_ar(valdir)
        idx_sorted, _ = zip(*idx_ar_sorted)
        idx2ar = map_idx2ar(idx_ar_sorted, bs)

        ar_tfms = [transforms.Resize(int(sz* 1.14)), CropArTfm(idx2ar, sz)]
        val_dataset = ValDataset(valdir, transform=ar_tfms)
        return PaddleDataLoader(val_dataset, concurrent=1, indices=idx_sorted, shuffle=False).reader()

    val_tfms = [transforms.Resize(int(sz* 1.14)), transforms.CenterCrop(sz)]
    val_dataset = datasets.ImageFolder(valdir, transforms.Compose(val_tfms))

    return PaddleDataLoader(val_dataset).reader()

class ValDataset(datasets.ImageFolder):
    def __init__(self, root, transform=None, target_transform=None):
        super(ValDataset, self).__init__(root, transform, target_transform)

    def __getitem__(self, index):
        path, target = self.imgs[index]
        sample = self.loader(path)
        if self.transform is not None:
            for tfm in self.transform:
                if isinstance(tfm, CropArTfm):
                    sample = tfm(sample, index)
                else:
                    sample = tfm(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target

class CropArTfm(object):
    def __init__(self, idx2ar, target_size):
        self.idx2ar, self.target_size = idx2ar, target_size

    def __call__(self, img, idx):
        target_ar = self.idx2ar[idx]
        if target_ar < 1:
            w = int(self.target_size / target_ar)
            size = (w // 8 * 8, self.target_size)
        else:
            h = int(self.target_size * target_ar)
            size = (self.target_size, h // 8 * 8)
        return torchvision.transforms.functional.center_crop(img, size)

def sort_ar(valdir):
    idx2ar_file = valdir + '/../sorted_idxar.p'
    if os.path.isfile(idx2ar_file):
        return pickle.load(open(idx2ar_file, 'rb'))
    print('Creating AR indexes. Please be patient this may take a couple minutes...')
    val_dataset = datasets.ImageFolder(valdir)  # AS: TODO: use Image.open instead of looping through dataset
    sizes = [img[0].size for img in tqdm(val_dataset, total=len(val_dataset))]
    idx_ar = [(i, round(s[0] * 1.0/ s[1], 5)) for i, s in enumerate(sizes)]
    sorted_idxar = sorted(idx_ar, key=lambda x: x[1])
    pickle.dump(sorted_idxar, open(idx2ar_file, 'wb'))
    print('Done')
    return sorted_idxar

def chunks(l, n):
    n = max(1, n)
    return (l[i:i + n] for i in range(0, len(l), n))

def map_idx2ar(idx_ar_sorted, batch_size):
    ar_chunks = list(chunks(idx_ar_sorted, batch_size))
    idx2ar = {}
    for chunk in ar_chunks:
        idxs, ars = list(zip(*chunk))
        mean = round(np.mean(ars), 5)
        for idx in idxs:
            idx2ar[idx] = mean
    return idx2ar
<commit_msg>modify single queue to a queue per process for data processing (#2136)<commit_after>from __future__ import division
import os

import numpy as np
import math
import random
import torch
import torch.utils.data
from torch.utils.data.distributed import DistributedSampler
import torchvision.transforms as transforms
import torchvision.datasets as datasets

from torch.utils.data.sampler import Sampler
import torchvision
import pickle
from tqdm import tqdm
import time
import multiprocessing

FINISH_EVENT = "FINISH_EVENT"

class PaddleDataLoader(object):
    def __init__(self,
                 torch_dataset,
                 indices=None,
                 concurrent=24,
                 queue_size=3072,
                 shuffle=True,
                 shuffle_seed=0):
        self.torch_dataset = torch_dataset
        self.indices = indices
        self.concurrent = concurrent
        self.shuffle = shuffle
        self.shuffle_seed = shuffle_seed
        self.queue_size = queue_size // self.concurrent

    def _worker_loop(self, queue, worker_indices, worker_id):
        cnt = 0
        for idx in worker_indices:
            cnt += 1
            img, label = self.torch_dataset[idx]
            img = np.array(img).astype('uint8').transpose((2, 0, 1))
            queue.put((img, label))
        print("worker: [%d] read [%d] samples. " % (worker_id, cnt))
        queue.put(FINISH_EVENT)

    def reader(self):
        def _reader_creator():
            worker_processes = []
            index_queues = []
            total_img = len(self.torch_dataset)
            print("total image: ", total_img)
            if self.shuffle:
                self.indices = [i for i in xrange(total_img)]
                random.seed(self.shuffle_seed)
                random.shuffle(self.indices)
                print("shuffle indices: %s ..." % self.indices[:10])

            imgs_per_worker = int(math.ceil(total_img / self.concurrent))
            for i in xrange(self.concurrent):
                start = i * imgs_per_worker
                end = (i + 1
                       ) * imgs_per_worker if i != self.concurrent - 1 else None
                sliced_indices = self.indices[start:end]
                index_queue = multiprocessing.Queue(self.queue_size)
                w = multiprocessing.Process(
                    target=self._worker_loop,
                    args=(index_queue, sliced_indices, i))
                w.daemon = True
                w.start()
                worker_processes.append(w)
                index_queues.append(index_queue)
            finish_workers = 0
            worker_cnt = len(worker_processes)
            recv_index = 0
            while finish_workers < worker_cnt:
                while (index_queues[recv_index].empty()):
                    recv_index = (recv_index + 1) % self.concurrent
                sample = index_queues[recv_index].get()
                recv_index = (recv_index + 1) % self.concurrent
                if sample == FINISH_EVENT:
                    finish_workers += 1
                else:
                    yield sample

        return _reader_creator

def train(traindir, sz, min_scale=0.08, shuffle_seed=0):
    train_tfms = [
        transforms.RandomResizedCrop(
            sz, scale=(min_scale, 1.0)), transforms.RandomHorizontalFlip()
    ]
    train_dataset = datasets.ImageFolder(traindir,
                                         transforms.Compose(train_tfms))
    return PaddleDataLoader(train_dataset, shuffle_seed=shuffle_seed).reader()

def test(valdir, bs, sz, rect_val=False):
    if rect_val:
        idx_ar_sorted = sort_ar(valdir)
        idx_sorted, _ = zip(*idx_ar_sorted)
        idx2ar = map_idx2ar(idx_ar_sorted, bs)

        ar_tfms = [transforms.Resize(int(sz * 1.14)), CropArTfm(idx2ar, sz)]
        val_dataset = ValDataset(valdir, transform=ar_tfms)
        return PaddleDataLoader(
            val_dataset, concurrent=1, indices=idx_sorted,
            shuffle=False).reader()

    val_tfms = [transforms.Resize(int(sz * 1.14)), transforms.CenterCrop(sz)]
    val_dataset = datasets.ImageFolder(valdir, transforms.Compose(val_tfms))

    return PaddleDataLoader(val_dataset).reader()

class ValDataset(datasets.ImageFolder):
    def __init__(self, root, transform=None, target_transform=None):
        super(ValDataset, self).__init__(root, transform, target_transform)

    def __getitem__(self, index):
        path, target = self.imgs[index]
        sample = self.loader(path)
        if self.transform is not None:
            for tfm in self.transform:
                if isinstance(tfm, CropArTfm):
                    sample = tfm(sample, index)
                else:
                    sample = tfm(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target

class CropArTfm(object):
    def __init__(self, idx2ar, target_size):
        self.idx2ar, self.target_size = idx2ar, target_size

    def __call__(self, img, idx):
        target_ar = self.idx2ar[idx]
        if target_ar < 1:
            w = int(self.target_size / target_ar)
            size = (w // 8 * 8, self.target_size)
        else:
            h = int(self.target_size * target_ar)
            size = (self.target_size, h // 8 * 8)
        return torchvision.transforms.functional.center_crop(img, size)

def sort_ar(valdir):
    idx2ar_file = valdir + '/../sorted_idxar.p'
    if os.path.isfile(idx2ar_file):
        return pickle.load(open(idx2ar_file, 'rb'))
    print(
        'Creating AR indexes. Please be patient this may take a couple minutes...'
    )
    val_dataset = datasets.ImageFolder(
        valdir)  # AS: TODO: use Image.open instead of looping through dataset
    sizes = [img[0].size for img in tqdm(val_dataset, total=len(val_dataset))]
    idx_ar = [(i, round(s[0] * 1.0 / s[1], 5)) for i, s in enumerate(sizes)]
    sorted_idxar = sorted(idx_ar, key=lambda x: x[1])
    pickle.dump(sorted_idxar, open(idx2ar_file, 'wb'))
    print('Done')
    return sorted_idxar

def chunks(l, n):
    n = max(1, n)
    return (l[i:i + n] for i in range(0, len(l), n))

def map_idx2ar(idx_ar_sorted, batch_size):
    ar_chunks = list(chunks(idx_ar_sorted, batch_size))
    idx2ar = {}
    for chunk in ar_chunks:
        idxs, ars = list(zip(*chunk))
        mean = round(np.mean(ars), 5)
        for idx in idxs:
            idx2ar[idx] = mean
    return idx2ar
<|endoftext|>


<commit_before>#!/usr/bin/env python3
import datetime, json, os, requests, shutil, socket, subprocess, time

HOSTS = '/node/etc/hosts'

# POD_NAMESPACE must be explicitly set in deployment yaml using downward api --
# see https://github.com/kubernetes/kubernetes/blob/release-1.0/docs/user-guide/downward-api.md
POD_NAMESPACE = os.environ.get('POD_NAMESPACE', 'default')

def run(v, shell=False, path='.', get_output=False, env=None, verbose=True):
    t = time.time()
    if isinstance(v, str):
        cmd = v
        shell = True
    else:
        cmd = ' '.join([(x if len(x.split())<=1 else '"%s"'%x) for x in v])
    if path != '.':
        cur = os.path.abspath(os.curdir)
        if verbose:
            print('chdir %s'%path)
        os.chdir(path)
    try:
        if verbose:
            print(cmd)
        if shell:
            kwds = {'shell':True, 'executable':'/bin/bash', 'env':env}
        else:
            kwds = {'env':env}
        if get_output:
            output = subprocess.Popen(v, stdout=subprocess.PIPE, **kwds).stdout.read().decode()
        else:
            if subprocess.call(v, **kwds):
                raise RuntimeError("error running '{cmd}'".format(cmd=cmd))
            output = None
        seconds = time.time() - t
        if verbose:
            print("TOTAL TIME: {seconds} seconds -- to run '{cmd}'".format(seconds=seconds, cmd=cmd))
        return output
    finally:
        if path != '.':
            os.chdir(cur)

def get_service(service):
    """
    Get in json format the kubernetes information about the given service.
    """
    if not os.environ['KUBERNETES_SERVICE_HOST']:
        print('KUBERNETES_SERVICE_HOST environment variable not set')
        return None
    URL = "https://{KUBERNETES_SERVICE_HOST}:{KUBERNETES_SERVICE_PORT}/api/v1/namespaces/{POD_NAMESPACE}/endpoints/{service}"
    URL = URL.format(KUBERNETES_SERVICE_HOST=os.environ['KUBERNETES_SERVICE_HOST'],
                     KUBERNETES_SERVICE_PORT=os.environ['KUBERNETES_SERVICE_PORT'],
                     POD_NAMESPACE=POD_NAMESPACE,
                     service=service)
    token = open('/var/run/secrets/kubernetes.io/serviceaccount/token').read()
    headers={'Authorization':'Bearer {token}'.format(token=token)}
    print("Getting k8s information about '{service}' from '{URL}'".format(service=service, URL=URL))
    x = requests.get(URL, headers=headers, verify='/var/run/secrets/kubernetes.io/serviceaccount/ca.crt').json()
    print("Got {x}".format(x=x))
    return x

def update_etc_hosts():
    print('udpate_etc_hosts')
    try:
        v = get_service('storage-projects')
    except Exception as err:
        # Expected to happen when node is starting up, etc. - we'll retry later soon!
        print("Failed getting storage service info", err)
        return
    if v.get('status', None) == 'Failure':
        return
    try:
        if 'addresses' not in v['subsets'][0]:
            return   # nothing to do; no known addresses
        namespace = v['metadata']['namespace']
        hosts = ["{ip}    {namespace}-{name}".format(ip=x['ip'], namespace=namespace,
                              name=x['targetRef']['name'].split('-')[0]) for x in v['subsets'][0]['addresses']]
        start = "# start smc-storage dns - namespace="+namespace+"\n\n"
        end = "# end smc-storage dns - namespace="+namespace+"\n\n"
        block = '\n'.join([start] + hosts + [end])
        current = open(HOSTS).read()
        if block in current:
            return
        i = current.find(start)
        j = current.find(end)
        if i == -1 or j == -1:
            new = current + '\n' + block
        else:
            new = current[:i] + block + current[j+len(end):]
        open(HOSTS,'w').write(new)
    except Exception as err:
        print("Problem in update_etc_hosts", err)

MINION_IP = 'unknown'
def enable_ssh_access_to_minion():
    global MINION_IP
    # create our own local ssh key
    if os.path.exists('/root/.ssh'):
        shutil.rmtree('/root/.ssh')
    run(['ssh-keygen', '-b', '2048', '-N', '', '-f', '/root/.ssh/id_rsa'])
    # make root user of minion allow login using this (and only this) key.
    run('cat /root/.ssh/id_rsa.pub >> /node/root/.ssh/authorized_keys')
    open("/root/.ssh/config",'w').write("StrictHostKeyChecking no\nUserKnownHostsFile=/dev/null\n")
    # record hostname of minion
    for x in open("/node/etc/hosts").readlines():
        if 'group' in x:
            MINION_IP = x.split()[0]
            open("/node/minion_ip",'w').write(MINION_IP)

def minion_ip():
    global MINION_IP
    if MINION_IP == 'unknown':
        if os.path.exists("/node/minion_ip"):
            MINION_IP = open("/node/minion_ip").read()
            return MINION_IP
        else:
            enable_ssh_access_to_minion()
            if MINION_IP == 'unknown':
                raise RuntimeError("first run enable_ssh_access_to_minion")
            else:
                return MINION_IP
    else:
        return MINION_IP

def run_on_minion(v, *args, **kwds):
    if isinstance(v, str):
        v = "ssh " + minion_ip() + " '%s'"%v
    else:
        v = ['ssh', minion_ip() ] + v
    return run(v, *args, **kwds)

def smc_storage(*args, **kwds):
    run_on_minion(["/usr/libexec/kubernetes/kubelet-plugins/volume/exec/smc~smc-storage/smc-storage"] + list(args), **kwds)

def install_flexvolume_plugin():
    # we always copy it over, which at least upgrades it if necessary.
    shutil.copyfile("/install/smc-storage", "/node/plugin/smc-storage")
    shutil.copymode("/install/smc-storage", "/node/plugin/smc-storage")

def is_plugin_loaded():
    try:
        if int(run_on_minion("zgrep Loaded /var/log/kubelet*|grep smc-storage|wc -l", get_output=True).strip()) > 0:
            return True
        else:
            return False
    except Exception as err:
        print(err)
        return False

def install_zfs():
    try:
        run_on_minion('zpool status')
        print("OK: zfs is installed")
    except:
        print("zfs not installed, so installing it")
        run(['scp', '-r', '/install/gke-zfs', minion_ip()+":"])
        run_on_minion("cd /root/gke-zfs/3.16.0-4-amd64/ && ./install.sh")

def install_bindfs():
    try:
        run_on_minion('which bindfs')
        print("OK: bindfs is installed")
    except:
        print("bindfs not installed, so installing it")
        run_on_minion(["apt-get", "update"])
        run_on_minion(["apt-get", "install", "-y", "bindfs"])

def install_sshfs():
    try:
        run_on_minion('which sshfs')
        print("OK: bindfs is installed")
    except:
        print("bindfs not installed, so installing it")
        run_on_minion(["apt-get", "update"])
        run_on_minion(["apt-get", "install", "-y", "sshfs"])

def install_ssh_keys():
    # Copy the shared secret ssh keys to the minion so that it is able to sshfs
    # mount the storage servers.
    path = '/node/root/.ssh/smc-storage/{POD_NAMESPACE}'.format(POD_NAMESPACE = POD_NAMESPACE)
    if not os.path.exists(path):
        os.makedirs(path)
    for x in ['id-rsa', 'id-rsa.pub']:
        src = os.path.join('/ssh', x); target = os.path.join(path, x.replace('-', '_'))
        shutil.copyfile(src, target)
        os.chmod(target, 0o600)

def restart_kubelet():
    # Sadly I don't know of any other way to properly restart the minion except to completely reboot it.
    # Just restarting the service leaves everything very broken.
    run_on_minion("reboot")

def create_snapshot(pool, name):
    snapshot = "{timestamp}-{name}".format(timestamp=time_to_timestamp(), name=name)
    run_on_minion(['zfs', 'snapshot', "{pool}@{snapshot}".format(pool=pool, snapshot=snapshot)])

def delete_snapshot(pool, snapshot):
    run_on_minion(['zfs', 'destroy', "{pool}@{snaphot}".format(pool=pool, snapshot=snapshot)])

# Lengths of time in minutes.
SNAPSHOT_INTERVALS = {
    'five'    : 5,
    'hourly'  : 60,
    'daily'   : 60*24,
    'weekly'  : 60*24*7,
    'monthly' : 60*24*7*4
}

# How many of each type of snapshot to retain
SNAPSHOT_COUNTS = {
    'five'    : 12*6,   # 6 hours worth of five-minute snapshots
    'hourly'  : 24*7,   # 1 week of hourly snapshots
    'daily'   : 30,     # 1 month of daily snapshots
    'weekly'  : 8,      # 2 months of weekly snapshots
    'monthly' : 6       # 6 months of monthly snapshots
}

TIMESTAMP_FORMAT = "%Y-%m-%d-%H%M%S"      # e.g., 2016-06-27-141131
TIMESTAMP_N = len("2016-06-27-141131")
def time_to_timestamp(tm=None):
    if tm is None:
        tm = time.time()
    return datetime.datetime.fromtimestamp(tm).strftime(TIMESTAMP_FORMAT)

def timestamp_to_time(timestamp):
    return datetime.datetime.strptime(timestamp, TIMESTAMP_FORMAT).timestamp()

def update_snapshots(pool, snapshots):
    """
    Update the rolling ZFS snapshots on the given pool.
    """
    # determine which snapshots we need to make
    now = time.time()
    for name, interval in SNAPSHOT_INTERVALS.items():
        if SNAPSHOT_COUNTS[name] <= 0: # not making any of these
            continue
        # Is there a snapshot with the given name that is within the given
        # interval of now?  If not, make snapshot.
        v = [s for s in snapshots if s.endswith('-'+name)]
        if len(v) > 0:
            newest = v[-1]
            t = timestamp_to_time(newest[:TIMESTAMP_N])
            age_m = (now - t)/60.0   # age in minutes since snapshot
        else:
            age_m = 999999999999  # 'infinite'
        if age_m > interval:
            # make this snapshot
            create_snapshot(pool, name)
        # Are there too many snapshots of the given type?  If so, delete them:
        if len(v) > SNAPSHOT_COUNTS[name]:
            for s in v[ : len(v) - SNAPSHOT_COUNTS[name]]:
                delete_snapshot(pool, s)

def snapshot_info():
    info = {}
    # Get all pools (some may not be in result of snapshot listing below!)
    for pool in run_on_minion(['zfs', 'list', '-r', '-H', '-o', 'name'], get_output=True).split():
        info[pool] = []
    # Get snapshot info for *all* snapshots on all pools
    for snapshot in sorted(run_on_minion(['zfs', 'list', '-r', '-H', '-t', 'snapshot', '-o', 'name'], get_output=True).split()):
        pool, snap = snapshot.split('@')
        info[pool].append(snap)
    return info

def update_all_snapshots():
    """
    Update the rolling ZFS snapshots on all mounted zpool's.
    """
    for pool, snaps in snapshot_info().items():
        update_snapshots(pool, snaps)

# TODO: this entire approach is pointless and broken because when multiple processes
# append to the same file, the result is broken corruption.
def update_zpool_active_log():
    """
    Update log file showing which ZFS filesystems are mounted, which is used by the backup system.
    """
    prefix = "/mnt/smc-storage/{namespace}/".format(namespace=POD_NAMESPACE)
    try:
        v = run_on_minion("zpool status -PL|grep {prefix}".format(prefix=prefix),
                          get_output=True).splitlines()
    except:
        # Nothing to do -- get error if no pools are mounted
        return
    for x in v:
        w = x.split()
        if w:
            path = w[0].strip()           # '/mnt/smc-storage/test/storage0/foo/bar/abc.zfs/00.img'
            path = path[len(prefix):]     # 'storage0/foo/bar/abc.zfs/00.img'
            path = os.path.split(path)[0] # 'storage0/foo/bar/abc.zfs'
            i = path.find('/')
            server = path[:i]
            image = path[i+1:]
            log = "{timestamp} {image}".format(timestamp=time_to_timestamp(), image=image)
            run_on_minion("echo '{log}' >> {prefix}/{server}/log/active.log".format(
                    log=log, prefix=prefix, server=server))

def update_all_lock_files():
    smc_storage("update-all-locks")

def start_storage_daemon():
    print("launching storage daemon")
    install_flexvolume_plugin()
    enable_ssh_access_to_minion()
    install_ssh_keys()
    install_zfs()
    install_bindfs()
    install_sshfs()
    if not is_plugin_loaded():
        restart_kubelet()
    last_snapshot_update = last_lock_update = 0
    while True:
        try:
            update_etc_hosts()
        except Exception as err:
            print("ERROR updating etc hosts -- ", err)
        if time.time() - last_snapshot_update >= 60*2.5:
            try:
                update_all_snapshots()
                last_snapshot_update = time.time()
            except Exception as err:
                print("ERROR updating snapshots -- ", err)
        if time.time() - last_lock_update >= 90:
            try:
                update_all_lock_files()
                last_lock_update = time.time()
            except Exception as err:
                print("ERROR updating locks -- ", err)
        time.sleep(10)

if __name__ == "__main__":
    start_storage_daemon()<commit_msg>k8s: better restarting of kubelet on installing volume plugin<commit_after>#!/usr/bin/env python3
import datetime, json, os, requests, shutil, socket, subprocess, time

HOSTS = '/node/etc/hosts'

# POD_NAMESPACE must be explicitly set in deployment yaml using downward api --
# see https://github.com/kubernetes/kubernetes/blob/release-1.0/docs/user-guide/downward-api.md
POD_NAMESPACE = os.environ.get('POD_NAMESPACE', 'default')

def run(v, shell=False, path='.', get_output=False, env=None, verbose=True):
    t = time.time()
    if isinstance(v, str):
        cmd = v
        shell = True
    else:
        cmd = ' '.join([(x if len(x.split())<=1 else '"%s"'%x) for x in v])
    if path != '.':
        cur = os.path.abspath(os.curdir)
        if verbose:
            print('chdir %s'%path)
        os.chdir(path)
    try:
        if verbose:
            print(cmd)
        if shell:
            kwds = {'shell':True, 'executable':'/bin/bash', 'env':env}
        else:
            kwds = {'env':env}
        if get_output:
            output = subprocess.Popen(v, stdout=subprocess.PIPE, **kwds).stdout.read().decode()
        else:
            if subprocess.call(v, **kwds):
                raise RuntimeError("error running '{cmd}'".format(cmd=cmd))
            output = None
        seconds = time.time() - t
        if verbose:
            print("TOTAL TIME: {seconds} seconds -- to run '{cmd}'".format(seconds=seconds, cmd=cmd))
        return output
    finally:
        if path != '.':
            os.chdir(cur)

def get_service(service):
    """
    Get in json format the kubernetes information about the given service.
    """
    if not os.environ['KUBERNETES_SERVICE_HOST']:
        print('KUBERNETES_SERVICE_HOST environment variable not set')
        return None
    URL = "https://{KUBERNETES_SERVICE_HOST}:{KUBERNETES_SERVICE_PORT}/api/v1/namespaces/{POD_NAMESPACE}/endpoints/{service}"
    URL = URL.format(KUBERNETES_SERVICE_HOST=os.environ['KUBERNETES_SERVICE_HOST'],
                     KUBERNETES_SERVICE_PORT=os.environ['KUBERNETES_SERVICE_PORT'],
                     POD_NAMESPACE=POD_NAMESPACE,
                     service=service)
    token = open('/var/run/secrets/kubernetes.io/serviceaccount/token').read()
    headers={'Authorization':'Bearer {token}'.format(token=token)}
    print("Getting k8s information about '{service}' from '{URL}'".format(service=service, URL=URL))
    x = requests.get(URL, headers=headers, verify='/var/run/secrets/kubernetes.io/serviceaccount/ca.crt').json()
    print("Got {x}".format(x=x))
    return x

def update_etc_hosts():
    print('udpate_etc_hosts')
    try:
        v = get_service('storage-projects')
    except Exception as err:
        # Expected to happen when node is starting up, etc. - we'll retry later soon!
        print("Failed getting storage service info", err)
        return
    if v.get('status', None) == 'Failure':
        return
    try:
        if 'addresses' not in v['subsets'][0]:
            return   # nothing to do; no known addresses
        namespace = v['metadata']['namespace']
        hosts = ["{ip}    {namespace}-{name}".format(ip=x['ip'], namespace=namespace,
                              name=x['targetRef']['name'].split('-')[0]) for x in v['subsets'][0]['addresses']]
        start = "# start smc-storage dns - namespace="+namespace+"\n\n"
        end = "# end smc-storage dns - namespace="+namespace+"\n\n"
        block = '\n'.join([start] + hosts + [end])
        current = open(HOSTS).read()
        if block in current:
            return
        i = current.find(start)
        j = current.find(end)
        if i == -1 or j == -1:
            new = current + '\n' + block
        else:
            new = current[:i] + block + current[j+len(end):]
        open(HOSTS,'w').write(new)
    except Exception as err:
        print("Problem in update_etc_hosts", err)

MINION_IP = 'unknown'
def enable_ssh_access_to_minion():
    global MINION_IP
    # create our own local ssh key
    if os.path.exists('/root/.ssh'):
        shutil.rmtree('/root/.ssh')
    run(['ssh-keygen', '-b', '2048', '-N', '', '-f', '/root/.ssh/id_rsa'])
    # make root user of minion allow login using this (and only this) key.
    run('cat /root/.ssh/id_rsa.pub >> /node/root/.ssh/authorized_keys')
    open("/root/.ssh/config",'w').write("StrictHostKeyChecking no\nUserKnownHostsFile=/dev/null\n")
    # record hostname of minion
    for x in open("/node/etc/hosts").readlines():
        if 'group' in x:
            MINION_IP = x.split()[0]
            open("/node/minion_ip",'w').write(MINION_IP)

def minion_ip():
    global MINION_IP
    if MINION_IP == 'unknown':
        if os.path.exists("/node/minion_ip"):
            MINION_IP = open("/node/minion_ip").read()
            return MINION_IP
        else:
            enable_ssh_access_to_minion()
            if MINION_IP == 'unknown':
                raise RuntimeError("first run enable_ssh_access_to_minion")
            else:
                return MINION_IP
    else:
        return MINION_IP

def run_on_minion(v, *args, **kwds):
    if isinstance(v, str):
        v = "ssh " + minion_ip() + " '%s'"%v
    else:
        v = ['ssh', minion_ip() ] + v
    return run(v, *args, **kwds)

def smc_storage(*args, **kwds):
    run_on_minion(["/usr/libexec/kubernetes/kubelet-plugins/volume/exec/smc~smc-storage/smc-storage"] + list(args), **kwds)

def install_flexvolume_plugin():
    # we always copy it over, which at least upgrades it if necessary.
    shutil.copyfile("/install/smc-storage", "/node/plugin/smc-storage")
    shutil.copymode("/install/smc-storage", "/node/plugin/smc-storage")

def is_plugin_loaded():
    try:
        if int(run_on_minion("zgrep Loaded /var/log/kubelet*|grep smc-storage|wc -l", get_output=True).strip()) > 0:
            return True
        else:
            return False
    except Exception as err:
        print(err)
        return False

def install_zfs():
    try:
        run_on_minion('zpool status')
        print("OK: zfs is installed")
    except:
        print("zfs not installed, so installing it")
        run(['scp', '-r', '/install/gke-zfs', minion_ip()+":"])
        run_on_minion("cd /root/gke-zfs/3.16.0-4-amd64/ && ./install.sh")

def install_bindfs():
    try:
        run_on_minion('which bindfs')
        print("OK: bindfs is installed")
    except:
        print("bindfs not installed, so installing it")
        run_on_minion(["apt-get", "update"])
        run_on_minion(["apt-get", "install", "-y", "bindfs"])

def install_sshfs():
    try:
        run_on_minion('which sshfs')
        print("OK: bindfs is installed")
    except:
        print("bindfs not installed, so installing it")
        run_on_minion(["apt-get", "update"])
        run_on_minion(["apt-get", "install", "-y", "sshfs"])

def install_ssh_keys():
    # Copy the shared secret ssh keys to the minion so that it is able to sshfs
    # mount the storage servers.
    path = '/node/root/.ssh/smc-storage/{POD_NAMESPACE}'.format(POD_NAMESPACE = POD_NAMESPACE)
    if not os.path.exists(path):
        os.makedirs(path)
    for x in ['id-rsa', 'id-rsa.pub']:
        src = os.path.join('/ssh', x); target = os.path.join(path, x.replace('-', '_'))
        shutil.copyfile(src, target)
        os.chmod(target, 0o600)

def restart_kubelet():
    run_on_minion("kill `pidof /usr/local/bin/kubelet`")

def create_snapshot(pool, name):
    snapshot = "{timestamp}-{name}".format(timestamp=time_to_timestamp(), name=name)
    run_on_minion(['zfs', 'snapshot', "{pool}@{snapshot}".format(pool=pool, snapshot=snapshot)])

def delete_snapshot(pool, snapshot):
    run_on_minion(['zfs', 'destroy', "{pool}@{snaphot}".format(pool=pool, snapshot=snapshot)])

# Lengths of time in minutes.
SNAPSHOT_INTERVALS = {
    'five'    : 5,
    'hourly'  : 60,
    'daily'   : 60*24,
    'weekly'  : 60*24*7,
    'monthly' : 60*24*7*4
}

# How many of each type of snapshot to retain
SNAPSHOT_COUNTS = {
    'five'    : 12*6,   # 6 hours worth of five-minute snapshots
    'hourly'  : 24*7,   # 1 week of hourly snapshots
    'daily'   : 30,     # 1 month of daily snapshots
    'weekly'  : 8,      # 2 months of weekly snapshots
    'monthly' : 6       # 6 months of monthly snapshots
}

TIMESTAMP_FORMAT = "%Y-%m-%d-%H%M%S"      # e.g., 2016-06-27-141131
TIMESTAMP_N = len("2016-06-27-141131")
def time_to_timestamp(tm=None):
    if tm is None:
        tm = time.time()
    return datetime.datetime.fromtimestamp(tm).strftime(TIMESTAMP_FORMAT)

def timestamp_to_time(timestamp):
    return datetime.datetime.strptime(timestamp, TIMESTAMP_FORMAT).timestamp()

def update_snapshots(pool, snapshots):
    """
    Update the rolling ZFS snapshots on the given pool.
    """
    # determine which snapshots we need to make
    now = time.time()
    for name, interval in SNAPSHOT_INTERVALS.items():
        if SNAPSHOT_COUNTS[name] <= 0: # not making any of these
            continue
        # Is there a snapshot with the given name that is within the given
        # interval of now?  If not, make snapshot.
        v = [s for s in snapshots if s.endswith('-'+name)]
        if len(v) > 0:
            newest = v[-1]
            t = timestamp_to_time(newest[:TIMESTAMP_N])
            age_m = (now - t)/60.0   # age in minutes since snapshot
        else:
            age_m = 999999999999  # 'infinite'
        if age_m > interval:
            # make this snapshot
            create_snapshot(pool, name)
        # Are there too many snapshots of the given type?  If so, delete them:
        if len(v) > SNAPSHOT_COUNTS[name]:
            for s in v[ : len(v) - SNAPSHOT_COUNTS[name]]:
                delete_snapshot(pool, s)

def snapshot_info():
    info = {}
    # Get all pools (some may not be in result of snapshot listing below!)
    for pool in run_on_minion(['zfs', 'list', '-r', '-H', '-o', 'name'], get_output=True).split():
        info[pool] = []
    # Get snapshot info for *all* snapshots on all pools
    for snapshot in sorted(run_on_minion(['zfs', 'list', '-r', '-H', '-t', 'snapshot', '-o', 'name'], get_output=True).split()):
        pool, snap = snapshot.split('@')
        info[pool].append(snap)
    return info

def update_all_snapshots():
    """
    Update the rolling ZFS snapshots on all mounted zpool's.
    """
    for pool, snaps in snapshot_info().items():
        update_snapshots(pool, snaps)

# TODO: this entire approach is pointless and broken because when multiple processes
# append to the same file, the result is broken corruption.
def update_zpool_active_log():
    """
    Update log file showing which ZFS filesystems are mounted, which is used by the backup system.
    """
    prefix = "/mnt/smc-storage/{namespace}/".format(namespace=POD_NAMESPACE)
    try:
        v = run_on_minion("zpool status -PL|grep {prefix}".format(prefix=prefix),
                          get_output=True).splitlines()
    except:
        # Nothing to do -- get error if no pools are mounted
        return
    for x in v:
        w = x.split()
        if w:
            path = w[0].strip()           # '/mnt/smc-storage/test/storage0/foo/bar/abc.zfs/00.img'
            path = path[len(prefix):]     # 'storage0/foo/bar/abc.zfs/00.img'
            path = os.path.split(path)[0] # 'storage0/foo/bar/abc.zfs'
            i = path.find('/')
            server = path[:i]
            image = path[i+1:]
            log = "{timestamp} {image}".format(timestamp=time_to_timestamp(), image=image)
            run_on_minion("echo '{log}' >> {prefix}/{server}/log/active.log".format(
                    log=log, prefix=prefix, server=server))

def update_all_lock_files():
    smc_storage("update-all-locks")

def start_storage_daemon():
    print("launching storage daemon")
    install_flexvolume_plugin()
    enable_ssh_access_to_minion()
    install_ssh_keys()
    install_zfs()
    install_bindfs()
    install_sshfs()
    if not is_plugin_loaded():
        restart_kubelet()
    last_snapshot_update = last_lock_update = 0
    while True:
        try:
            update_etc_hosts()
        except Exception as err:
            print("ERROR updating etc hosts -- ", err)
        if time.time() - last_snapshot_update >= 60*2.5:
            try:
                update_all_snapshots()
                last_snapshot_update = time.time()
            except Exception as err:
                print("ERROR updating snapshots -- ", err)
        if time.time() - last_lock_update >= 90:
            try:
                update_all_lock_files()
                last_lock_update = time.time()
            except Exception as err:
                print("ERROR updating locks -- ", err)
        time.sleep(10)

if __name__ == "__main__":
    start_storage_daemon()<|endoftext|>