<commit_before>#! /usr/bin/python

from __future__ import print_function, division
import time
import logging
import subprocess
import os
import smtplib
from email.mime.text import MIMEText
from abc import ABCMeta, abstractproperty
import xml.etree.ElementTree as ET # for XML parsing
import signal

# TODO: Send daily heartbeat

# TODO: check git is updating (maybe just check age of .git??)

"""
***********************************
********* DESCRIPTION *************

    Script for monitoring disk space, multiple files and multiple processes.
    If errors are found then an email is sent.
    Configuration is stored in a babysitter_config.xml file (see below)

***********************************
********* REQUIREMENTS ************

    ---------------------------------
    ADD A babysitter_config.xml FILE
    USING THE FOLLOWING FORMAT:

    Add as many <process> or <file>
    elements as you want to check.
    --------------------------------

        <config>

            <smtp_server></smtp_server>
            <email_from></email_from>
            <email_to></email_to>
            <username></username>
            <password></password>

            <disk_space_threshold>20</disk_space_threshold> <!-- MBytes -->

            <file>
                <location></location>
                <timeout>300</timeout> <!-- Seconds -->
            </file>

            <process>
                <name>iam_logger.py</name>
                <restart_command>nohup ./iam_logger.py</restart_command>
            </process>

            <process>
                <name>ntpd</name>
                <restart_command>sudo service ntp restart</restart_command>
            </process>

        </config>

    ----------------------------------
    SETUP SUDO FOR service restart ntp
    ----------------------------------

        This Python script needs to be able to run 'service restart ntp'
        without requiring a password.  Follow these steps:

        1. run 'sudo visudo'
        2. add the following line:
           'USER   ALL=NOPASSWD: /usr/sbin/service ntp *'
           (replace USER with your unix username!)

"""

class Checker:
    """Abstract base class (ABC) for classes which check on the state of
    a particular part of the system. """

    __metaclass__ = ABCMeta

    # Define some constants for tracking state
    FAIL = 0
    OK = 1

    def __init__(self, name):
        self.name = name
        self.last_state = self.state

    @abstractproperty
    def state(self):
        pass

    @property
    def state_as_str(self):
        return ['FAIL', 'OK'][self.state]

    @property
    def just_changed_state(self):
        state = self.state # cache to avoid this changing under us
        if state == self.last_state:
            return False
        elif state == Checker.FAIL:
            logger.warning('state change to FAIL: {}'.format(self))
        elif state == Checker.OK:
            logger.info('state change: {}'.format(self))

        self.last_state = state
        return True

    def __str__(self):
        return '{} = {}'.format(self.name.rpartition('/')[2], # remove path
                                self.state_as_str)

class Process(Checker):
    """Class for monitoring a unix process.

    Attributes:
        name (str): the process name as it appears in `ps -A`
        restart_string (str): the command used to restart this process

    """

    def __init__(self, name):
        """
        Args:
            name (str): the process name as it appears in `ps -A`
        """
        self.restart_string = None
        super(Process, self).__init__(name)

    @property
    def pid(self):
        pid_string = subprocess.check_output(['pidof', '-x', self.name])
        return pid_string.strip()

    def restart(self):
        if self.restart_string is None:
            logger.info("No restart string for {}".format(self.name))
            return

        if self.state == Checker.OK:
            return

        logger.info("Attempting to restart {}".format(self.name))
        try:
            subprocess.Popen(self.restart_command.split())
        except Exception:
            logger.exception("Failed to restart. {}".format(self))
        else:
            logger.info("Successfully restarted. {}".format(self) )

    @property
    def state(self):
        try:
            self.pid
        except subprocess.CalledProcessError:
            return Checker.FAIL
        else:
            return Checker.OK

class File(Checker):
    def __init__(self, name, timeout=120):
        """File constructor

        Args:
            name (str) : including full path
            timeout (int or str) : time in seconds after which this file is
                considered overdue.
        """
        self.timeout = int(timeout)
        super(File, self).__init__(name)

    @property
    def state(self):
        return self.seconds_since_modified < self.timeout

    @property
    def seconds_since_modified(self):
        return time.time() - self.last_modified

    @property
    def last_modified(self):
        return os.path.getmtime(self.name)

    def __str__(self):
        msg = super(File, self).__str__()
        msg += ", last modified {:.1f}s ago.".format(self.seconds_since_modified)
        return msg

class DiskSpaceRemaining(Checker):

    def __init__(self, threshold):
        """
        Args:
            threshold (int or str): number of MBytes of free space below which
                state will change to FAIL.
        """
        self.threshold = int(threshold)
        super(DiskSpaceRemaining, self).__init__('disk space')

    @property
    def state(self):
        return self.available_space > self.threshold

    @property
    def available_space(self):
        """Returns available disk space in MBytes."""
        # From http://stackoverflow.com/a/787832/732596
        s = os.statvfs('/')
        return (s.f_bavail * s.f_frsize) / 1024**2

    def __str__(self):
        msg = super(DiskSpaceRemaining, self).__str__()
        msg += ", remaining={:.0f} MB.".format(self.available_space)
        return msg

class Manager(object):
    """Manages multiple Checker objects."""

    def __init__(self):
        self._checkers = []

    def append(self, checker):
        self._checkers.append(checker)
        logger.info('Added {} to Manager: {}'.format(checker.__class__.__name__,
                                                     self._checkers[-1]))

    def run(self):
        msg = "IAM logger babysitter running.\n{}".format(self)
        self.send_email(body=msg, subject="babysitter.py running")

        while True:
            msg = ""
            for checker in self._checkers:
                if checker.just_changed_state:
                    msg += "STATE CHANGED:\n"
                    msg += str(checker) + "\n"
                    if isinstance(checker, Process):
                        msg += "Attempting to restart...\n"
                        checker.restart()
                        time.sleep(5)
                        msg += str(checker) + "\n"

            if msg != "":
                msg += "CURRENT STATE OF ALL CHECKERS:\n" + str(self)
                self.send_email(body=msg, subject="iam_logger errors.")

            time.sleep(10)

    def load_config(self, config_file):
        config_tree = ET.parse(config_file)

        self.SMTP_SERVER = config_tree.findtext("smtp_server")
        self.EMAIL_FROM  = config_tree.findtext("email_from")
        self.EMAIL_TO    = config_tree.findtext("email_to")
        self.USERNAME    = config_tree.findtext("username")
        self.PASSWORD    = config_tree.findtext("password")

        logger.debug('\nSMTP_SERVER={}\nEMAIL_FROM={}\nEMAIL_TO={}'
                     .format(self.SMTP_SERVER, self.EMAIL_FROM, self.EMAIL_TO))

        # Disk space checker
        disk_space_threshold = config_tree.findtext("disk_space_threshold")
        if disk_space_threshold is not None:
            self.append(DiskSpaceRemaining(disk_space_threshold))

        # Load files
        files_etree = config_tree.findall("file")
        for f in files_etree:
            self.append(File(f.findtext('location'),
                             int(f.findtext('timeout'))))

        # Load processes
        processes_etree = config_tree.findall("process")
        for process in processes_etree:
            p = Process(process.findtext('name'))
            p.restart_command = process.findtext('restart_command')
            self.append(p)

    def send_email(self, body, subject):
        hostname = os.uname()[1]
        me = hostname + '<' + self.EMAIL_FROM + '>'
        body += '\nUnixtime = ' + str(time.time()) + '\n'
        msg = MIMEText(body)
        msg['Subject'] = subject
        msg['From'] =  me
        msg['To'] = self.EMAIL_TO

        logger.debug('sending message: \n{}'.format(msg.as_string()))

        retry = True
        while retry is True:
            try:
                logger.debug("SMPT_SSL")
                s = smtplib.SMTP_SSL(self.SMTP_SERVER)
                logger.debug("logging in")
                s.login(self.USERNAME, self.PASSWORD) # TODO take these from config!

                logger.debug("sendmail")
                s.sendmail(me, [self.EMAIL_TO], msg.as_string())
                logger.debug("quit")
                s.quit()
            except (smtplib.SMTPServerDisconnected, smtplib.SMTPConnectError):
                logger.exception("")
                time.sleep(2)
            except smtplib.SMTPAuthenticationError:
                error_msg = "SMTP authentication error. Please check username and password in config file."
                print(error_msg)
                logger.exception(error_msg)
                raise
            else:
                logger.info("Successfully sent message")
                retry = False

    def __str__(self):
        msg = ""
        for checker in self._checkers:
            msg += '{}\n'.format(checker)
        return msg

def _init_logger():
    global logger

    # create logger
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)

    # create console handler for stderr
    ch_stderr = logging.StreamHandler()
    ch_stderr.setLevel(logging.INFO)
    stderr_formatter = logging.Formatter('%(asctime)s %(levelname)s: '
                        '%(message)s', datefmt='%d/%m/%y %H:%M:%S')
    ch_stderr.setFormatter(stderr_formatter)
    logger.addHandler(ch_stderr)

    # create file handler for babysitter.log
    fh = logging.FileHandler('babysitter.log')
    fh.setLevel(logging.DEBUG)
    fh_formatter = logging.Formatter('%(asctime)s level=%(levelname)s: '
                        'function=%(funcName)s, thread=%(threadName)s'
                        '\n   %(message)s')
    fh.setFormatter(fh_formatter)
    logger.addHandler(fh)

def _shutdown():
    logger.info("Shutting down.")
    logging.shutdown()

def _signal_handler(signal_number, frame):
    raise KeyboardInterrupt()

def main():

    _init_logger()
    logger.debug('MAIN: babysitter.py starting up. Unixtime = {:.0f}'
                  .format(time.time()))

    # register SIGINT and SIGTERM handler
    logger.info("MAIN: setting signal handlers")
    signal.signal(signal.SIGINT,  _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    # Wrap this in try... except so we can send any unexpected exceptions
    # to logging
    try:
        manager = Manager()
        manager.load_config("babysitter_config.xml")
        manager.run()
    except KeyboardInterrupt:
        _shutdown()
    except Exception:
        logger.exception("")
        _shutdown()
        raise

if __name__ == "__main__":
    main()<commit_msg>fix bug with restart_command<commit_after>#! /usr/bin/python

from __future__ import print_function, division
import time
import logging
import subprocess
import os
import smtplib
from email.mime.text import MIMEText
from abc import ABCMeta, abstractproperty
import xml.etree.ElementTree as ET # for XML parsing
import signal

# TODO: Send daily heartbeat

# TODO: check git is updating (maybe just check age of .git??)

"""
***********************************
********* DESCRIPTION *************

    Script for monitoring disk space, multiple files and multiple processes.
    If errors are found then an email is sent.
    Configuration is stored in a babysitter_config.xml file (see below)

***********************************
********* REQUIREMENTS ************

    ---------------------------------
    ADD A babysitter_config.xml FILE
    USING THE FOLLOWING FORMAT:

    Add as many <process> or <file>
    elements as you want to check.
    --------------------------------

        <config>

            <smtp_server></smtp_server>
            <email_from></email_from>
            <email_to></email_to>
            <username></username>
            <password></password>

            <disk_space_threshold>20</disk_space_threshold> <!-- MBytes -->

            <file>
                <location></location>
                <timeout>300</timeout> <!-- Seconds -->
            </file>

            <process>
                <name>iam_logger.py</name>
                <restart_command>nohup ./iam_logger.py</restart_command>
            </process>

            <process>
                <name>ntpd</name>
                <restart_command>sudo service ntp restart</restart_command>
            </process>

        </config>

    ----------------------------------
    SETUP SUDO FOR service restart ntp
    ----------------------------------

        This Python script needs to be able to run 'service restart ntp'
        without requiring a password.  Follow these steps:

        1. run 'sudo visudo'
        2. add the following line:
           'USER   ALL=NOPASSWD: /usr/sbin/service ntp *'
           (replace USER with your unix username!)

"""

class Checker:
    """Abstract base class (ABC) for classes which check on the state of
    a particular part of the system. """

    __metaclass__ = ABCMeta

    # Define some constants for tracking state
    FAIL = 0
    OK = 1

    def __init__(self, name):
        self.name = name
        self.last_state = self.state

    @abstractproperty
    def state(self):
        pass

    @property
    def state_as_str(self):
        return ['FAIL', 'OK'][self.state]

    @property
    def just_changed_state(self):
        state = self.state # cache to avoid this changing under us
        if state == self.last_state:
            return False
        elif state == Checker.FAIL:
            logger.warning('state change to FAIL: {}'.format(self))
        elif state == Checker.OK:
            logger.info('state change: {}'.format(self))

        self.last_state = state
        return True

    def __str__(self):
        return '{} = {}'.format(self.name.rpartition('/')[2], # remove path
                                self.state_as_str)

class Process(Checker):
    """Class for monitoring a unix process.

    Attributes:
        name (str): the process name as it appears in `ps -A`
        restart_command (str): the command used to restart this process

    """

    def __init__(self, name):
        """
        Args:
            name (str): the process name as it appears in `ps -A`
        """
        self.restart_command = None
        super(Process, self).__init__(name)

    @property
    def pid(self):
        pid_string = subprocess.check_output(['pidof', '-x', self.name])
        return pid_string.strip()

    def restart(self):
        if self.restart_command is None:
            logger.info("No restart string for {}".format(self.name))
            return

        if self.state == Checker.OK:
            return

        logger.info("Attempting to restart {}".format(self.name))
        try:
            subprocess.Popen(self.restart_command.split())
        except Exception:
            logger.exception("Failed to restart. {}".format(self))
        else:
            logger.info("Successfully restarted. {}".format(self) )

    @property
    def state(self):
        try:
            self.pid
        except subprocess.CalledProcessError:
            return Checker.FAIL
        else:
            return Checker.OK

class File(Checker):
    def __init__(self, name, timeout=120):
        """File constructor

        Args:
            name (str) : including full path
            timeout (int or str) : time in seconds after which this file is
                considered overdue.
        """
        self.timeout = int(timeout)
        super(File, self).__init__(name)

    @property
    def state(self):
        return self.seconds_since_modified < self.timeout

    @property
    def seconds_since_modified(self):
        return time.time() - self.last_modified

    @property
    def last_modified(self):
        return os.path.getmtime(self.name)

    def __str__(self):
        msg = super(File, self).__str__()
        msg += ", last modified {:.1f}s ago.".format(self.seconds_since_modified)
        return msg

class DiskSpaceRemaining(Checker):

    def __init__(self, threshold):
        """
        Args:
            threshold (int or str): number of MBytes of free space below which
                state will change to FAIL.
        """
        self.threshold = int(threshold)
        super(DiskSpaceRemaining, self).__init__('disk space')

    @property
    def state(self):
        return self.available_space > self.threshold

    @property
    def available_space(self):
        """Returns available disk space in MBytes."""
        # From http://stackoverflow.com/a/787832/732596
        s = os.statvfs('/')
        return (s.f_bavail * s.f_frsize) / 1024**2

    def __str__(self):
        msg = super(DiskSpaceRemaining, self).__str__()
        msg += ", remaining={:.0f} MB.".format(self.available_space)
        return msg

class Manager(object):
    """Manages multiple Checker objects."""

    def __init__(self):
        self._checkers = []

    def append(self, checker):
        self._checkers.append(checker)
        logger.info('Added {} to Manager: {}'.format(checker.__class__.__name__,
                                                     self._checkers[-1]))

    def run(self):
        msg = "IAM logger babysitter running.\n{}".format(self)
        self.send_email(body=msg, subject="babysitter.py running")

        while True:
            msg = ""
            for checker in self._checkers:
                if checker.just_changed_state:
                    msg += "STATE CHANGED:\n"
                    msg += str(checker) + "\n"
                    if isinstance(checker, Process):
                        msg += "Attempting to restart...\n"
                        checker.restart()
                        time.sleep(5)
                        msg += str(checker) + "\n"

            if msg != "":
                msg += "CURRENT STATE OF ALL CHECKERS:\n" + str(self)
                self.send_email(body=msg, subject="iam_logger errors.")

            time.sleep(10)

    def load_config(self, config_file):
        config_tree = ET.parse(config_file)

        self.SMTP_SERVER = config_tree.findtext("smtp_server")
        self.EMAIL_FROM  = config_tree.findtext("email_from")
        self.EMAIL_TO    = config_tree.findtext("email_to")
        self.USERNAME    = config_tree.findtext("username")
        self.PASSWORD    = config_tree.findtext("password")

        logger.debug('\nSMTP_SERVER={}\nEMAIL_FROM={}\nEMAIL_TO={}'
                     .format(self.SMTP_SERVER, self.EMAIL_FROM, self.EMAIL_TO))

        # Disk space checker
        disk_space_threshold = config_tree.findtext("disk_space_threshold")
        if disk_space_threshold is not None:
            self.append(DiskSpaceRemaining(disk_space_threshold))

        # Load files
        files_etree = config_tree.findall("file")
        for f in files_etree:
            self.append(File(f.findtext('location'),
                             int(f.findtext('timeout'))))

        # Load processes
        processes_etree = config_tree.findall("process")
        for process in processes_etree:
            p = Process(process.findtext('name'))
            p.restart_command = process.findtext('restart_command')
            self.append(p)

    def send_email(self, body, subject):
        hostname = os.uname()[1]
        me = hostname + '<' + self.EMAIL_FROM + '>'
        body += '\nUnixtime = ' + str(time.time()) + '\n'
        msg = MIMEText(body)
        msg['Subject'] = subject
        msg['From'] =  me
        msg['To'] = self.EMAIL_TO

        logger.debug('sending message: \n{}'.format(msg.as_string()))

        retry = True
        while retry is True:
            try:
                logger.debug("SMPT_SSL")
                s = smtplib.SMTP_SSL(self.SMTP_SERVER)
                logger.debug("logging in")
                s.login(self.USERNAME, self.PASSWORD) # TODO take these from config!

                logger.debug("sendmail")
                s.sendmail(me, [self.EMAIL_TO], msg.as_string())
                logger.debug("quit")
                s.quit()
            except (smtplib.SMTPServerDisconnected, smtplib.SMTPConnectError):
                logger.exception("")
                time.sleep(2)
            except smtplib.SMTPAuthenticationError:
                error_msg = "SMTP authentication error. Please check username and password in config file."
                print(error_msg)
                logger.exception(error_msg)
                raise
            else:
                logger.info("Successfully sent message")
                retry = False

    def __str__(self):
        msg = ""
        for checker in self._checkers:
            msg += '{}\n'.format(checker)
        return msg

def _init_logger():
    global logger

    # create logger
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)

    # create console handler for stderr
    ch_stderr = logging.StreamHandler()
    ch_stderr.setLevel(logging.INFO)
    stderr_formatter = logging.Formatter('%(asctime)s %(levelname)s: '
                        '%(message)s', datefmt='%d/%m/%y %H:%M:%S')
    ch_stderr.setFormatter(stderr_formatter)
    logger.addHandler(ch_stderr)

    # create file handler for babysitter.log
    fh = logging.FileHandler('babysitter.log')
    fh.setLevel(logging.DEBUG)
    fh_formatter = logging.Formatter('%(asctime)s level=%(levelname)s: '
                        'function=%(funcName)s, thread=%(threadName)s'
                        '\n   %(message)s')
    fh.setFormatter(fh_formatter)
    logger.addHandler(fh)

def _shutdown():
    logger.info("Shutting down.")
    logging.shutdown()

def _signal_handler(signal_number, frame):
    raise KeyboardInterrupt()

def main():

    _init_logger()
    logger.debug('MAIN: babysitter.py starting up. Unixtime = {:.0f}'
                  .format(time.time()))

    # register SIGINT and SIGTERM handler
    logger.info("MAIN: setting signal handlers")
    signal.signal(signal.SIGINT,  _signal_handler)
    signal.signal(signal.SIGTERM, _signal_handler)

    # Wrap this in try... except so we can send any unexpected exceptions
    # to logging
    try:
        manager = Manager()
        manager.load_config("babysitter_config.xml")
        manager.run()
    except KeyboardInterrupt:
        _shutdown()
    except Exception:
        logger.exception("")
        _shutdown()
        raise

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


<commit_before>##########################################################################
#          (C) Vrije Universiteit, Amsterdam (the Netherlands)            #
#                                                                         #
# This file is part of AmCAT - The Amsterdam Content Analysis Toolkit     #
#                                                                         #
# AmCAT is free software: you can redistribute it and/or modify it under  #
# the terms of the GNU Affero General Public License as published by the  #
# Free Software Foundation, either version 3 of the License, or (at your  #
# option) any later version.                                              #
#                                                                         #
# AmCAT is distributed in the hope that it will be useful, but WITHOUT    #
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or   #
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public     #
# License for more details.                                               #
#                                                                         #
# You should have received a copy of the GNU Affero General Public        #
# License along with AmCAT.  If not, see <http://www.gnu.org/licenses/>.  #
###########################################################################
from django import forms
from django.core.exceptions import ValidationError
from django.shortcuts import render
from django.template import Context
from django.template.loader import get_template
from amcat.models import ArticleSet

from amcat.scripts.query import QueryAction, QueryActionForm
from amcat.tools.keywordsearch import SelectionSearch

OK_TEMPLATE = get_template("query/ok.html")

class SaveAsSetForm(QueryActionForm):
    name = forms.CharField(required=True)
    provenance = forms.CharField(required=False)

    def clean_name(self):
        new_name = self.cleaned_data["name"]

        all_articlesets = self.project.articlesets_set.all()
        if all_articlesets.filter(name__iexact=new_name).exists():
            error_msg = "An articleset called {} already exists."
            raise ValidationError(error_msg.format(self.cleaned_data["name"]))

        return new_name

class SaveAsSetAction(QueryAction):
    form_class = SaveAsSetForm
    output_types = (("text/html", "Result"),)

    def run(self, form):
        name = form.cleaned_data["name"]
        provenance = form.cleaned_data["provenance"]
        aset = ArticleSet.objects.create(name=name, provenance=provenance, project=self.project)
        self.monitor.update(10, "Executing query..")
        article_ids = list(SelectionSearch(form).get_article_ids())
        self.monitor.update(60, "Saving to set..")
        aset.add_articles(article_ids)

        return OK_TEMPLATE.render(Context({
            "project": self.project,
            "aset": aset,
            "len": len(article_ids)
        }))

<commit_msg>Added 'project' argument in SaveAsSetAction<commit_after>##########################################################################
#          (C) Vrije Universiteit, Amsterdam (the Netherlands)            #
#                                                                         #
# This file is part of AmCAT - The Amsterdam Content Analysis Toolkit     #
#                                                                         #
# AmCAT is free software: you can redistribute it and/or modify it under  #
# the terms of the GNU Affero General Public License as published by the  #
# Free Software Foundation, either version 3 of the License, or (at your  #
# option) any later version.                                              #
#                                                                         #
# AmCAT is distributed in the hope that it will be useful, but WITHOUT    #
# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or   #
# FITNESS FOR A PARTICULAR PURPOSE. See the GNU Affero General Public     #
# License for more details.                                               #
#                                                                         #
# You should have received a copy of the GNU Affero General Public        #
# License along with AmCAT.  If not, see <http://www.gnu.org/licenses/>.  #
###########################################################################
from django import forms
from django.core.exceptions import ValidationError
from django.shortcuts import render
from django.template import Context
from django.template.loader import get_template
from amcat.models import ArticleSet
from amcat.models import Project
from amcat.models import Role
from amcat.models.authorisation import ROLE_PROJECT_WRITER

from amcat.scripts.query import QueryAction, QueryActionForm
from amcat.tools.keywordsearch import SelectionSearch

OK_TEMPLATE = get_template("query/ok.html")

class SaveAsSetForm(QueryActionForm):
    project = forms.ModelChoiceField(queryset=Project.objects.none(), empty_label=None)
    name = forms.CharField(required=True)
    #provenance = forms.CharField(required=False)

    def __init__(self, *args, **kwargs):
        super(SaveAsSetForm, self).__init__(*args, **kwargs)
        rw_role = Role.objects.get(id=ROLE_PROJECT_WRITER)
        self.fields["project"].queryset = self.user.userprofile.get_projects(rw_role)
        self.fields["project"].initial = self.project

    def clean_name(self):
        new_name = self.cleaned_data["name"]
        project = self.cleaned_data["project"]

        all_articlesets = project.articlesets_set.all()
        if all_articlesets.filter(name__iexact=new_name).exists():
            error_msg = "An articleset called {0!r} already exists in project {1.name!r}."
            raise ValidationError(error_msg.format(new_name, project))

        return new_name

class SaveAsSetAction(QueryAction):
    form_class = SaveAsSetForm
    output_types = (("text/html", "Result"),)

    def run(self, form):
        name = form.cleaned_data["name"]
        #provenance = form.cleaned_data["provenance"]
        project = form.cleaned_data["project"]
        aset = ArticleSet.objects.create(name=name, project=project)
        self.monitor.update(10, "Executing query..")
        article_ids = list(SelectionSearch(form).get_article_ids())
        self.monitor.update(60, "Saving to set..")
        aset.add_articles(article_ids)

        return OK_TEMPLATE.render(Context({
            "project": project,
            "aset": aset,
            "len": len(article_ids)
        }))

<|endoftext|>


<commit_before>### This script calculates various statistics on all tiles in input folders and saves them to a txt.
### Users can input as many folders as they want for calculating statistics on each tile

import multiprocessing
import tile_statistics
import subprocess
import os
import glob
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

# The column names for the tile summary statistics.
# If the statistics calculations are changed in tile_statistics.py, the list here needs to be changed, too.
headers = ['tile_id', 'tile_type', 'tile_name', 'pixel_count', 'mean', 'median', 'percentile10', 'percentile25',
           'percentile75', 'percentile90', 'min', 'max', 'sum']
header_no_brackets = ', '.join(headers)

tile_stats = '{}_{}'.format(uu.date_today, cn.tile_stats_pattern)

# Creates the output text file with the column names
with open(cn.tile_stats_pattern, 'w+') as f:
    f.write(header_no_brackets  +'\r\n')
f.close()

# # Creates list of tiles to iterate through, for testing
# download_tile_list = ['00N_110E', '30N_080W', '40N_050E', '50N_100E', '80N_020E'] # test tiles
# print download_tile_list
#
# # For copying individual tiles to spot machine for testing
# for tile in download_tile_list:
#     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.mangrove_biomass_2000_dir, cn.pattern_mangrove_biomass_2000, tile), '.')
#     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.natrl_forest_biomass_2000_dir, tile, cn.pattern_natrl_forest_biomass_2000), '.')
#     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.pixel_area_dir, cn.pattern_pixel_area, tile), '.')

# Pixel area tiles-- necessary for calculating sum of pixels for any set of tiles
uu.s3_folder_download(cn.pixel_area_dir, '.')

# For downloading all tiles in selected folders
download_list = [
                # # cn.WHRC_biomass_2000_unmasked_dir,
                # # cn.mangrove_biomass_2000_dir,
                # # cn.cont_eco_dir,
                # # cn.WHRC_biomass_2000_non_mang_non_planted_dir,
                # # cn.gain_year_count_mangrove_dir,
                # cn.annual_gain_AGB_mangrove_dir,
                # cn.annual_gain_BGB_mangrove_dir,
                # cn.cumul_gain_AGCO2_mangrove_dir,
                # cn.cumul_gain_BGCO2_mangrove_dir,
                #
                # # cn.gain_year_count_planted_forest_non_mangrove_dir,
                # cn.annual_gain_AGC_BGC_planted_forest_unmasked_dir,
                # cn.planted_forest_type_unmasked_dir,
                # cn.annual_gain_AGB_planted_forest_non_mangrove_dir,
                # cn.annual_gain_BGB_planted_forest_non_mangrove_dir,
                # cn.cumul_gain_AGCO2_planted_forest_non_mangrove_dir,
                # cn.cumul_gain_BGCO2_planted_forest_non_mangrove_dir

                # cn.age_cat_natrl_forest_dir,
                # # cn.gain_year_count_natrl_forest_dir,
                # cn.annual_gain_AGB_natrl_forest_dir,
                # cn.annual_gain_BGB_natrl_forest_dir,
                # cn.cumul_gain_AGCO2_natrl_forest_dir,
                # cn.cumul_gain_BGCO2_natrl_forest_dir,
                # cn.annual_gain_AGB_BGB_all_types_dir,
                # cn.cumul_gain_AGCO2_BGCO2_all_types_dir

                # cn.elevation_processed_dir,
                # cn.precip_processed_dir,
                # cn.planted_forest_type_unmasked_dir,
                # cn.bor_tem_trop_processed_dir,
                # cn.soil_C_full_extent_2000_dir,
                # cn.drivers_processed_dir,
                # cn.climate_zone_processed_dir

                # cn.AGC_emis_year_dir,
                # cn.BGC_emis_year_dir,
                # cn.deadwood_emis_year_2000_dir,
                # cn.litter_emis_year_2000_dir,
                # cn.soil_C_emis_year_2000_dir,
                # cn.total_C_emis_year_dir

                # cn.gross_emis_all_gases_all_drivers_biomass_soil_dir,
                cn.net_flux_dir,
                cn.gross_emis_co2_only_all_drivers_biomass_soil_dir,
                cn.gross_emis_non_co2_all_drivers_biomass_soil_dir,
                cn.gross_emis_commod_biomass_soil_dir,
                cn.gross_emis_shifting_ag_biomass_soil_dir,
                cn.gross_emis_forestry_biomass_soil_dir,
                cn.gross_emis_wildfire_biomass_soil_dir,
                cn.gross_emis_urban_biomass_soil_dir,
                cn.gross_emis_no_driverv_biomass_soil_dir,
                cn.gross_emis_nodes_biomass_soil_dir

]

# Iterates through each set of tiles and gets statistics of it
for input in download_list:

    uu.s3_folder_download(input, '.')

    # List of all the tiles on the spot machine to be summarized (excludes pixel area tiles and tiles created by gdal_calc
    # (in case this script was already run on this spot machine and created output from gdal_calc)
    tile_list = uu.tile_list_spot_machine(".", ".tif")
    # from https://stackoverflow.com/questions/12666897/removing-an-item-from-list-matching-a-substring
    tile_list = [i for i in tile_list if not ('hanson_2013' in i or 'value_per_pixel' in i)]
    # tile_list = ['00N_000E_biomass.tif']
    # tile_list = download_tile_list
    print tile_list

    # For multiprocessor use.
    count = multiprocessing.cpu_count()
    pool = multiprocessing.Pool(processes=9)
    pool.map(tile_statistics.create_tile_statistics, tile_list)
    # Added these in response to error12: Cannot allocate memory error.
    # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
    # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
    pool.close()
    pool.join()

    # # For single processor use
    # for tile in tile_list:
    #     tile_statistics.create_tile_statistics(tile)

    # Even an m4.16xlarge spot machine can't handle all these sets of tiles, so this deletes each set of tiles after it is analyzed
    print "Deleting tiles..."
    for tile in tile_list:
        os.remove(tile)
        tile_short = tile[:-4]
        outname = '{0}_value_per_pixel.tif'.format(tile_short)
        os.remove(outname)
        print "  Tiles deleted"

    # Copies the text file to the location on s3 that the tiles are from
    cmd = ['aws', 's3', 'cp', tile_stats, cn.tile_stats_dir]
    subprocess.check_call(cmd)<commit_msg>Tile statistics for gross emissions soil only, model v1.1.<commit_after>### This script calculates various statistics on all tiles in input folders and saves them to a txt.
### Users can input as many folders as they want for calculating statistics on each tile

import multiprocessing
import tile_statistics
import subprocess
import os
import glob
import sys
sys.path.append('../')
import constants_and_names as cn
import universal_util as uu

# The column names for the tile summary statistics.
# If the statistics calculations are changed in tile_statistics.py, the list here needs to be changed, too.
headers = ['tile_id', 'tile_type', 'tile_name', 'pixel_count', 'mean', 'median', 'percentile10', 'percentile25',
           'percentile75', 'percentile90', 'min', 'max', 'sum']
header_no_brackets = ', '.join(headers)

tile_stats = '{}_{}'.format(uu.date_today, cn.tile_stats_pattern)

# Creates the output text file with the column names
with open(cn.tile_stats_pattern, 'w+') as f:
    f.write(header_no_brackets  +'\r\n')
f.close()

# # Creates list of tiles to iterate through, for testing
# download_tile_list = ['00N_110E', '30N_080W', '40N_050E', '50N_100E', '80N_020E'] # test tiles
# print download_tile_list
#
# # For copying individual tiles to spot machine for testing
# for tile in download_tile_list:
#     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.mangrove_biomass_2000_dir, cn.pattern_mangrove_biomass_2000, tile), '.')
#     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.natrl_forest_biomass_2000_dir, tile, cn.pattern_natrl_forest_biomass_2000), '.')
#     uu.s3_file_download('{0}{1}_{2}.tif'.format(cn.pixel_area_dir, cn.pattern_pixel_area, tile), '.')

# Pixel area tiles-- necessary for calculating sum of pixels for any set of tiles
uu.s3_folder_download(cn.pixel_area_dir, '.')

# For downloading all tiles in selected folders
download_list = [
                # # cn.WHRC_biomass_2000_unmasked_dir,
                # # cn.mangrove_biomass_2000_dir,
                # # cn.cont_eco_dir,
                # # cn.WHRC_biomass_2000_non_mang_non_planted_dir,
                # # cn.gain_year_count_mangrove_dir,
                # cn.annual_gain_AGB_mangrove_dir,
                # cn.annual_gain_BGB_mangrove_dir,
                # cn.cumul_gain_AGCO2_mangrove_dir,
                # cn.cumul_gain_BGCO2_mangrove_dir,
                #
                # # cn.gain_year_count_planted_forest_non_mangrove_dir,
                # cn.annual_gain_AGC_BGC_planted_forest_unmasked_dir,
                # cn.planted_forest_type_unmasked_dir,
                # cn.annual_gain_AGB_planted_forest_non_mangrove_dir,
                # cn.annual_gain_BGB_planted_forest_non_mangrove_dir,
                # cn.cumul_gain_AGCO2_planted_forest_non_mangrove_dir,
                # cn.cumul_gain_BGCO2_planted_forest_non_mangrove_dir

                # cn.age_cat_natrl_forest_dir,
                # # cn.gain_year_count_natrl_forest_dir,
                # cn.annual_gain_AGB_natrl_forest_dir,
                # cn.annual_gain_BGB_natrl_forest_dir,
                # cn.cumul_gain_AGCO2_natrl_forest_dir,
                # cn.cumul_gain_BGCO2_natrl_forest_dir,
                # cn.annual_gain_AGB_BGB_all_types_dir,
                # cn.cumul_gain_AGCO2_BGCO2_all_types_dir

                # cn.elevation_processed_dir,
                # cn.precip_processed_dir,
                # cn.planted_forest_type_unmasked_dir,
                # cn.bor_tem_trop_processed_dir,
                # cn.soil_C_full_extent_2000_dir,
                # cn.drivers_processed_dir,
                # cn.climate_zone_processed_dir

                # cn.AGC_emis_year_dir,
                # cn.BGC_emis_year_dir,
                # cn.deadwood_emis_year_2000_dir,
                # cn.litter_emis_year_2000_dir,
                # cn.soil_C_emis_year_2000_dir,
                # cn.total_C_emis_year_dir

                # cn.gross_emis_all_gases_all_drivers_biomass_soil_dir,
                # cn.net_flux_dir,
                # cn.gross_emis_co2_only_all_drivers_biomass_soil_dir,
                # cn.gross_emis_non_co2_all_drivers_biomass_soil_dir,
                # cn.gross_emis_commod_biomass_soil_dir,
                # cn.gross_emis_shifting_ag_biomass_soil_dir,
                # cn.gross_emis_forestry_biomass_soil_dir,
                # cn.gross_emis_wildfire_biomass_soil_dir,
                # cn.gross_emis_urban_biomass_soil_dir,
                # cn.gross_emis_no_driverv_biomass_soil_dir,
                # cn.gross_emis_nodes_biomass_soil_dir

                cn.gross_emis_co2_only_all_drivers_soil_only_dir,
                cn.gross_emis_non_co2_all_drivers_soil_only_dir,
                cn.gross_emis_commod_soil_only_dir,
                cn.gross_emis_shifting_ag_soil_only_dir,
                cn.gross_emis_forestry_soil_only_dir,
                cn.gross_emis_wildfire_soil_only_dir,
                cn.gross_emis_urban_soil_only_dir,
                cn.gross_emis_no_driverv_soil_only_dir,
                cn.gross_emis_nodes_soil_only_dir

]

# Iterates through each set of tiles and gets statistics of it
for input in download_list:

    uu.s3_folder_download(input, '.')

    # List of all the tiles on the spot machine to be summarized (excludes pixel area tiles and tiles created by gdal_calc
    # (in case this script was already run on this spot machine and created output from gdal_calc)
    tile_list = uu.tile_list_spot_machine(".", ".tif")
    # from https://stackoverflow.com/questions/12666897/removing-an-item-from-list-matching-a-substring
    tile_list = [i for i in tile_list if not ('hanson_2013' in i or 'value_per_pixel' in i)]
    # tile_list = ['00N_000E_biomass.tif']
    # tile_list = download_tile_list
    print tile_list

    # For multiprocessor use.
    count = multiprocessing.cpu_count()
    pool = multiprocessing.Pool(processes=9)
    pool.map(tile_statistics.create_tile_statistics, tile_list)
    # Added these in response to error12: Cannot allocate memory error.
    # This fix was mentioned here: of https://stackoverflow.com/questions/26717120/python-cannot-allocate-memory-using-multiprocessing-pool
    # Could also try this: https://stackoverflow.com/questions/42584525/python-multiprocessing-debugging-oserror-errno-12-cannot-allocate-memory
    pool.close()
    pool.join()

    # # For single processor use
    # for tile in tile_list:
    #     tile_statistics.create_tile_statistics(tile)

    # Even an m4.16xlarge spot machine can't handle all these sets of tiles, so this deletes each set of tiles after it is analyzed
    print "Deleting tiles..."
    for tile in tile_list:
        os.remove(tile)
        tile_short = tile[:-4]
        outname = '{0}_value_per_pixel.tif'.format(tile_short)
        os.remove(outname)
        print "  Tiles deleted"

    # Copies the text file to the location on s3 that the tiles are from
    cmd = ['aws', 's3', 'cp', tile_stats, cn.tile_stats_dir]
    subprocess.check_call(cmd)<|endoftext|>


<commit_before>import os
import csv
import numpy as np

def get_scalars_csv_filename(model_dir):
    return os.path.join(model_dir, 'stats', 'scalars.csv')

def summarize_training(jobdir):
    csv_filenames = []
    for run in os.listdir(jobdir):
        csv_filename = os.path.join(jobdir, get_scalars_csv_filename(run))
        if os.path.isdir(os.path.join(jobdir, run)):
            csv_filenames.append(csv_filename)

    # collect final lines of training stats
    for i, csv_filename in enumerate(csv_filenames):
        with open(csv_filename, 'r', newline='') as f:
            #reader = csv.DictReader(f)
            reader = f.readlines()

            #for i, row in enumerate(reader)
            if i == 0:
                headers = reader[0].split(',')
                stats_dict = {name: [] for name in headers}
            for j, name in enumerate(headers):
                stats_dict[name].append(reader[-1][j])

    # compute aggregate stats
    aggregate_stats_dict = {}
    for name in headers:
        try:
            mean_stat = np.mean(np.array(stats_dict[name]))
            std_stat = np.std(np.array(stats_dict[name]))
        except TypeError as e:
            print(e, name, stats_dict[name])
            mean_stat = None
            std_stat = None
        aggregate_stats_dict[name] = mean_stat, std_stat

    # pretty print to results file
    with open(os.path.join(jobdir, 'stats.txt'), 'w') as f:
        for name, (mean, std) in aggregate_stats_dict.items():
            if mean is not None:
                # first print the aggregated stats
                f.write('{}: mean = {}, std = {}\n'.format(name, mean, std))
                # then the collected stats
                f.write('Individual statistics\n')
                for s in stats_dict[name]: f.write('{}\n'.format(s))
<commit_msg>debug<commit_after>import os
import csv
import numpy as np

def get_scalars_csv_filename(model_dir):
    return os.path.join(model_dir, 'stats', 'scalars.csv')

def summarize_training(jobdir):
    csv_filenames = []
    for run in os.listdir(jobdir):
        csv_filename = os.path.join(jobdir, get_scalars_csv_filename(run))
        if os.path.isdir(os.path.join(jobdir, run)):
            csv_filenames.append(csv_filename)

    # collect final lines of training stats
    for i, csv_filename in enumerate(csv_filenames):
        with open(csv_filename, 'r', newline='') as f:
            #reader = csv.DictReader(f)
            reader = f.readlines()

            #for i, row in enumerate(reader)
            if i == 0:
                headers = reader[0].split(',')
                stats_dict = {name: [] for name in headers}
            for j, name in enumerate(headers):
                stats_dict[name].append(reader[-1].split(',')[j])

    # compute aggregate stats
    aggregate_stats_dict = {}
    for name in headers:
        try:
            mean_stat = np.mean(np.array(stats_dict[name]))
            std_stat = np.std(np.array(stats_dict[name]))
        except TypeError as e:
            print(e, name, stats_dict[name])
            mean_stat = None
            std_stat = None
        aggregate_stats_dict[name] = mean_stat, std_stat

    # pretty print to results file
    with open(os.path.join(jobdir, 'stats.txt'), 'w') as f:
        for name, (mean, std) in aggregate_stats_dict.items():
            if mean is not None:
                # first print the aggregated stats
                f.write('{}: mean = {}, std = {}\n'.format(name, mean, std))
                # then the collected stats
                f.write('Individual statistics\n')
                for s in stats_dict[name]: f.write('{}\n'.format(s))
<|endoftext|>


<commit_before><commit_msg>Add the __version__ variable to __init__ See: http://goo.gl/ixKV5T<commit_after>__version__ = '1.3'
<|endoftext|>


<commit_before># Copyright 2017-2018 TensorHub, Inc.
#
# 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.

from __future__ import absolute_import
from __future__ import division

import hashlib
import logging
import os
import pkg_resources
import re
import subprocess

import yaml

import guild

from guild import cli
from guild import config
from guild import init
from guild import namespace
from guild import util

log = logging.getLogger("guild")

class Config(object):

    def __init__(self, args):
        self.env_dir = os.path.abspath(args.dir)
        self.env_name = self._init_env_name(args.name, self.env_dir)
        self.venv_python_ver = self._init_venv_python(args)
        self.guild = args.guild
        self.guild_pkg_reqs = self._init_guild_pkg_reqs(args)
        self.user_reqs = self._init_user_reqs(args)
        self.paths = args.path
        self.tensorflow = args.tf_package
        self.local_resource_cache = args.local_resource_cache
        self.prompt_params = self._init_prompt_params()
        self.no_progress = args.no_progress

    @staticmethod
    def _init_env_name(name, abs_env_dir):
        if name:
            return name
        elif os.path.exists(abs_env_dir):
            return os.path.basename(os.path.dirname(abs_env_dir))
        else:
            return os.path.basename(abs_env_dir)

    @staticmethod
    def _init_venv_python(args):
        arg = args.python
        if not arg or arg.startswith("python"):
            return arg
        return "python{}".format(arg)

    @staticmethod
    def _init_guild_pkg_reqs(args):
        if args.no_reqs:
            return ()
        return tuple(sorted(
            _iter_all_guild_pkg_reqs(config.cwd(), args.path)))

    @staticmethod
    def _init_user_reqs(args):
        # -r options may be used with --no-reqs, in which case -r
        # takes precedence (--no-reqs may still be used to suppress
        # installation of Guild package reqs)
        if args.requirement:
            _validate_req_files(args.requirement)
            return args.requirement
        elif not args.no_reqs:
            default_reqs = os.path.join(config.cwd(), "requirements.txt")
            if os.path.exists(default_reqs):
                return (default_reqs,)
        return ()

    def _init_prompt_params(self):
        params = []
        params.append(("Location", _shorten_path(self.env_dir)))
        params.append(("Name", self.env_name))
        if self.venv_python_ver:
            params.append(("Python version", self.venv_python_ver))
        else:
            params.append(("Python version", "default"))
        if self.guild:
            params.append(("Guild version", self.guild))
        else:
            params.append(("Guild version", _implicit_guild_version()))
        if self.guild_pkg_reqs:
            params.append(("Guild package requirements", self.guild_pkg_reqs))
        if self.user_reqs:
            params.append(("Python requirements", self.user_reqs))
        if self.paths:
            params.append(("Additional paths", self.paths))
        if self.tensorflow == "no":
            params.append(("TensorFlow support", "no"))
        elif self.tensorflow == "tensorflow":
            params.append(("TensorFlow support", "non-GPU"))
        elif self.tensorflow == "tensorflow-gpu":
            params.append(("TensorFlow support", "GPU"))
        elif self.tensorflow is None:
            params.append(("TensorFlow support", "auto-detect GPU"))
        else:
            raise AssertionError(self.tensorflow)
        if self.local_resource_cache:
            params.append(("Resource cache", "local"))
        else:
            params.append(("Resource cache", "shared"))
        return params

    def _maybe_guild_pkg_reqs(self):
        if self.no_guild_pkg_reqs:
            return []

    def as_kw(self):
        return self.__dict__

def _iter_all_guild_pkg_reqs(dir, search_path, seen=None):
    seen = seen or set()
    src = os.path.abspath(os.path.join(dir, "guild.yml"))
    if not os.path.exists(src):
        return
    if src in seen:
        return
    seen.add(src)
    for req in _guild_pkg_reqs(src):
        dir_on_path = _find_req_on_path(req, search_path)
        if dir_on_path:
            for req in _iter_all_guild_pkg_reqs(dir_on_path, search_path, seen):
                yield req
        else:
            yield req

def _guild_pkg_reqs(src):
    pkg = _pkg_for_guildfile(src)
    if not pkg:
        return []
    return _pkg_requires(pkg, src)

def _pkg_for_guildfile(src):
    data = _guildfile_data(src)
    if not isinstance(data, list):
        data = [data]
    for item in data:
        if isinstance(item, dict) and "package" in item:
            return item
    return None

def _guildfile_data(src):
    # Use low level parsing to bypass path-related errors.
    try:
        f = open(src, "r")
        return yaml.load(f)
    except Exception as e:
        log.warning(
            "cannot read Guild package requirements for %s (%s) - ignoring",
            src, e)
        return []

def _pkg_requires(pkg_data, src):
    requires = pkg_data.get("requires") or []
    if not isinstance(requires, list):
        log.warning(
            "invalid package requires list in %s (%r) - ignoring",
            src, pkg_data)
        return []
    return requires

def _find_req_on_path(req, path):
    req_subpath = req.replace(".", os.path.sep)
    for root in path:
        full_path = os.path.join(root, req_subpath)
        if os.path.exists(full_path):
            return full_path
    return None

def _validate_guild_reqs(reqs, guildfile_path):
    # Make sure we can convert each to pip reqs
    for req in reqs:
        try:
            namespace.pip_info(req)
        except namespace.NamespaceError:
            cli.error(
                "invalid required Guild package '%s' in %s\n"
                "Either correct the package name or use --no-reqs "
                "to skip installing required Guild packages"
                % (req, guildfile_path))
    return reqs

def _validate_req_files(paths):
    for p in paths:
        if not os.path.exists(p):
            cli.error("requirement file %s does not exist" % p)

def _shorten_path(path):
    return path.replace(os.path.expanduser("~"), "~")

def _implicit_guild_version():
    reqs_file = _guild_reqs_file()
    if reqs_file:
        return "from source (%s)" % os.path.dirname(reqs_file)
    else:
        return guild.__version__

def main(args):
    _error_if_active_env()
    config = Config(args)
    if args.yes or _confirm(config):
        _init(config)

def _error_if_active_env():
    active_env = os.getenv("VIRTUAL_ENV")
    if active_env:
        cli.error(
            "cannot run init from an activate environment (%s)\n"
            "Deactivate the environment by running 'deactivate' "
            "and try again."
            % active_env)

def _confirm(config):
    cli.out("You are about to initialize a Guild environment:")
    for name, val in config.prompt_params:
        if isinstance(val, tuple):
            cli.out("  {}:".format(name))
            for x in val:
                cli.out("    {}".format(x))
        else:
            cli.out("  {}: {}".format(name, val))
    return cli.confirm("Continue?", default=True)

def _init(config):
    _init_guild_env(config)
    _init_venv(config)
    _install_guild(config)
    _install_guild_pkg_reqs(config)
    _install_user_reqs(config)
    _install_paths(config)
    _ensure_tensorflow(config)
    _initialized_msg(config)

def _init_guild_env(config):
    cli.out("Initializing Guild environment in {}".format(config.env_dir))
    init.init_env(config.env_dir, config.local_resource_cache)

def _init_venv(config):
    cmd_args = _venv_cmd_args(config)
    if not cmd_args:
        cli.out("Skipping virtual env")
        return
    cli.out("Creating virtual environment")
    try:
        subprocess.check_call(cmd_args)
    except subprocess.CalledProcessError as e:
        cli.error(str(e), exit_status=e.returncode)

def _venv_cmd_args(config):
    args = ["virtualenv", config.env_dir]
    args.extend(["--prompt", "({}) ".format(config.env_name)])
    if config.venv_python_ver:
        args.extend(["--python", "python{}".format(config.venv_python)])
    return args

def _install_guild(config):
    """Install Guild into env.

    If the running Guild program is from `guild/scripts` (i.e. running
    in dev mode), only the Guild requirements, as defined in
    `guild/requirements.txt` are installed, under the assumption that
    `guild/scripts` is in `PATH` or an alias is being used and
    therefore Guild will be available from the env. In this case only
    the requirements need to be installed.

    If the running Guild program is not from `guild/scripts`, Guild
    will be installed using the env `pip` program.
    """
    if config.guild:
        _install_guild_dist(config)
    else:
        guild_reqs = _guild_reqs_file()
        if guild_reqs:
            _install_guild_reqs(guild_reqs, config)
        else:
            _install_default_guild_dist(config)

def _install_guild_dist(config):
    assert config.guild
    # If config.guild can be parsed as a version, use as
    # 'guildai==VERSION'
    if config.guild[0].isdigit():
        cli.out("Installing Guild %s" % config.guild)
        req = "guildai==%s" % config.guild
    else:
        cli.out("Installing %s" % config.guild)
        req = config.guild
    _install_reqs([req], config)

def _guild_reqs_file():
    guild_location = pkg_resources.resource_filename("guild", "")
    guild_parent = os.path.dirname(guild_location)
    path = os.path.join(guild_parent, "requirements.txt")
    try:
        f = open(path, "r")
    except OSError:
        pass
    else:
        with f:
            if "guildai" in f.readline():
                return path
    return None

def _install_guild_reqs(req_files, config):
    cli.out("Installing Guild requirements")
    _install_req_files([req_files], config)

def _install_default_guild_dist(config):
    req = "guildai==%s" % guild.__version__
    cli.out("Installing Guild %s" % req)
    _install_reqs([req], config)

def _install_reqs(reqs, config):
    cmd = (
        [_pip_bin(config.env_dir), "install"] +
        _pip_extra_opts(config) +
        reqs
    )
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        cli.error(str(e), exit_status=e.returncode)

def _pip_bin(env_dir):
    pip_bin = os.path.join(env_dir, "bin", "pip")
    assert os.path.exists(pip_bin), pip_bin
    return pip_bin

def _pip_extra_opts(config):
    if config.no_progress:
        return ["--progress", "off"]
    else:
        return []

def _install_req_files(req_files, config):
    cmd_args = [_pip_bin(config.env_dir), "install"] + _pip_extra_opts(config)
    for path in req_files:
        cmd_args.extend(["-r", path])
    try:
        subprocess.check_call(cmd_args)
    except subprocess.CalledProcessError as e:
        cli.error(str(e), exit_status=e.returncode)

def _install_guild_pkg_reqs(config):
    if config.guild_pkg_reqs:
        cli.out("Installing Python requirements")
        reqs = _pip_package_reqs(config.guild_pkg_reqs)
        _install_reqs(reqs, config)

def _pip_package_reqs(guild_reqs):
    reqs = []
    for guild_req in guild_reqs:
        pip_info = namespace.pip_info(guild_req)
        reqs.append(pip_info.project_name)
    return reqs

def _install_user_reqs(config):
    if config.user_reqs:
        cli.out("Installing Python requirements")
        _install_req_files(config.user_reqs, config)

def _install_paths(config):
    if not config.paths:
        return
    site_packages = _env_site_packages(config.env_dir)
    for path in config.paths:
        _write_path(path, site_packages)

def _env_site_packages(env_dir):
    python_bin = os.path.join(env_dir, "bin", "python")
    out = subprocess.check_output([python_bin, "-m", "site"])
    site_packages_pattern = r"(%s.+?site-packages)" % env_dir
    m = re.search(site_packages_pattern, out)
    assert m, out
    return m.group(1)

def _write_path(path, target_dir):
    path = os.path.abspath(path)
    with open(_pth_filename(target_dir, path), "w") as f:
        f.write(path)

def _pth_filename(target_dir, path):
    digest = hashlib.md5(path).hexdigest()[:8]
    pth_name = "%s-%s.pth" % (os.path.basename(path), digest)
    return os.path.join(target_dir, pth_name)

def _ensure_tensorflow(config):
    if config.tensorflow == "no":
        cli.out("Skipping TensorFlow installation")
        return
    if _tensorflow_installed(config.env_dir):
        cli.out("TensorFlow already installed, skipping installation")
        return
    tf_pkg = _tensorflow_package(config)
    cli.out("Installing TensorFlow (%s)" % tf_pkg)
    _install_reqs([tf_pkg], config)

def _tensorflow_installed(env_dir):
    cli.out("Checking for TensorFlow")
    python_bin = os.path.join(env_dir, "bin", "python")
    assert os.path.exists(python_bin)
    cmd_args = [python_bin, "-c", "import tensorflow; tensorflow.__version__"]
    try:
        subprocess.check_output(cmd_args, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError:
        return False
    else:
        return True

def _tensorflow_package(config):
    if config.tensorflow is not None:
        return config.tensorflow
    if util.gpu_available():
        return "tensorflow-gpu"
    else:
        return "tensorflow"

def _initialized_msg(config):
    cli.out(
        "Guild environment initialized in {}."
        "\n".format(_shorten_path(config.env_dir)))
    cli.out("To activate it " "run:\n")
    env_parent = os.path.dirname(config.env_dir)
    if env_parent != os.getcwd():
        cli.out("  source guild-env {}".format(config.env_dir))
    else:
        env_name = os.path.basename(config.env_dir)
        if env_name != "env":
            cli.out("  source guild-env {}".format(env_name))
        else:
            cli.out("  source guild-env")
    cli.out()

def _source_cmd(env_dir):
    name = os.path.basename(env_dir)
    if name == "env":
        return "source guild-env"
    else:
        return "source guild-env {}".format(name)
<commit_msg>Handle IOError for file check<commit_after># Copyright 2017-2018 TensorHub, Inc.
#
# 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.

from __future__ import absolute_import
from __future__ import division

import hashlib
import logging
import os
import pkg_resources
import re
import subprocess

import yaml

import guild

from guild import cli
from guild import config
from guild import init
from guild import namespace
from guild import util

log = logging.getLogger("guild")

class Config(object):

    def __init__(self, args):
        self.env_dir = os.path.abspath(args.dir)
        self.env_name = self._init_env_name(args.name, self.env_dir)
        self.venv_python_ver = self._init_venv_python(args)
        self.guild = args.guild
        self.guild_pkg_reqs = self._init_guild_pkg_reqs(args)
        self.user_reqs = self._init_user_reqs(args)
        self.paths = args.path
        self.tensorflow = args.tf_package
        self.local_resource_cache = args.local_resource_cache
        self.prompt_params = self._init_prompt_params()
        self.no_progress = args.no_progress

    @staticmethod
    def _init_env_name(name, abs_env_dir):
        if name:
            return name
        elif os.path.exists(abs_env_dir):
            return os.path.basename(os.path.dirname(abs_env_dir))
        else:
            return os.path.basename(abs_env_dir)

    @staticmethod
    def _init_venv_python(args):
        arg = args.python
        if not arg or arg.startswith("python"):
            return arg
        return "python{}".format(arg)

    @staticmethod
    def _init_guild_pkg_reqs(args):
        if args.no_reqs:
            return ()
        return tuple(sorted(
            _iter_all_guild_pkg_reqs(config.cwd(), args.path)))

    @staticmethod
    def _init_user_reqs(args):
        # -r options may be used with --no-reqs, in which case -r
        # takes precedence (--no-reqs may still be used to suppress
        # installation of Guild package reqs)
        if args.requirement:
            _validate_req_files(args.requirement)
            return args.requirement
        elif not args.no_reqs:
            default_reqs = os.path.join(config.cwd(), "requirements.txt")
            if os.path.exists(default_reqs):
                return (default_reqs,)
        return ()

    def _init_prompt_params(self):
        params = []
        params.append(("Location", _shorten_path(self.env_dir)))
        params.append(("Name", self.env_name))
        if self.venv_python_ver:
            params.append(("Python version", self.venv_python_ver))
        else:
            params.append(("Python version", "default"))
        if self.guild:
            params.append(("Guild version", self.guild))
        else:
            params.append(("Guild version", _implicit_guild_version()))
        if self.guild_pkg_reqs:
            params.append(("Guild package requirements", self.guild_pkg_reqs))
        if self.user_reqs:
            params.append(("Python requirements", self.user_reqs))
        if self.paths:
            params.append(("Additional paths", self.paths))
        if self.tensorflow == "no":
            params.append(("TensorFlow support", "no"))
        elif self.tensorflow == "tensorflow":
            params.append(("TensorFlow support", "non-GPU"))
        elif self.tensorflow == "tensorflow-gpu":
            params.append(("TensorFlow support", "GPU"))
        elif self.tensorflow is None:
            params.append(("TensorFlow support", "auto-detect GPU"))
        else:
            raise AssertionError(self.tensorflow)
        if self.local_resource_cache:
            params.append(("Resource cache", "local"))
        else:
            params.append(("Resource cache", "shared"))
        return params

    def _maybe_guild_pkg_reqs(self):
        if self.no_guild_pkg_reqs:
            return []

    def as_kw(self):
        return self.__dict__

def _iter_all_guild_pkg_reqs(dir, search_path, seen=None):
    seen = seen or set()
    src = os.path.abspath(os.path.join(dir, "guild.yml"))
    if not os.path.exists(src):
        return
    if src in seen:
        return
    seen.add(src)
    for req in _guild_pkg_reqs(src):
        dir_on_path = _find_req_on_path(req, search_path)
        if dir_on_path:
            for req in _iter_all_guild_pkg_reqs(dir_on_path, search_path, seen):
                yield req
        else:
            yield req

def _guild_pkg_reqs(src):
    pkg = _pkg_for_guildfile(src)
    if not pkg:
        return []
    return _pkg_requires(pkg, src)

def _pkg_for_guildfile(src):
    data = _guildfile_data(src)
    if not isinstance(data, list):
        data = [data]
    for item in data:
        if isinstance(item, dict) and "package" in item:
            return item
    return None

def _guildfile_data(src):
    # Use low level parsing to bypass path-related errors.
    try:
        f = open(src, "r")
        return yaml.load(f)
    except Exception as e:
        log.warning(
            "cannot read Guild package requirements for %s (%s) - ignoring",
            src, e)
        return []

def _pkg_requires(pkg_data, src):
    requires = pkg_data.get("requires") or []
    if not isinstance(requires, list):
        log.warning(
            "invalid package requires list in %s (%r) - ignoring",
            src, pkg_data)
        return []
    return requires

def _find_req_on_path(req, path):
    req_subpath = req.replace(".", os.path.sep)
    for root in path:
        full_path = os.path.join(root, req_subpath)
        if os.path.exists(full_path):
            return full_path
    return None

def _validate_guild_reqs(reqs, guildfile_path):
    # Make sure we can convert each to pip reqs
    for req in reqs:
        try:
            namespace.pip_info(req)
        except namespace.NamespaceError:
            cli.error(
                "invalid required Guild package '%s' in %s\n"
                "Either correct the package name or use --no-reqs "
                "to skip installing required Guild packages"
                % (req, guildfile_path))
    return reqs

def _validate_req_files(paths):
    for p in paths:
        if not os.path.exists(p):
            cli.error("requirement file %s does not exist" % p)

def _shorten_path(path):
    return path.replace(os.path.expanduser("~"), "~")

def _implicit_guild_version():
    reqs_file = _guild_reqs_file()
    if reqs_file:
        return "from source (%s)" % os.path.dirname(reqs_file)
    else:
        return guild.__version__

def main(args):
    _error_if_active_env()
    config = Config(args)
    if args.yes or _confirm(config):
        _init(config)

def _error_if_active_env():
    active_env = os.getenv("VIRTUAL_ENV")
    if active_env:
        cli.error(
            "cannot run init from an activate environment (%s)\n"
            "Deactivate the environment by running 'deactivate' "
            "and try again."
            % active_env)

def _confirm(config):
    cli.out("You are about to initialize a Guild environment:")
    for name, val in config.prompt_params:
        if isinstance(val, tuple):
            cli.out("  {}:".format(name))
            for x in val:
                cli.out("    {}".format(x))
        else:
            cli.out("  {}: {}".format(name, val))
    return cli.confirm("Continue?", default=True)

def _init(config):
    _init_guild_env(config)
    _init_venv(config)
    _install_guild(config)
    _install_guild_pkg_reqs(config)
    _install_user_reqs(config)
    _install_paths(config)
    _ensure_tensorflow(config)
    _initialized_msg(config)

def _init_guild_env(config):
    cli.out("Initializing Guild environment in {}".format(config.env_dir))
    init.init_env(config.env_dir, config.local_resource_cache)

def _init_venv(config):
    cmd_args = _venv_cmd_args(config)
    if not cmd_args:
        cli.out("Skipping virtual env")
        return
    cli.out("Creating virtual environment")
    try:
        subprocess.check_call(cmd_args)
    except subprocess.CalledProcessError as e:
        cli.error(str(e), exit_status=e.returncode)

def _venv_cmd_args(config):
    args = ["virtualenv", config.env_dir]
    args.extend(["--prompt", "({}) ".format(config.env_name)])
    if config.venv_python_ver:
        args.extend(["--python", "python{}".format(config.venv_python)])
    return args

def _install_guild(config):
    """Install Guild into env.

    If the running Guild program is from `guild/scripts` (i.e. running
    in dev mode), only the Guild requirements, as defined in
    `guild/requirements.txt` are installed, under the assumption that
    `guild/scripts` is in `PATH` or an alias is being used and
    therefore Guild will be available from the env. In this case only
    the requirements need to be installed.

    If the running Guild program is not from `guild/scripts`, Guild
    will be installed using the env `pip` program.
    """
    if config.guild:
        _install_guild_dist(config)
    else:
        guild_reqs = _guild_reqs_file()
        if guild_reqs:
            _install_guild_reqs(guild_reqs, config)
        else:
            _install_default_guild_dist(config)

def _install_guild_dist(config):
    assert config.guild
    # If config.guild can be parsed as a version, use as
    # 'guildai==VERSION'
    if config.guild[0].isdigit():
        cli.out("Installing Guild %s" % config.guild)
        req = "guildai==%s" % config.guild
    else:
        cli.out("Installing %s" % config.guild)
        req = config.guild
    _install_reqs([req], config)

def _guild_reqs_file():
    guild_location = pkg_resources.resource_filename("guild", "")
    guild_parent = os.path.dirname(guild_location)
    path = os.path.join(guild_parent, "requirements.txt")
    try:
        f = open(path, "r")
    except (IOError, OSError):
        pass
    else:
        with f:
            if "guildai" in f.readline():
                return path
    return None

def _install_guild_reqs(req_files, config):
    cli.out("Installing Guild requirements")
    _install_req_files([req_files], config)

def _install_default_guild_dist(config):
    req = "guildai==%s" % guild.__version__
    cli.out("Installing Guild %s" % req)
    _install_reqs([req], config)

def _install_reqs(reqs, config):
    cmd = (
        [_pip_bin(config.env_dir), "install"] +
        _pip_extra_opts(config) +
        reqs
    )
    try:
        subprocess.check_call(cmd)
    except subprocess.CalledProcessError as e:
        cli.error(str(e), exit_status=e.returncode)

def _pip_bin(env_dir):
    pip_bin = os.path.join(env_dir, "bin", "pip")
    assert os.path.exists(pip_bin), pip_bin
    return pip_bin

def _pip_extra_opts(config):
    if config.no_progress:
        return ["--progress", "off"]
    else:
        return []

def _install_req_files(req_files, config):
    cmd_args = [_pip_bin(config.env_dir), "install"] + _pip_extra_opts(config)
    for path in req_files:
        cmd_args.extend(["-r", path])
    try:
        subprocess.check_call(cmd_args)
    except subprocess.CalledProcessError as e:
        cli.error(str(e), exit_status=e.returncode)

def _install_guild_pkg_reqs(config):
    if config.guild_pkg_reqs:
        cli.out("Installing Python requirements")
        reqs = _pip_package_reqs(config.guild_pkg_reqs)
        _install_reqs(reqs, config)

def _pip_package_reqs(guild_reqs):
    reqs = []
    for guild_req in guild_reqs:
        pip_info = namespace.pip_info(guild_req)
        reqs.append(pip_info.project_name)
    return reqs

def _install_user_reqs(config):
    if config.user_reqs:
        cli.out("Installing Python requirements")
        _install_req_files(config.user_reqs, config)

def _install_paths(config):
    if not config.paths:
        return
    site_packages = _env_site_packages(config.env_dir)
    for path in config.paths:
        _write_path(path, site_packages)

def _env_site_packages(env_dir):
    python_bin = os.path.join(env_dir, "bin", "python")
    out = subprocess.check_output([python_bin, "-m", "site"])
    site_packages_pattern = r"(%s.+?site-packages)" % env_dir
    m = re.search(site_packages_pattern, out)
    assert m, out
    return m.group(1)

def _write_path(path, target_dir):
    path = os.path.abspath(path)
    with open(_pth_filename(target_dir, path), "w") as f:
        f.write(path)

def _pth_filename(target_dir, path):
    digest = hashlib.md5(path).hexdigest()[:8]
    pth_name = "%s-%s.pth" % (os.path.basename(path), digest)
    return os.path.join(target_dir, pth_name)

def _ensure_tensorflow(config):
    if config.tensorflow == "no":
        cli.out("Skipping TensorFlow installation")
        return
    if _tensorflow_installed(config.env_dir):
        cli.out("TensorFlow already installed, skipping installation")
        return
    tf_pkg = _tensorflow_package(config)
    cli.out("Installing TensorFlow (%s)" % tf_pkg)
    _install_reqs([tf_pkg], config)

def _tensorflow_installed(env_dir):
    cli.out("Checking for TensorFlow")
    python_bin = os.path.join(env_dir, "bin", "python")
    assert os.path.exists(python_bin)
    cmd_args = [python_bin, "-c", "import tensorflow; tensorflow.__version__"]
    try:
        subprocess.check_output(cmd_args, stderr=subprocess.STDOUT)
    except subprocess.CalledProcessError:
        return False
    else:
        return True

def _tensorflow_package(config):
    if config.tensorflow is not None:
        return config.tensorflow
    if util.gpu_available():
        return "tensorflow-gpu"
    else:
        return "tensorflow"

def _initialized_msg(config):
    cli.out(
        "Guild environment initialized in {}."
        "\n".format(_shorten_path(config.env_dir)))
    cli.out("To activate it " "run:\n")
    env_parent = os.path.dirname(config.env_dir)
    if env_parent != os.getcwd():
        cli.out("  source guild-env {}".format(config.env_dir))
    else:
        env_name = os.path.basename(config.env_dir)
        if env_name != "env":
            cli.out("  source guild-env {}".format(env_name))
        else:
            cli.out("  source guild-env")
    cli.out()

def _source_cmd(env_dir):
    name = os.path.basename(env_dir)
    if name == "env":
        return "source guild-env"
    else:
        return "source guild-env {}".format(name)
<|endoftext|>


<commit_before>import numpy as np
import os
import gym
from gym import error, spaces
from gym import utils
from gym.utils import seeding

try:
    import atari_py
except ImportError as e:
    raise error.DependencyNotInstalled("{}. (HINT: you can install Atari dependencies by running 'pip install gym[atari]'.)".format(e))

import logging
logger = logging.getLogger(__name__)

def to_ram(ale):
    ram_size = ale.getRAMSize()
    ram = np.zeros((ram_size),dtype=np.uint8)
    ale.getRAM(ram)
    return ram

class AtariEnv(gym.Env, utils.EzPickle):
    metadata = {'render.modes': ['human', 'rgb_array']}

    def __init__(self, game='pong', obs_type='ram', frameskip=(2, 5), repeat_action_probability=0.):
        """Frameskip should be either a tuple (indicating a random range to
        choose from, with the top value exclude), or an int."""

        utils.EzPickle.__init__(self, game, obs_type)
        assert obs_type in ('ram', 'image')

        self.game_path = atari_py.get_game_path(game)
        if not os.path.exists(self.game_path):
            raise IOError('You asked for game %s but path %s does not exist'%(game, self.game_path))
        self._obs_type = obs_type
        self.frameskip = frameskip
        self.ale = atari_py.ALEInterface()
        self.viewer = None

        # Tune (or disable) ALE's action repeat:
        # https://github.com/openai/gym/issues/349
        assert isinstance(repeat_action_probability, (float, int)), "Invalid repeat_action_probability: {!r}".format(repeat_action_probability)
        self.ale.setFloat('repeat_action_probability'.encode('utf-8'), repeat_action_probability)

        self._seed()

        (screen_width, screen_height) = self.ale.getScreenDims()
        self._buffer = np.empty((screen_height, screen_width, 4), dtype=np.uint8)

        self._action_set = self.ale.getMinimalActionSet()
        self.action_space = spaces.Discrete(len(self._action_set))

        (screen_width,screen_height) = self.ale.getScreenDims()
        if self._obs_type == 'ram':
            self.observation_space = spaces.Box(low=np.zeros(128), high=np.zeros(128)+255)
        elif self._obs_type == 'image':
            self.observation_space = spaces.Box(low=0, high=255, shape=(screen_height, screen_width, 3))
        else:
            raise error.Error('Unrecognized observation type: {}'.format(self._obs_type))

    def _seed(self, seed=None):
        self.np_random, seed1 = seeding.np_random(seed)
        # Derive a random seed. This gets passed as a uint, but gets
        # checked as an int elsewhere, so we need to keep it below
        # 2**31.
        seed2 = seeding.hash_seed(seed1 + 1) % 2**31
        # Empirically, we need to seed before loading the ROM.
        self.ale.setInt(b'random_seed', seed2)
        self.ale.loadROM(self.game_path)
        return [seed1, seed2]

    def _step(self, a):
        reward = 0.0
        action = self._action_set[a]

        if isinstance(self.frameskip, int):
            num_steps = self.frameskip
        else:
            num_steps = self.np_random.randint(self.frameskip[0], self.frameskip[1])
        for _ in range(num_steps):
            reward += self.ale.act(action)
        ob = self._get_obs()

        return ob, reward, self.ale.game_over(), {"ale.lives": self.ale.lives()}

    def _get_image(self):
        self.ale.getScreenRGB(self._buffer)  # says rgb but actually bgr on little-endian systems like x86
        return self._buffer[:, :, [2, 1, 0]]

    def _get_ram(self):
        return to_ram(self.ale)

    @property
    def _n_actions(self):
        return len(self._action_set)

    def _get_obs(self):
        if self._obs_type == 'ram':
            return self._get_ram()
        elif self._obs_type == 'image':
            img = self._get_image()
        return img

    # return: (states, observations)
    def _reset(self):
        self.ale.reset_game()
        return self._get_obs()

    def _render(self, mode='human', close=False):
        if close:
            if self.viewer is not None:
                self.viewer.close()
                self.viewer = None
            return
        img = self._get_image()
        if mode == 'rgb_array':
            return img
        elif mode == 'human':
            from gym.envs.classic_control import rendering
            if self.viewer is None:
                self.viewer = rendering.SimpleImageViewer()
            self.viewer.imshow(img)

    def get_action_meanings(self):
        return [ACTION_MEANING[i] for i in self._action_set]

    def get_keys_to_action(self):
        KEYWORD_TO_KEY = {
            'UP':      ord('w'),
            'DOWN':    ord('s'),
            'LEFT':    ord('a'),
            'RIGHT':   ord('d'),
            'FIRE':    ord(' '),
        }

        keys_to_action = {}

        for action_id, action_meaning in enumerate(self.get_action_meanings()):
            keys = []
            for keyword, key in KEYWORD_TO_KEY.items():
                if keyword in action_meaning:
                    keys.append(key)
            keys = tuple(sorted(keys))

            assert keys not in keys_to_action
            keys_to_action[keys] = action_id

        return keys_to_action

    # def save_state(self):
    #     return self.ale.saveState()

    # def load_state(self):
    #     return self.ale.loadState()

    # def clone_state(self):
    #     return self.ale.cloneState()

    # def restore_state(self, state):
    #     return self.ale.restoreState(state)

ACTION_MEANING = {
    0 : "NOOP",
    1 : "FIRE",
    2 : "UP",
    3 : "RIGHT",
    4 : "LEFT",
    5 : "DOWN",
    6 : "UPRIGHT",
    7 : "UPLEFT",
    8 : "DOWNRIGHT",
    9 : "DOWNLEFT",
    10 : "UPFIRE",
    11 : "RIGHTFIRE",
    12 : "LEFTFIRE",
    13 : "DOWNFIRE",
    14 : "UPRIGHTFIRE",
    15 : "UPLEFTFIRE",
    16 : "DOWNRIGHTFIRE",
    17 : "DOWNLEFTFIRE",
}
<commit_msg>Revert "Revert "Optimize Atari envs""<commit_after>import numpy as np
import os
import gym
from gym import error, spaces
from gym import utils
from gym.utils import seeding

try:
    import atari_py
except ImportError as e:
    raise error.DependencyNotInstalled("{}. (HINT: you can install Atari dependencies by running 'pip install gym[atari]'.)".format(e))

import logging
logger = logging.getLogger(__name__)

def to_ram(ale):
    ram_size = ale.getRAMSize()
    ram = np.zeros((ram_size),dtype=np.uint8)
    ale.getRAM(ram)
    return ram

class AtariEnv(gym.Env, utils.EzPickle):
    metadata = {'render.modes': ['human', 'rgb_array']}

    def __init__(self, game='pong', obs_type='ram', frameskip=(2, 5), repeat_action_probability=0.):
        """Frameskip should be either a tuple (indicating a random range to
        choose from, with the top value exclude), or an int."""

        utils.EzPickle.__init__(self, game, obs_type)
        assert obs_type in ('ram', 'image')

        self.game_path = atari_py.get_game_path(game)
        if not os.path.exists(self.game_path):
            raise IOError('You asked for game %s but path %s does not exist'%(game, self.game_path))
        self._obs_type = obs_type
        self.frameskip = frameskip
        self.ale = atari_py.ALEInterface()
        self.viewer = None

        # Tune (or disable) ALE's action repeat:
        # https://github.com/openai/gym/issues/349
        assert isinstance(repeat_action_probability, (float, int)), "Invalid repeat_action_probability: {!r}".format(repeat_action_probability)
        self.ale.setFloat('repeat_action_probability'.encode('utf-8'), repeat_action_probability)

        self._seed()

        (screen_width, screen_height) = self.ale.getScreenDims()
        self._buffer = np.empty((screen_height, screen_width, 3), dtype=np.uint8)

        self._action_set = self.ale.getMinimalActionSet()
        self.action_space = spaces.Discrete(len(self._action_set))

        (screen_width,screen_height) = self.ale.getScreenDims()
        if self._obs_type == 'ram':
            self.observation_space = spaces.Box(low=np.zeros(128), high=np.zeros(128)+255)
        elif self._obs_type == 'image':
            self.observation_space = spaces.Box(low=0, high=255, shape=(screen_height, screen_width, 3))
        else:
            raise error.Error('Unrecognized observation type: {}'.format(self._obs_type))

    def _seed(self, seed=None):
        self.np_random, seed1 = seeding.np_random(seed)
        # Derive a random seed. This gets passed as a uint, but gets
        # checked as an int elsewhere, so we need to keep it below
        # 2**31.
        seed2 = seeding.hash_seed(seed1 + 1) % 2**31
        # Empirically, we need to seed before loading the ROM.
        self.ale.setInt(b'random_seed', seed2)
        self.ale.loadROM(self.game_path)
        return [seed1, seed2]

    def _step(self, a):
        reward = 0.0
        action = self._action_set[a]

        if isinstance(self.frameskip, int):
            num_steps = self.frameskip
        else:
            num_steps = self.np_random.randint(self.frameskip[0], self.frameskip[1])
        for _ in range(num_steps):
            reward += self.ale.act(action)
        ob = self._get_obs()

        return ob, reward, self.ale.game_over(), {"ale.lives": self.ale.lives()}

    def _get_image(self):
        self.ale.getScreenRGB2(self._buffer) # New in atari-py 2017-3-28
        return self._buffer

    def _get_ram(self):
        return to_ram(self.ale)

    @property
    def _n_actions(self):
        return len(self._action_set)

    def _get_obs(self):
        if self._obs_type == 'ram':
            return self._get_ram()
        elif self._obs_type == 'image':
            img = self._get_image()
        return img

    # return: (states, observations)
    def _reset(self):
        self.ale.reset_game()
        return self._get_obs()

    def _render(self, mode='human', close=False):
        if close:
            if self.viewer is not None:
                self.viewer.close()
                self.viewer = None
            return
        img = self._get_image()
        if mode == 'rgb_array':
            return img
        elif mode == 'human':
            from gym.envs.classic_control import rendering
            if self.viewer is None:
                self.viewer = rendering.SimpleImageViewer()
            self.viewer.imshow(img)

    def get_action_meanings(self):
        return [ACTION_MEANING[i] for i in self._action_set]

    def get_keys_to_action(self):
        KEYWORD_TO_KEY = {
            'UP':      ord('w'),
            'DOWN':    ord('s'),
            'LEFT':    ord('a'),
            'RIGHT':   ord('d'),
            'FIRE':    ord(' '),
        }

        keys_to_action = {}

        for action_id, action_meaning in enumerate(self.get_action_meanings()):
            keys = []
            for keyword, key in KEYWORD_TO_KEY.items():
                if keyword in action_meaning:
                    keys.append(key)
            keys = tuple(sorted(keys))

            assert keys not in keys_to_action
            keys_to_action[keys] = action_id

        return keys_to_action

    # def save_state(self):
    #     return self.ale.saveState()

    # def load_state(self):
    #     return self.ale.loadState()

    # def clone_state(self):
    #     return self.ale.cloneState()

    # def restore_state(self, state):
    #     return self.ale.restoreState(state)

ACTION_MEANING = {
    0 : "NOOP",
    1 : "FIRE",
    2 : "UP",
    3 : "RIGHT",
    4 : "LEFT",
    5 : "DOWN",
    6 : "UPRIGHT",
    7 : "UPLEFT",
    8 : "DOWNRIGHT",
    9 : "DOWNLEFT",
    10 : "UPFIRE",
    11 : "RIGHTFIRE",
    12 : "LEFTFIRE",
    13 : "DOWNFIRE",
    14 : "UPRIGHTFIRE",
    15 : "UPLEFTFIRE",
    16 : "DOWNRIGHTFIRE",
    17 : "DOWNLEFTFIRE",
}
<|endoftext|>


<commit_before>__version__ = (1, 2, 0)

def get_version():
    return '.'.join(str(i) for i in __version__)
<commit_msg>Bumped version to 1.3.0.<commit_after>__version__ = (1, 3, 0)

def get_version():
    return '.'.join(str(i) for i in __version__)
<|endoftext|>