<commit_before>import json, os.path, glob
import jsonschema.exceptions
from jsonschema import validate
from PIL import Image

from shared_funcs import yn_input
from hammerhal import ConfigLoader
from logging import getLogger
logger = getLogger('hammerhal.compilers.compiler_base')

class CompilerException(Exception):
    pass

class CompilerBase():

    compiler_type = None
    schema_path = None
    raw_directory = None
    sources_directory = None
    output_directory = None
    output_name = None

    raw = None
    compiled = None

    def __init__(self):
        self.schema_path = "{directory}{type}.json".format(directory=ConfigLoader.get_from_config('schemasDirectory', 'compilers'), type=self.compiler_type)
        self.raw_directory = "{rawRoot}{rawOffset}".format(rawRoot=ConfigLoader.get_from_config('rawDirectoryRoot'), rawOffset=ConfigLoader.get_from_config('compilerTypeSpecific/{type}/rawDirectory'.format(type=self.compiler_type), 'compilers'))
        self.output_directory = "{rawRoot}{rawOffset}".format(rawRoot=ConfigLoader.get_from_config('outputDirectoryRoot', 'compilers'), rawOffset=ConfigLoader.get_from_config('compilerTypeSpecific/{type}/outputDirectory'.format(type=self.compiler_type), 'compilers'))
        self.sources_directory = ConfigLoader.get_from_config('sourcesDirectory', 'compilers')

    def search(self):
        return glob.glob(self.raw_directory + "*.json")

    def find(self, name):
        filename = name
        if (os.path.isfile(filename)):
            return filename

        filename = "{directory}{name}.json".format(directory=self.raw_directory, name=name)
        if (os.path.isfile(filename)):
            return filename

        for filename in self.search():
            try:
                _file = open(filename)
                _json = json.load(_file)
                _file.close()

                _name = _json.get('name', None)
                if (_name and isinstance(_name, str) and _name.lower() in name.lower()):
                    message = "{type} found - {filename}:\n  {name}{subtitle}".format \
                    (
                        type = self.compiler_type.capitalize(),
                        filename = filename,
                        name = _name,
                        subtitle = " ({subtitle})".format(**_json) if ('subtitle' in _json) else '',
                    )
                    print(message)
                    if (yn_input('Is it what you need?')):
                        return filename

            except (OSError, FileNotFoundError, json.JSONDecodeError):
                pass

        return None

    def open(self, name):
        filename = self.find(name)
        if (not filename):
            logger.error("Cannot find {type}: '{name}'".format(type=self.compiler_type, name=name))
            return None
        logger.info("Reading '{filename}'...".format(filename=filename))
        file = open(filename)
        raw = json.load(file)
        file.close()

        filename = self.schema_path
        logger.debug("Reading '{filename}'...".format(filename=filename))
        file = open(filename)
        schema = json.load(file)
        file.close()

        try:
            logger.debug("Validating {name}...".format(name=name))
            validate(raw, schema)
        except jsonschema.exceptions.ValidationError as e:
            logger.error("Raw file is not valid: {msg}".format(msg=e.message))
            self.raw = None
        else:
            logger.debug("Raw file is valid")
            self.raw = raw

        return self.raw

    def compile(self):
        raise NotImplementedError

    def save(self, forced_width=None):
        if (not self.compiled):
            logger.error("Could not save not compiled result")
            return None

        name = self.output_name or self.raw.get('name', None)
        if (not name):
            logger.error("Could find proper name")
            return None

        _image = self.compiled
        if (forced_width):
            _estimated_width = forced_width

            _actual_width = _image.width
            _actual_height = _image.height
            _image = _image.resize((int(_estimated_width), int(_estimated_width / _actual_width * _actual_height)), Image.ANTIALIAS)

        filename = "{directory}{name}.png".format(directory=self.output_directory, name=name)
        try:
            logger.info("Saving compiled file: '{filename}'".format(filename=filename))
            _image.save(filename)
        except:
            logger.exception("Error while saving file '{filename}'".format(filename=filename))
            return None
        else:
            return filename

    def insert_table \
                    (
                    self, vertical_columns, top, cell_height, data,
                    body_row_template, body_text_drawer, body_row_interval, body_capitalization=None, body_bold=None, body_italic=None,
                    header_row=None, header_text_drawer=None, header_row_interval=None, header_capitalization=None,header_bold=None, header_italic=None,
            ):
        if (top < 0):
            direction = -1
            indirect = True
        else:
            direction = 1
            indirect = False

        total_height = 0
        y1 = top
        y2 = cell_height and (top + cell_height)

        if (header_row and not indirect):
            y1, y2, total_height = self.__print_header_row \
                    (
                    vertical_columns=vertical_columns, y1=y1, y2=y2, total_height=total_height,
                    body_row_template=body_row_template, body_text_drawer=body_text_drawer, body_row_interval=body_row_interval, body_capitalization=body_capitalization, body_bold=body_bold, body_italic=body_italic,
                    header_row=header_row, header_text_drawer=header_text_drawer, header_row_interval=header_row_interval, header_capitalization=header_capitalization, header_bold=header_bold, header_italic=header_italic,
                )

        text_drawer = body_text_drawer
        _font = text_drawer.get_font()
        text_drawer.set_font(capitalization=body_capitalization, bold=body_bold, italic=body_italic)

        dy = body_row_interval
        _data = reversed(data) if (indirect) else data
        for data_row in _data:
            table_row = body_row_template
            _h = self.__print_table_row(y1=y1, y2=y2, text_drawer=text_drawer, vertical_columns=vertical_columns, table_row=table_row, data_row=data_row)
            y1 += _h + dy
            total_height += _h + dy
            if (y2):
                y2 += _h + dy

        logger.debug("Restoring font {font}".format(font=_font))
        text_drawer.set_font(**_font)

        if (header_row and indirect):
            y1, y2, total_height = self.__print_header_row \
                    (
                    vertical_columns=vertical_columns, y1=y1, y2=y2, total_height=total_height,
                    body_row_template=body_row_template, body_text_drawer=body_text_drawer, body_row_interval=body_row_interval, body_capitalization=body_capitalization, body_bold=body_bold, body_italic=body_italic,
                    header_row=header_row, header_text_drawer=header_text_drawer, header_row_interval=header_row_interval, header_capitalization=header_capitalization, header_bold=header_bold, header_italic=header_italic,
                )

        return total_height

    def __print_header_row \
                    (
                    self, vertical_columns, y1, y2, total_height,
                    body_row_template, body_text_drawer, body_row_interval, body_capitalization=None, body_bold=None, body_italic=None,
                    header_row=None, header_text_drawer=None, header_row_interval=None, header_capitalization=None,header_bold=None, header_italic=None,
            ):
        text_drawer = header_text_drawer or body_text_drawer

        _font = text_drawer.get_font()
        text_drawer.set_font(capitalization=header_capitalization, bold=header_bold, italic=header_italic)

        dy = header_row_interval or body_row_interval
        data_row = { }
        table_row = header_row
        _h = self.__print_table_row(y1=y1, y2=y2, text_drawer=text_drawer, vertical_columns=vertical_columns, table_row=table_row, data_row=data_row)
        y1 += _h + dy
        total_height += _h + dy
        if (y2):
            y2 += _h + dy

        logger.debug("Restoring font {font}".format(font=_font))
        text_drawer.set_font(**_font)
        return y1, y2, total_height

    def __print_table_row(self, y1, y2, text_drawer, vertical_columns, table_row, data_row):
        max_h = 0
        for i, _cell in enumerate(table_row):
            x1 = vertical_columns[i]
            x2 = vertical_columns[i + 1]
            _, _h = text_drawer.get_text_size((x1, y1, x2, y2), _cell.format(**data_row), offset_borders=False)
            if (max_h < _h):
                max_h = _h

        _y1 = -y1 - max_h if (y1 < 0) else y1
        _y2 = (-y2 - max_h if (y2 < 0) else y2) if y2 else _y1 + max_h

        for i, _cell in enumerate(table_row):
            x1 = vertical_columns[i]
            x2 = vertical_columns[i + 1]

            text_drawer.print_in_region((x1, _y1, x2, _y2), _cell.format(**data_row), offset_borders=False)

        return max_h

    def insert_image_scaled(self, base_image, region, image_path, offset_borders=True):
        if (offset_borders):
            x, y, w, h = region
        else:
            x, y, x2, y2 = region
            w = x2 - x
            h = y2 - y

        logger.debug("Inserting image '{path}' to position {region} with scaling and reversed mask".format(path=image_path, region=region))
        image = Image.open(image_path)

        _estimated_width = w
        _estimated_height = h

        _actual_width = image.width
        _actual_height = image.height

        _width_scale = _estimated_width / _actual_width
        _height_scale = _estimated_height / _actual_height

        _new_scale = max(_width_scale, _height_scale)
        _image = image.resize((int(_new_scale * _actual_width), int(_new_scale * _actual_height)), Image.ANTIALIAS)

        result_image = base_image.copy()
        result_image.paste(_image, (x, y))
        result_image.paste(base_image, (0, 0), base_image)
        return result_image

    def get_image_size(self, image_path):
        image = Image.open(image_path)
        _w = image.width
        _h = image.height
        return _w, _h

    def insert_image_centered(self, base_image, position, image_path, offset_borders=True):
        x, y = position

        logger.debug("Inserting image '{path}' to position {position}".format(path=image_path, position=position))
        image = Image.open(image_path)
        _w = image.width
        _h = image.height

        _x = x - _w // 2; _y = y - _h // 2
        base_image.paste(image, (_x, _y), image)
        return _x, _y, _w, _h
<commit_msg>Fix: small typo in raw finder<commit_after>import json, os.path, glob
import jsonschema.exceptions
from jsonschema import validate
from PIL import Image

from shared_funcs import yn_input
from hammerhal import ConfigLoader
from logging import getLogger
logger = getLogger('hammerhal.compilers.compiler_base')

class CompilerException(Exception):
    pass

class CompilerBase():

    compiler_type = None
    schema_path = None
    raw_directory = None
    sources_directory = None
    output_directory = None
    output_name = None

    raw = None
    compiled = None

    def __init__(self):
        self.schema_path = "{directory}{type}.json".format(directory=ConfigLoader.get_from_config('schemasDirectory', 'compilers'), type=self.compiler_type)
        self.raw_directory = "{rawRoot}{rawOffset}".format(rawRoot=ConfigLoader.get_from_config('rawDirectoryRoot'), rawOffset=ConfigLoader.get_from_config('compilerTypeSpecific/{type}/rawDirectory'.format(type=self.compiler_type), 'compilers'))
        self.output_directory = "{rawRoot}{rawOffset}".format(rawRoot=ConfigLoader.get_from_config('outputDirectoryRoot', 'compilers'), rawOffset=ConfigLoader.get_from_config('compilerTypeSpecific/{type}/outputDirectory'.format(type=self.compiler_type), 'compilers'))
        self.sources_directory = ConfigLoader.get_from_config('sourcesDirectory', 'compilers')

    def search(self):
        return glob.glob(self.raw_directory + "*.json")

    def find(self, name):
        filename = name
        if (os.path.isfile(filename)):
            return filename

        filename = "{directory}{name}.json".format(directory=self.raw_directory, name=name)
        if (os.path.isfile(filename)):
            return filename

        for filename in self.search():
            try:
                _file = open(filename)
                _json = json.load(_file)
                _file.close()

                _name = _json.get('name', None)
                if (_name and isinstance(_name, str) and name.lower() in _name.lower()):
                    message = "{type} found - {filename}:\n  {name}{subtitle}".format \
                    (
                        type = self.compiler_type.capitalize(),
                        filename = filename,
                        name = _name,
                        subtitle = " ({subtitle})".format(**_json) if ('subtitle' in _json) else '',
                    )
                    print(message)
                    if (yn_input('Is it what you need?')):
                        return filename

            except (OSError, FileNotFoundError, json.JSONDecodeError):
                pass

        return None

    def open(self, name):
        filename = self.find(name)
        if (not filename):
            logger.error("Cannot find {type}: '{name}'".format(type=self.compiler_type, name=name))
            return None
        logger.info("Reading '{filename}'...".format(filename=filename))
        file = open(filename)
        raw = json.load(file)
        file.close()

        filename = self.schema_path
        logger.debug("Reading '{filename}'...".format(filename=filename))
        file = open(filename)
        schema = json.load(file)
        file.close()

        try:
            logger.debug("Validating {name}...".format(name=name))
            validate(raw, schema)
        except jsonschema.exceptions.ValidationError as e:
            logger.error("Raw file is not valid: {msg}".format(msg=e.message))
            self.raw = None
        else:
            logger.debug("Raw file is valid")
            self.raw = raw

        return self.raw

    def compile(self):
        raise NotImplementedError

    def save(self, forced_width=None):
        if (not self.compiled):
            logger.error("Could not save not compiled result")
            return None

        name = self.output_name or self.raw.get('name', None)
        if (not name):
            logger.error("Could find proper name")
            return None

        _image = self.compiled
        if (forced_width):
            _estimated_width = forced_width

            _actual_width = _image.width
            _actual_height = _image.height
            _image = _image.resize((int(_estimated_width), int(_estimated_width / _actual_width * _actual_height)), Image.ANTIALIAS)

        filename = "{directory}{name}.png".format(directory=self.output_directory, name=name)
        try:
            logger.info("Saving compiled file: '{filename}'".format(filename=filename))
            _image.save(filename)
        except:
            logger.exception("Error while saving file '{filename}'".format(filename=filename))
            return None
        else:
            return filename

    def insert_table \
                    (
                    self, vertical_columns, top, cell_height, data,
                    body_row_template, body_text_drawer, body_row_interval, body_capitalization=None, body_bold=None, body_italic=None,
                    header_row=None, header_text_drawer=None, header_row_interval=None, header_capitalization=None,header_bold=None, header_italic=None,
            ):
        if (top < 0):
            direction = -1
            indirect = True
        else:
            direction = 1
            indirect = False

        total_height = 0
        y1 = top
        y2 = cell_height and (top + cell_height)

        if (header_row and not indirect):
            y1, y2, total_height = self.__print_header_row \
                    (
                    vertical_columns=vertical_columns, y1=y1, y2=y2, total_height=total_height,
                    body_row_template=body_row_template, body_text_drawer=body_text_drawer, body_row_interval=body_row_interval, body_capitalization=body_capitalization, body_bold=body_bold, body_italic=body_italic,
                    header_row=header_row, header_text_drawer=header_text_drawer, header_row_interval=header_row_interval, header_capitalization=header_capitalization, header_bold=header_bold, header_italic=header_italic,
                )

        text_drawer = body_text_drawer
        _font = text_drawer.get_font()
        text_drawer.set_font(capitalization=body_capitalization, bold=body_bold, italic=body_italic)

        dy = body_row_interval
        _data = reversed(data) if (indirect) else data
        for data_row in _data:
            table_row = body_row_template
            _h = self.__print_table_row(y1=y1, y2=y2, text_drawer=text_drawer, vertical_columns=vertical_columns, table_row=table_row, data_row=data_row)
            y1 += _h + dy
            total_height += _h + dy
            if (y2):
                y2 += _h + dy

        logger.debug("Restoring font {font}".format(font=_font))
        text_drawer.set_font(**_font)

        if (header_row and indirect):
            y1, y2, total_height = self.__print_header_row \
                    (
                    vertical_columns=vertical_columns, y1=y1, y2=y2, total_height=total_height,
                    body_row_template=body_row_template, body_text_drawer=body_text_drawer, body_row_interval=body_row_interval, body_capitalization=body_capitalization, body_bold=body_bold, body_italic=body_italic,
                    header_row=header_row, header_text_drawer=header_text_drawer, header_row_interval=header_row_interval, header_capitalization=header_capitalization, header_bold=header_bold, header_italic=header_italic,
                )

        return total_height

    def __print_header_row \
                    (
                    self, vertical_columns, y1, y2, total_height,
                    body_row_template, body_text_drawer, body_row_interval, body_capitalization=None, body_bold=None, body_italic=None,
                    header_row=None, header_text_drawer=None, header_row_interval=None, header_capitalization=None,header_bold=None, header_italic=None,
            ):
        text_drawer = header_text_drawer or body_text_drawer

        _font = text_drawer.get_font()
        text_drawer.set_font(capitalization=header_capitalization, bold=header_bold, italic=header_italic)

        dy = header_row_interval or body_row_interval
        data_row = { }
        table_row = header_row
        _h = self.__print_table_row(y1=y1, y2=y2, text_drawer=text_drawer, vertical_columns=vertical_columns, table_row=table_row, data_row=data_row)
        y1 += _h + dy
        total_height += _h + dy
        if (y2):
            y2 += _h + dy

        logger.debug("Restoring font {font}".format(font=_font))
        text_drawer.set_font(**_font)
        return y1, y2, total_height

    def __print_table_row(self, y1, y2, text_drawer, vertical_columns, table_row, data_row):
        max_h = 0
        for i, _cell in enumerate(table_row):
            x1 = vertical_columns[i]
            x2 = vertical_columns[i + 1]
            _, _h = text_drawer.get_text_size((x1, y1, x2, y2), _cell.format(**data_row), offset_borders=False)
            if (max_h < _h):
                max_h = _h

        _y1 = -y1 - max_h if (y1 < 0) else y1
        _y2 = (-y2 - max_h if (y2 < 0) else y2) if y2 else _y1 + max_h

        for i, _cell in enumerate(table_row):
            x1 = vertical_columns[i]
            x2 = vertical_columns[i + 1]

            text_drawer.print_in_region((x1, _y1, x2, _y2), _cell.format(**data_row), offset_borders=False)

        return max_h

    def insert_image_scaled(self, base_image, region, image_path, offset_borders=True):
        if (offset_borders):
            x, y, w, h = region
        else:
            x, y, x2, y2 = region
            w = x2 - x
            h = y2 - y

        logger.debug("Inserting image '{path}' to position {region} with scaling and reversed mask".format(path=image_path, region=region))
        image = Image.open(image_path)

        _estimated_width = w
        _estimated_height = h

        _actual_width = image.width
        _actual_height = image.height

        _width_scale = _estimated_width / _actual_width
        _height_scale = _estimated_height / _actual_height

        _new_scale = max(_width_scale, _height_scale)
        _image = image.resize((int(_new_scale * _actual_width), int(_new_scale * _actual_height)), Image.ANTIALIAS)

        result_image = base_image.copy()
        result_image.paste(_image, (x, y))
        result_image.paste(base_image, (0, 0), base_image)
        return result_image

    def get_image_size(self, image_path):
        image = Image.open(image_path)
        _w = image.width
        _h = image.height
        return _w, _h

    def insert_image_centered(self, base_image, position, image_path, offset_borders=True):
        x, y = position

        logger.debug("Inserting image '{path}' to position {position}".format(path=image_path, position=position))
        image = Image.open(image_path)
        _w = image.width
        _h = image.height

        _x = x - _w // 2; _y = y - _h // 2
        base_image.paste(image, (_x, _y), image)
        return _x, _y, _w, _h
<|endoftext|>


<commit_before># Copyright (C) 2014-2015 New York University
# This file is part of ReproZip which is released under the Revised BSD License
# See file LICENSE for full license details.

"""Miscellaneous utilities for unpacker plugins.
"""

from __future__ import unicode_literals

import functools
import logging
import os
import random
from rpaths import PosixPath, Path
import signal
import string
import subprocess
import sys
import tarfile

import reprounzip.common
from reprounzip.utils import irange

COMPAT_OK = 0
COMPAT_NO = 1
COMPAT_MAYBE = 2

class UsageError(Exception):
    def __init__(self, msg="Invalid command-line"):
        Exception.__init__(self, msg)

def composite_action(*functions):
    """Makes an action that just calls several other actions in sequence.

    Useful to implement ``myplugin setup`` in terms of ``myplugin setup/part1``
    and ``myplugin setup/part2``: simply use
    ``act1n2 = composite_action(act1, act2)``.
    """
    def wrapper(args):
        for function in functions:
            function(args)
    return wrapper

def target_must_exist(func):
    """Decorator that checks that ``args.target`` exists.
    """
    @functools.wraps(func)
    def wrapper(args):
        target = Path(args.target[0])
        if not target.is_dir():
            logging.critical("Error: Target directory doesn't exist")
            raise UsageError
        return func(args)
    return wrapper

def unique_names():
    """Generates unique sequences of bytes.
    """
    characters = (b"abcdefghijklmnopqrstuvwxyz"
                  b"0123456789")
    characters = [characters[i:i + 1] for i in irange(len(characters))]
    rng = random.Random()
    while True:
        letters = [rng.choice(characters) for i in irange(10)]
        yield b''.join(letters)
unique_names = unique_names()

def make_unique_name(prefix):
    """Makes a unique (random) bytestring name, starting with the given prefix.
    """
    assert isinstance(prefix, bytes)
    return prefix + next(unique_names)

safe_shell_chars = set("ABCDEFGHIJKLMNOPQRSTUVWXYZ"
                       "abcdefghijklmnopqrstuvwxyz"
                       "0123456789"
                       "-+=/:.,%_")

def shell_escape(s):
    """Given bl"a, returns "bl\\"a".
    """
    if isinstance(s, bytes):
        s = s.decode('utf-8')
    if any(c not in safe_shell_chars for c in s):
        return '"%s"' % (s.replace('\\', '\\\\')
                          .replace('"', '\\"')
                          .replace('$', '\\$'))
    else:
        return s

def load_config(pack):
    """Utility method loading the YAML configuration from inside a pack file.

    Decompresses the config.yml file from the tarball to a temporary file then
    loads it. Note that decompressing a single file is inefficient, thus
    calling this method can be slow.
    """
    tmp = Path.tempdir(prefix='reprozip_')
    try:
        # Loads info from package
        tar = tarfile.open(str(pack), 'r:*')
        f = tar.extractfile('METADATA/version')
        version = f.read()
        f.close()
        if version != b'REPROZIP VERSION 1\n':
            logging.critical("Unknown pack format")
            sys.exit(1)
        tar.extract('METADATA/config.yml', path=str(tmp))
        tar.close()
        configfile = tmp / 'METADATA/config.yml'
        ret = reprounzip.common.load_config(configfile, canonical=True)
    finally:
        tmp.rmtree()

    return ret

def busybox_url(arch):
    """Gets the correct URL for the busybox binary given the architecture.
    """
    return 'http://www.busybox.net/downloads/binaries/latest/busybox-%s' % arch

def join_root(root, path):
    """Prepends `root` to the absolute path `path`.
    """
    p_root, p_loc = path.split_root()
    assert p_root == b'/'
    return root / p_loc

class FileUploader(object):
    """Common logic for 'upload' commands.
    """
    def __init__(self, target, input_files, files):
        self.target = target
        self.input_files = input_files
        self.run(files)

    def run(self, files):
        reprounzip.common.record_usage(upload_files=len(files))
        runs = self.get_runs_from_config()

        # No argument: list all the input files and exit
        if not files:
            print("Input files:")
            for i, run in enumerate(runs):
                if len(runs) > 1:
                    print("  Run %d:" % i)
                for input_name in run['input_files']:
                    if self.input_files.get(input_name) is not None:
                        assigned = PosixPath(self.input_files[input_name])
                    else:
                        assigned = "(original)"
                    print("    %s: %s" % (input_name, assigned))
            return

        self.prepare_upload(files)

        # Get the path of each input file
        all_input_files = {}
        for run in runs:
            all_input_files.update(run['input_files'])

        try:
            # Upload files
            for filespec in files:
                filespec_split = filespec.rsplit(':', 1)
                if len(filespec_split) != 2:
                    logging.critical("Invalid file specification: %r",
                                     filespec)
                    sys.exit(1)
                local_path, input_name = filespec_split

                try:
                    input_path = PosixPath(all_input_files[input_name])
                except KeyError:
                    logging.critical("Invalid input file: %r", input_name)
                    sys.exit(1)

                temp = None

                if not local_path:
                    # Restore original file from pack
                    logging.debug("Restoring input file %s", input_path)
                    fd, temp = Path.tempfile(prefix='reprozip_input_')
                    os.close(fd)
                    local_path = self.extract_original_input(input_name,
                                                             input_path,
                                                             temp)
                else:
                    local_path = Path(local_path)
                    logging.debug("Uploading file %s to %s",
                                  local_path, input_path)
                    if not local_path.exists():
                        logging.critical("Local file %s doesn't exist",
                                         local_path)
                        sys.exit(1)

                self.upload_file(local_path, input_path)

                if temp is not None:
                    temp.remove()
                    self.input_files.pop(input_name, None)
                else:
                    self.input_files[input_name] = local_path.absolute().path
        finally:
            self.finalize()

    def get_runs_from_config(self):
        # Loads config
        runs, packages, other_files = reprounzip.common.load_config(
                self.target / 'config.yml',
                canonical=True)
        return runs

    def prepare_upload(self, files):
        pass

    def extract_original_input(self, input_name, input_path, temp):
        tar = tarfile.open(str(self.target / 'experiment.rpz'), 'r:*')
        member = tar.getmember(str(join_root(PosixPath('DATA'), input_path)))
        member.name = str(temp.name)
        tar.extract(member, str(temp.parent))
        tar.close()
        return temp

    def upload_file(self, local_path, input_path):
        raise NotImplementedError

    def finalize(self):
        pass

class FileDownloader(object):
    """Common logic for 'download' commands.
    """
    def __init__(self, target, files):
        self.target = target
        self.run(files)

    def run(self, files):
        reprounzip.common.record_usage(download_files=len(files))
        runs = self.get_runs_from_config()

        # No argument: list all the output files and exit
        if not files:
            print("Output files:")
            for i, run in enumerate(runs):
                if len(runs) > 1:
                    print("  Run %d:" % i)
                for output_name in run['output_files']:
                    print("    %s" % output_name)
            return

        self.prepare_download(files)

        # Get the path of each output file
        all_output_files = {}
        for run in runs:
            all_output_files.update(run['output_files'])

        try:
            # Download files
            for filespec in files:
                filespec_split = filespec.split(':', 1)
                if len(filespec_split) != 2:
                    logging.critical("Invalid file specification: %r",
                                     filespec)
                    sys.exit(1)
                output_name, local_path = filespec_split

                try:
                    remote_path = PosixPath(all_output_files[output_name])
                except KeyError:
                    logging.critical("Invalid output file: %r", output_name)
                    sys.exit(1)

                logging.debug("Downloading file %s", remote_path)
                if not local_path:
                    self.download_and_print(remote_path)
                else:
                    self.download(remote_path, Path(local_path))
        finally:
            self.finalize()

    def get_runs_from_config(self):
        # Loads config
        runs, packages, other_files = reprounzip.common.load_config(
                self.target / 'config.yml',
                canonical=True)
        return runs

    def prepare_download(self, files):
        pass

    def download_and_print(self, remote_path):
        # Download to temporary file
        fd, temp = Path.tempfile(prefix='reprozip_output_')
        os.close(fd)
        self.download(remote_path, temp)
        # Output to stdout
        with temp.open('rb') as fp:
            chunk = fp.read(1024)
            if chunk:
                sys.stdout.buffer.write(chunk)
            while len(chunk) == 1024:
                chunk = fp.read(1024)
                if chunk:
                    sys.stdout.buffer.write(chunk)
        temp.remove()

    def download(self, remote_path, local_path):
        raise NotImplementedError

    def finalize(self):
        pass

def get_runs(runs, selected_run, cmdline):
    """Selects which run(s) to execute based on parts of the command-line.

    Will return an iterable of run numbers. Might also fail loudly or exit
    after printing the original command-line.
    """
    if selected_run is None and len(runs) == 1:
        selected_run = 0
    elif selected_run is not None:
        try:
            selected_run = int(selected_run)
        except ValueError:
            logging.critical("Error: Run is not a number")
            raise UsageError
        if selected_run < 0 or selected_run >= len(runs):
            logging.critical("Error: Expected 0 <= run <= %d, got %d",
                             len(runs) - 1, selected_run)
            sys.exit(1)

    # --cmdline without arguments: display the original command-line
    if cmdline == []:
        if selected_run is None:
            logging.critical("There are several runs in this pack -- you have "
                             "to choose which one to use with --cmdline")
            sys.exit(1)
        print("Original command-line:")
        print(' '.join(shell_escape(arg)
                       for arg in runs[selected_run]['argv']))
        sys.exit(0)

    if selected_run is None:
        selected_run = irange(len(runs))
    else:
        selected_run = (int(selected_run),)

    return selected_run

def interruptible_call(*args, **kwargs):
    assert signal.getsignal(signal.SIGINT) == signal.default_int_handler
    proc = [None]

    def _sigint_handler(signum, frame):
        if proc[0] is not None:
            try:
                proc[0].send_signal(signum)
            except OSError:
                pass

    signal.signal(signal.SIGINT, _sigint_handler)

    try:
        proc[0] = subprocess.Popen(*args, **kwargs)
        return proc[0].wait()
    finally:
        signal.signal(signal.SIGINT, signal.default_int_handler)
<commit_msg>Removes unused import<commit_after># Copyright (C) 2014-2015 New York University
# This file is part of ReproZip which is released under the Revised BSD License
# See file LICENSE for full license details.

"""Miscellaneous utilities for unpacker plugins.
"""

from __future__ import unicode_literals

import functools
import logging
import os
import random
from rpaths import PosixPath, Path
import signal
import subprocess
import sys
import tarfile

import reprounzip.common
from reprounzip.utils import irange

COMPAT_OK = 0
COMPAT_NO = 1
COMPAT_MAYBE = 2

class UsageError(Exception):
    def __init__(self, msg="Invalid command-line"):
        Exception.__init__(self, msg)

def composite_action(*functions):
    """Makes an action that just calls several other actions in sequence.

    Useful to implement ``myplugin setup`` in terms of ``myplugin setup/part1``
    and ``myplugin setup/part2``: simply use
    ``act1n2 = composite_action(act1, act2)``.
    """
    def wrapper(args):
        for function in functions:
            function(args)
    return wrapper

def target_must_exist(func):
    """Decorator that checks that ``args.target`` exists.
    """
    @functools.wraps(func)
    def wrapper(args):
        target = Path(args.target[0])
        if not target.is_dir():
            logging.critical("Error: Target directory doesn't exist")
            raise UsageError
        return func(args)
    return wrapper

def unique_names():
    """Generates unique sequences of bytes.
    """
    characters = (b"abcdefghijklmnopqrstuvwxyz"
                  b"0123456789")
    characters = [characters[i:i + 1] for i in irange(len(characters))]
    rng = random.Random()
    while True:
        letters = [rng.choice(characters) for i in irange(10)]
        yield b''.join(letters)
unique_names = unique_names()

def make_unique_name(prefix):
    """Makes a unique (random) bytestring name, starting with the given prefix.
    """
    assert isinstance(prefix, bytes)
    return prefix + next(unique_names)

safe_shell_chars = set("ABCDEFGHIJKLMNOPQRSTUVWXYZ"
                       "abcdefghijklmnopqrstuvwxyz"
                       "0123456789"
                       "-+=/:.,%_")

def shell_escape(s):
    """Given bl"a, returns "bl\\"a".
    """
    if isinstance(s, bytes):
        s = s.decode('utf-8')
    if any(c not in safe_shell_chars for c in s):
        return '"%s"' % (s.replace('\\', '\\\\')
                          .replace('"', '\\"')
                          .replace('$', '\\$'))
    else:
        return s

def load_config(pack):
    """Utility method loading the YAML configuration from inside a pack file.

    Decompresses the config.yml file from the tarball to a temporary file then
    loads it. Note that decompressing a single file is inefficient, thus
    calling this method can be slow.
    """
    tmp = Path.tempdir(prefix='reprozip_')
    try:
        # Loads info from package
        tar = tarfile.open(str(pack), 'r:*')
        f = tar.extractfile('METADATA/version')
        version = f.read()
        f.close()
        if version != b'REPROZIP VERSION 1\n':
            logging.critical("Unknown pack format")
            sys.exit(1)
        tar.extract('METADATA/config.yml', path=str(tmp))
        tar.close()
        configfile = tmp / 'METADATA/config.yml'
        ret = reprounzip.common.load_config(configfile, canonical=True)
    finally:
        tmp.rmtree()

    return ret

def busybox_url(arch):
    """Gets the correct URL for the busybox binary given the architecture.
    """
    return 'http://www.busybox.net/downloads/binaries/latest/busybox-%s' % arch

def join_root(root, path):
    """Prepends `root` to the absolute path `path`.
    """
    p_root, p_loc = path.split_root()
    assert p_root == b'/'
    return root / p_loc

class FileUploader(object):
    """Common logic for 'upload' commands.
    """
    def __init__(self, target, input_files, files):
        self.target = target
        self.input_files = input_files
        self.run(files)

    def run(self, files):
        reprounzip.common.record_usage(upload_files=len(files))
        runs = self.get_runs_from_config()

        # No argument: list all the input files and exit
        if not files:
            print("Input files:")
            for i, run in enumerate(runs):
                if len(runs) > 1:
                    print("  Run %d:" % i)
                for input_name in run['input_files']:
                    if self.input_files.get(input_name) is not None:
                        assigned = PosixPath(self.input_files[input_name])
                    else:
                        assigned = "(original)"
                    print("    %s: %s" % (input_name, assigned))
            return

        self.prepare_upload(files)

        # Get the path of each input file
        all_input_files = {}
        for run in runs:
            all_input_files.update(run['input_files'])

        try:
            # Upload files
            for filespec in files:
                filespec_split = filespec.rsplit(':', 1)
                if len(filespec_split) != 2:
                    logging.critical("Invalid file specification: %r",
                                     filespec)
                    sys.exit(1)
                local_path, input_name = filespec_split

                try:
                    input_path = PosixPath(all_input_files[input_name])
                except KeyError:
                    logging.critical("Invalid input file: %r", input_name)
                    sys.exit(1)

                temp = None

                if not local_path:
                    # Restore original file from pack
                    logging.debug("Restoring input file %s", input_path)
                    fd, temp = Path.tempfile(prefix='reprozip_input_')
                    os.close(fd)
                    local_path = self.extract_original_input(input_name,
                                                             input_path,
                                                             temp)
                else:
                    local_path = Path(local_path)
                    logging.debug("Uploading file %s to %s",
                                  local_path, input_path)
                    if not local_path.exists():
                        logging.critical("Local file %s doesn't exist",
                                         local_path)
                        sys.exit(1)

                self.upload_file(local_path, input_path)

                if temp is not None:
                    temp.remove()
                    self.input_files.pop(input_name, None)
                else:
                    self.input_files[input_name] = local_path.absolute().path
        finally:
            self.finalize()

    def get_runs_from_config(self):
        # Loads config
        runs, packages, other_files = reprounzip.common.load_config(
                self.target / 'config.yml',
                canonical=True)
        return runs

    def prepare_upload(self, files):
        pass

    def extract_original_input(self, input_name, input_path, temp):
        tar = tarfile.open(str(self.target / 'experiment.rpz'), 'r:*')
        member = tar.getmember(str(join_root(PosixPath('DATA'), input_path)))
        member.name = str(temp.name)
        tar.extract(member, str(temp.parent))
        tar.close()
        return temp

    def upload_file(self, local_path, input_path):
        raise NotImplementedError

    def finalize(self):
        pass

class FileDownloader(object):
    """Common logic for 'download' commands.
    """
    def __init__(self, target, files):
        self.target = target
        self.run(files)

    def run(self, files):
        reprounzip.common.record_usage(download_files=len(files))
        runs = self.get_runs_from_config()

        # No argument: list all the output files and exit
        if not files:
            print("Output files:")
            for i, run in enumerate(runs):
                if len(runs) > 1:
                    print("  Run %d:" % i)
                for output_name in run['output_files']:
                    print("    %s" % output_name)
            return

        self.prepare_download(files)

        # Get the path of each output file
        all_output_files = {}
        for run in runs:
            all_output_files.update(run['output_files'])

        try:
            # Download files
            for filespec in files:
                filespec_split = filespec.split(':', 1)
                if len(filespec_split) != 2:
                    logging.critical("Invalid file specification: %r",
                                     filespec)
                    sys.exit(1)
                output_name, local_path = filespec_split

                try:
                    remote_path = PosixPath(all_output_files[output_name])
                except KeyError:
                    logging.critical("Invalid output file: %r", output_name)
                    sys.exit(1)

                logging.debug("Downloading file %s", remote_path)
                if not local_path:
                    self.download_and_print(remote_path)
                else:
                    self.download(remote_path, Path(local_path))
        finally:
            self.finalize()

    def get_runs_from_config(self):
        # Loads config
        runs, packages, other_files = reprounzip.common.load_config(
                self.target / 'config.yml',
                canonical=True)
        return runs

    def prepare_download(self, files):
        pass

    def download_and_print(self, remote_path):
        # Download to temporary file
        fd, temp = Path.tempfile(prefix='reprozip_output_')
        os.close(fd)
        self.download(remote_path, temp)
        # Output to stdout
        with temp.open('rb') as fp:
            chunk = fp.read(1024)
            if chunk:
                sys.stdout.buffer.write(chunk)
            while len(chunk) == 1024:
                chunk = fp.read(1024)
                if chunk:
                    sys.stdout.buffer.write(chunk)
        temp.remove()

    def download(self, remote_path, local_path):
        raise NotImplementedError

    def finalize(self):
        pass

def get_runs(runs, selected_run, cmdline):
    """Selects which run(s) to execute based on parts of the command-line.

    Will return an iterable of run numbers. Might also fail loudly or exit
    after printing the original command-line.
    """
    if selected_run is None and len(runs) == 1:
        selected_run = 0
    elif selected_run is not None:
        try:
            selected_run = int(selected_run)
        except ValueError:
            logging.critical("Error: Run is not a number")
            raise UsageError
        if selected_run < 0 or selected_run >= len(runs):
            logging.critical("Error: Expected 0 <= run <= %d, got %d",
                             len(runs) - 1, selected_run)
            sys.exit(1)

    # --cmdline without arguments: display the original command-line
    if cmdline == []:
        if selected_run is None:
            logging.critical("There are several runs in this pack -- you have "
                             "to choose which one to use with --cmdline")
            sys.exit(1)
        print("Original command-line:")
        print(' '.join(shell_escape(arg)
                       for arg in runs[selected_run]['argv']))
        sys.exit(0)

    if selected_run is None:
        selected_run = irange(len(runs))
    else:
        selected_run = (int(selected_run),)

    return selected_run

def interruptible_call(*args, **kwargs):
    assert signal.getsignal(signal.SIGINT) == signal.default_int_handler
    proc = [None]

    def _sigint_handler(signum, frame):
        if proc[0] is not None:
            try:
                proc[0].send_signal(signum)
            except OSError:
                pass

    signal.signal(signal.SIGINT, _sigint_handler)

    try:
        proc[0] = subprocess.Popen(*args, **kwargs)
        return proc[0].wait()
    finally:
        signal.signal(signal.SIGINT, signal.default_int_handler)
<|endoftext|>


<commit_before>"""Modifications to the flow table from the controller"""

# System imports
import enum

# Third-party imports

# Local source tree imports
from pyof.v0x01.common import flow_match
from pyof.v0x01.common import header as of_header
from pyof.v0x01.controller2switch import common
from pyof.v0x01.foundation import base
from pyof.v0x01.foundation import basic_types

# Enums

class FlowModCommand(enum.Enum):
    """
    List the possible commands for a flow.

    Enums:
        OFPFC_ADD           # New Flow
        OFPFC_MODIFY        # Modify all flows
        OFPFC_MODIFY_STRICT # Modify entry strictly matching wildcards
        OFPFC_DELETE        # Delete all matching flows
        OFPFC_DELETE_STRICT # Strictly match wildcards and priority

    """
    OFPFC_ADD = 1
    OFPFC_MODIFY = 2
    OFPFC_MODIFY_STRICT = 3
    OFPFC_DELETE = 4
    OFPFC_DELETE_STRICT = 5

class FlowModFlags(enum.Enum):
    """Types to be used in Flags field"""

    #: Send flow removed message when flow expires or is deleted
    OFPFF_SEND_FLOW_REM = 1 << 0
    #: Check for overlapping entries first
    OFPFF_CHECK_OVERLAP = 1 << 1
    #: Remark this is for emergency
    OFPFF_EMERG = 1 << 2

# Classes

class FlowMod(base.GenericMessage):
    """
    Modifies the flow table from the controller.

        :param xid:          xid to be used on the message header
        :param match:        Fields to match
        :param cookie:       Opaque controller-issued identifier
        :param command:      One of OFPFC_*
        :param idle_timeout: Idle time before discarding (seconds)
        :param hard_timeout: Max time before discarding (seconds)
        :param priority:     Priority level of flow entry
        :param buffer_idle:  Buffered packet to apply to (or -1).
                             Not meaningful for OFPFC_DELETE*
        :param out_port:     For OFPFC_DELETE* commands, require matching
                             entries to include this as an output port.
                             A value of OFPP_NONE indicates no restriction.
        :param flags:        One of OFPFF_*
        :param actions:      The action length is inferred from the length
                             field in the header
    """
    header = of_header.Header()
    match = flow_match.Match()
    cookie = basic_types.UBInt64()
    command = basic_types.UBInt16()
    idle_timeout = basic_types.UBInt16()
    hard_timeout = basic_types.UBInt16()
    priority = basic_types.UBInt16()
    buffer_id = basic_types.UBInt32()
    out_port = basic_types.UBInt16()
    flags = basic_types.UBInt16()
    actions = common.ListOfActions()

    def __init__(self, xid=None, match=None, cookie=None, command=None,
                 idle_timeout=None, hard_timeout=None, priority=None,
                 buffer_id=None, out_port=None, flags=None, actions=[]):
        self.header.message_type = of_header.Type.OFPT_FLOW_MOD
        self.header.length = 0
        self.header.xid = xid
        self.match = match
        self.cookie = cookie
        self.command = command
        self.idle_timeout = idle_timeout
        self.hard_timeout = hard_timeout
        self.priority = priority
        self.buffer_id = buffer_id
        self.out_port = out_port
        self.flags = flags
        self.actions = actions
<commit_msg>Fix flow_mod enum. It should start at 0.<commit_after>"""Modifications to the flow table from the controller"""

# System imports
import enum

# Third-party imports

# Local source tree imports
from pyof.v0x01.common import flow_match
from pyof.v0x01.common import header as of_header
from pyof.v0x01.controller2switch import common
from pyof.v0x01.foundation import base
from pyof.v0x01.foundation import basic_types

# Enums

class FlowModCommand(enum.Enum):
    """
    List the possible commands for a flow.

    Enums:
        OFPFC_ADD           # New Flow
        OFPFC_MODIFY        # Modify all flows
        OFPFC_MODIFY_STRICT # Modify entry strictly matching wildcards
        OFPFC_DELETE        # Delete all matching flows
        OFPFC_DELETE_STRICT # Strictly match wildcards and priority

    """
    OFPFC_ADD = 0
    OFPFC_MODIFY = 1
    OFPFC_MODIFY_STRICT = 2
    OFPFC_DELETE = 3
    OFPFC_DELETE_STRICT = 4

class FlowModFlags(enum.Enum):
    """Types to be used in Flags field"""

    #: Send flow removed message when flow expires or is deleted
    OFPFF_SEND_FLOW_REM = 1 << 0
    #: Check for overlapping entries first
    OFPFF_CHECK_OVERLAP = 1 << 1
    #: Remark this is for emergency
    OFPFF_EMERG = 1 << 2

# Classes

class FlowMod(base.GenericMessage):
    """
    Modifies the flow table from the controller.

        :param xid:          xid to be used on the message header
        :param match:        Fields to match
        :param cookie:       Opaque controller-issued identifier
        :param command:      One of OFPFC_*
        :param idle_timeout: Idle time before discarding (seconds)
        :param hard_timeout: Max time before discarding (seconds)
        :param priority:     Priority level of flow entry
        :param buffer_idle:  Buffered packet to apply to (or -1).
                             Not meaningful for OFPFC_DELETE*
        :param out_port:     For OFPFC_DELETE* commands, require matching
                             entries to include this as an output port.
                             A value of OFPP_NONE indicates no restriction.
        :param flags:        One of OFPFF_*
        :param actions:      The action length is inferred from the length
                             field in the header
    """
    header = of_header.Header()
    match = flow_match.Match()
    cookie = basic_types.UBInt64()
    command = basic_types.UBInt16()
    idle_timeout = basic_types.UBInt16()
    hard_timeout = basic_types.UBInt16()
    priority = basic_types.UBInt16()
    buffer_id = basic_types.UBInt32()
    out_port = basic_types.UBInt16()
    flags = basic_types.UBInt16()
    actions = common.ListOfActions()

    def __init__(self, xid=None, match=None, cookie=None, command=None,
                 idle_timeout=None, hard_timeout=None, priority=None,
                 buffer_id=None, out_port=None, flags=None, actions=[]):
        self.header.message_type = of_header.Type.OFPT_FLOW_MOD
        self.header.length = 0
        self.header.xid = xid
        self.match = match
        self.cookie = cookie
        self.command = command
        self.idle_timeout = idle_timeout
        self.hard_timeout = hard_timeout
        self.priority = priority
        self.buffer_id = buffer_id
        self.out_port = out_port
        self.flags = flags
        self.actions = actions
<|endoftext|>


<commit_before>import boto
import sure
from nose import tools
from dnsadmin53 import core
from moto import mock_route53

@mock_route53
def test_get_zone():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel.huit.harvard.edu.')

    r53 = core.Route53()
    zone = r53.get_zone('devel.huit.harvard.edu')
    assert zone.name == 'devel.huit.harvard.edu.'

@mock_route53
def test_list_zone():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel1.huit.harvard.edu.')
    conn.create_hosted_zone('devel2.huit.harvard.edu.')

    r53 = core.Route53()
    zones = r53.list_zones()
    len(zones).should.equal(2)

@mock_route53
def test_does_zone_exist():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel.huit.harvard.edu.')

    r53 = core.Route53()
    result = r53.does_zone_exist('devel.huit.harvard.edu')
    print result
    tools.assert_true(result)

@mock_route53
def test_get_zone_id():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel.huit.harvard.edu.')

    r53 = core.Route53()
    zone_id = r53.get_zone_id('devel.huit.harvard.edu')
    tools.assert_true(zone_id)
<commit_msg>flake8 syntax update<commit_after>import boto
import sure
from nose import tools
from dnsadmin53 import core
from moto import mock_route53

@mock_route53
def test_get_zone():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel.huit.harvard.edu.')

    r53 = core.Route53()
    zone = r53.get_zone('devel.huit.harvard.edu')
    assert zone.name == 'devel.huit.harvard.edu.'

@mock_route53
def test_list_zone():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel1.huit.harvard.edu.')
    conn.create_hosted_zone('devel2.huit.harvard.edu.')

    r53 = core.Route53()
    zones = r53.list_zones()
    len(zones).should.equal(2)

@mock_route53
def test_does_zone_exist():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel.huit.harvard.edu.')

    r53 = core.Route53()
    result = r53.does_zone_exist('devel.huit.harvard.edu')
    print result
    tools.assert_true(result)

@mock_route53
def test_get_zone_id():
    conn = boto.connect_route53()
    conn.create_hosted_zone('devel.huit.harvard.edu.')

    r53 = core.Route53()
    zone_id = r53.get_zone_id('devel.huit.harvard.edu')
    tools.assert_true(zone_id)
<|endoftext|>


<commit_before>"""
Tests for the abstraction from continuous dynamics to logic
"""
import logging
logging.basicConfig(level=logging.INFO)

import numpy as np

from tulip import abstract, hybrid
from tulip.polytope import box2poly

input_bound = 0.4

def subsys0():
    dom = box2poly([[0., 3.], [0., 2.]])

    A = np.eye(2)
    B = np.eye(2)

    U = box2poly([[0., 1.],
                  [0., 1.]])
    U.scale(input_bound)

    sys_dyn = hybrid.LtiSysDyn(A, B, Uset=U, domain=dom)

    return sys_dyn

def subsys1():
    dom = box2poly([[0., 3.], [0., 2.]])

    A = np.eye(2)
    B = np.eye(2)

    U = box2poly([[0., 0.],
                  [-1., 0.]])
    U.scale(input_bound)

    sys_dyn = hybrid.LtiSysDyn(A, B, Uset=U, domain=dom)

    return sys_dyn

def transition_directions_test():
    """
    unit test for correctness of abstracted transition directions, with:

      - uni-directional control authority
      - no disturbance
    """
    modes = []
    modes.append(('normal', 'fly'))
    modes.append(('refuel', 'fly'))
    env_modes, sys_modes = zip(*modes)

    cont_state_space = box2poly([[0., 3.], [0., 2.]])
    pwa_sys = dict()
    pwa_sys[('normal', 'fly')] = hybrid.PwaSysDyn(
        [subsys0()], cont_state_space
    )
    pwa_sys[('refuel', 'fly')] = hybrid.PwaSysDyn(
        [subsys1()], cont_state_space
    )

    switched_dynamics = hybrid.HybridSysDyn(
        disc_domain_size=(len(env_modes), len(sys_modes)),
        dynamics=pwa_sys,
        env_labels=env_modes,
        disc_sys_labels=sys_modes,
        cts_ss=cont_state_space
    )

    cont_props = {}
    cont_props['home'] = box2poly([[0., 1.], [0., 1.]])
    cont_props['lot'] = box2poly([[2., 3.], [1., 2.]])

    ppp = abstract.prop2part(cont_state_space, cont_props)
    ppp, new2old = abstract.part2convex(ppp)

    N = 8
    trans_len=1

    disc_params = {}
    for mode in modes:
        disc_params[mode] = {'N':N, 'trans_length':trans_len}

    swab = abstract.discretize_switched(
        ppp, switched_dynamics, disc_params,
        plot=True, show_ts=True, only_adjacent=False
    )

    ts = swab.modes[('normal', 'fly')].ts
    edges = {('s0', 's0'), ('s1', 's1'), ('s2', 's2'), ('s3', 's3'),
             ('s4', 's4'), ('s5', 's5'),
             ('s1', 's2'), ('s1', 's4'),
             ('s2', 's3'), ('s2', 's5'), ('s2', 's0'),
             ('s3', 's0'),
             ('s4', 's5'),
             ('s5', 's0')}
    assert(set(ts.edges() ) == edges)

    ts = swab.ts
    edges.remove(('s2', 's0'))
    assert(set(ts.edges() ) == edges)
    for i, j in edges:
        assert(ts[i][j][0]['env_actions'] == 'normal')
        assert(ts[i][j][0]['sys_actions'] == 'fly')
<commit_msg>add unit test to guard againt self-loops that don't exist<commit_after>"""
Tests for the abstraction from continuous dynamics to logic
"""
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

import numpy as np

from tulip import abstract, hybrid
from tulip.polytope import box2poly
from tulip import polytope as pc

input_bound = 0.4

def subsys0():
    dom = box2poly([[0., 3.], [0., 2.]])

    A = np.eye(2)
    B = np.eye(2)

    U = box2poly([[0., 1.],
                  [0., 1.]])
    U.scale(input_bound)

    sys_dyn = hybrid.LtiSysDyn(A, B, Uset=U, domain=dom)

    return sys_dyn

def subsys1():
    dom = box2poly([[0., 3.], [0., 2.]])

    A = np.eye(2)
    B = np.eye(2)

    U = box2poly([[0., 0.],
                  [-1., 0.]])
    U.scale(input_bound)

    sys_dyn = hybrid.LtiSysDyn(A, B, Uset=U, domain=dom)

    return sys_dyn

def transition_directions_test():
    """
    unit test for correctness of abstracted transition directions, with:

      - uni-directional control authority
      - no disturbance
    """
    modes = []
    modes.append(('normal', 'fly'))
    modes.append(('refuel', 'fly'))
    env_modes, sys_modes = zip(*modes)

    cont_state_space = box2poly([[0., 3.], [0., 2.]])
    pwa_sys = dict()
    pwa_sys[('normal', 'fly')] = hybrid.PwaSysDyn(
        [subsys0()], cont_state_space
    )
    pwa_sys[('refuel', 'fly')] = hybrid.PwaSysDyn(
        [subsys1()], cont_state_space
    )

    switched_dynamics = hybrid.HybridSysDyn(
        disc_domain_size=(len(env_modes), len(sys_modes)),
        dynamics=pwa_sys,
        env_labels=env_modes,
        disc_sys_labels=sys_modes,
        cts_ss=cont_state_space
    )

    cont_props = {}
    cont_props['home'] = box2poly([[0., 1.], [0., 1.]])
    cont_props['lot'] = box2poly([[2., 3.], [1., 2.]])

    ppp = abstract.prop2part(cont_state_space, cont_props)
    ppp, new2old = abstract.part2convex(ppp)

    N = 8
    trans_len=1

    disc_params = {}
    for mode in modes:
        disc_params[mode] = {'N':N, 'trans_length':trans_len}

    swab = abstract.discretize_switched(
        ppp, switched_dynamics, disc_params,
        plot=True, show_ts=True, only_adjacent=False
    )

    ts = swab.modes[('normal', 'fly')].ts
    edges = {('s0', 's0'), ('s1', 's1'), ('s2', 's2'), ('s3', 's3'),
             ('s4', 's4'), ('s5', 's5'),
             ('s1', 's2'), ('s1', 's4'),
             ('s2', 's3'), ('s2', 's5'), ('s2', 's0'),
             ('s3', 's0'),
             ('s4', 's5'),
             ('s5', 's0')}
    assert(set(ts.edges() ) == edges)

    ts = swab.ts
    edges.remove(('s2', 's0'))
    assert(set(ts.edges() ) == edges)
    for i, j in edges:
        assert(ts[i][j][0]['env_actions'] == 'normal')
        assert(ts[i][j][0]['sys_actions'] == 'fly')

def test_transient_regions():
    """drift is too strong, so no self-loop must exist

    This bug caused when running union is taken between Presets
    during solve_feasible, as happened with old use_all_horizon,
    cf:
        - 5b1e9681918739b276a221fcc1fd6eebfd058ce3
        - f5f4934ab9d21062f633eef3861ad935c3d3b54b
    """
    dom = pc.box2poly([[0.0, 4.0], [0.0, 3.0]])

    def cont_predicates():
        p = dict()
        p['safe'] = pc.box2poly([[0.5, 3.5], [0.5, 2.5]])

        ppp = abstract.prop2part(dom, p)
        ppp, new2old_reg = abstract.part2convex(ppp)
        ppp.plot()
        return ppp

    def drifting_dynamics():
        A = np.array([[1.0, 0.0],
                      [0.0, 1.0]])

        B = np.array([[1.0],
                      [0.0]])

        U = pc.box2poly([[0.0, 1.0]])

        K = np.array([[-100.0],
                      [0.0]])

        sys = hybrid.LtiSysDyn(A, B, None, K, U, None, dom)
        return sys

    ppp = cont_predicates()
    sys = drifting_dynamics()
    logger.info(sys)

    ab = abstract.discretize(ppp, sys, N=1, use_all_horizon=True,
                             trans_length=1)

    logger.debug(ab.ts)
    self_loops = {i for i,j in ab.ts.transitions() if i==j}
    logger.debug('self loops at states: ' + str(self_loops))

    assert(not self_loops)

    #ax = ab.plot(show_ts=True)
    #ax.figure.savefig('./very_simple.pdf')
<|endoftext|>


<commit_before>#!/usr/bin/env python
import logging
import os, sys, time
import signal
import os.path
import marshal
import cPickle
import multiprocessing
import threading
import SocketServer
import SimpleHTTPServer
import shutil
import socket
import urllib2
import platform
import gc

import zmq

# ignore INFO and DEBUG log
os.environ['GLOG_logtostderr'] = '1'
os.environ['GLOG_minloglevel'] = '1'
import mesos
import mesos_pb2

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from dpark.util import compress, decompress, getproctitle, setproctitle, spawn
from dpark.serialize import marshalable
from dpark.accumulator import Accumulator
from dpark.schedule import Success, OtherFailure
from dpark.env import env
from dpark.shuffle import LocalFileShuffle
from dpark.broadcast import Broadcast

logger = logging.getLogger("executor@%s" % socket.gethostname())

TASK_RESULT_LIMIT = 1024 * 256
DEFAULT_WEB_PORT = 5055
MAX_TASKS_PER_WORKER = 50
MAX_IDLE_TIME = 60
MAX_IDLE_WORKERS = 8
MAX_IDLE_CYCLE = 60 * 30

Script = ''

def reply_status(driver, task, status, data=None):
    update = mesos_pb2.TaskStatus()
    update.task_id.value = task.task_id.value
    update.state = status
    if data is not None:
        update.data = data
    driver.sendStatusUpdate(update)

def run_task(task, ntry):
    try:
        setproctitle('dpark worker %s: run task %s' % (Script, task))
        Accumulator.clear()
        gc.disable()
        result = task.run(ntry)
        accUpdate = Accumulator.values()

        if marshalable(result):
            flag, data = 0, marshal.dumps(result)
        else:
            flag, data = 1, cPickle.dumps(result, -1)
        data = compress(data)

        if len(data) > TASK_RESULT_LIMIT:
            workdir = env.get('WORKDIR')
            name = 'task_%s_%s.result' % (task.id, ntry)
            path = os.path.join(workdir, name)
            f = open(path, 'w')
            f.write(data)
            f.close()
            data = LocalFileShuffle.getServerUri() + '/' + name
            flag += 2

        return mesos_pb2.TASK_FINISHED, cPickle.dumps((task.id, Success(), (flag, data), accUpdate), -1)
    except Exception, e:
        import traceback
        msg = traceback.format_exc()
        return mesos_pb2.TASK_FAILED, cPickle.dumps((task.id, OtherFailure(msg), None, None), -1)
    finally:
        setproctitle('dpark worker: idle')
        gc.collect()
        gc.enable()

def cleanup(workdir):
    while os.getppid() > 1:
        time.sleep(1)

    while os.path.exists(workdir):
        try: shutil.rmtree(workdir, True)
        except: pass

    sys.exit(0)

def init_env(args, workdir):
    setproctitle('dpark worker: idle')
    env.start(False, args)
    spawn(cleanup, workdir)

class LocalizedHTTP(SimpleHTTPServer.SimpleHTTPRequestHandler):
    basedir = None
    def translate_path(self, path):
        out = SimpleHTTPServer.SimpleHTTPRequestHandler.translate_path(self, path)
        return self.basedir + '/' + os.path.relpath(out)

    def log_message(self, format, *args):
        pass

def startWebServer(path):
    # check the default web server
    testpath = os.path.join(path, 'test')
    with open(testpath, 'w') as f:
        f.write(path)
    default_uri = 'http://%s:%d/%s' % (socket.gethostname(), DEFAULT_WEB_PORT,
            os.path.basename(path))
    try:
        data = urllib2.urlopen(default_uri + '/' + 'test').read()
        if data == path:
            return default_uri
    except IOError, e:
        pass

    logger.warning("default webserver at %s not available", DEFAULT_WEB_PORT)
    LocalizedHTTP.basedir = os.path.dirname(path)
    ss = SocketServer.TCPServer(('0.0.0.0', 0), LocalizedHTTP)
    spawn(ss.serve_forever)
    uri = "http://%s:%d/%s" % (socket.gethostname(), ss.server_address[1],
            os.path.basename(path))
    return uri

def forword(fd, addr, prefix=''):
    f = os.fdopen(fd, 'r')
    ctx = zmq.Context()
    out = [None]
    buf = []
    def send(buf):
        if not out[0]:
            out[0] = ctx.socket(zmq.PUSH)
            out[0].connect(addr)
        out[0].send(prefix+''.join(buf))

    while True:
        try:
            line = f.readline()
            if not line: break
            buf.append(line)
            if line.endswith('\n'):
                send(buf)
                buf = []
        except IOError:
            break
    if buf:
        send(buf)
    if out[0]:
        out[0].close()
    f.close()
    ctx.shutdown()

def start_forword(addr, prefix=''):
    rfd, wfd = os.pipe()
    t = spawn(forword, rfd, addr, prefix)
    return t, os.fdopen(wfd, 'w', 0)

def get_pool_memory(pool):
    try:
        import psutil
        p = psutil.Process(pool._pool[0].pid)
        return p.get_memory_info()[0] >> 20
    except Exception:
        return 0

def get_task_memory(task):
    for r in task.resources:
        if r.name == 'mem':
            return r.scalar.value
    logger.error("no memory in resource: %s", task.resources)
    return 100 # 100M

def safe(f):
    def _(self, *a, **kw):
        with self.lock:
            r = f(self, *a, **kw)
        return r
    return _

# cleaner process
def clean_work_dir(path):
    setproctitle('dpark cleaner %s' % path)
    spawn(cleanup, path)

class MyExecutor(mesos.Executor):
    def __init__(self):
        self.workdir = None
        self.idle_workers = []
        self.busy_workers = {}
        self.lock = threading.RLock()

    def check_memory(self, driver):
        try:
            import psutil
        except ImportError:
            logger.error("no psutil module")
            return

        mem_limit = {}
        idle_cycle = 0

        while True:
            self.lock.acquire()

            for tid, (task, pool) in self.busy_workers.items():
                pid = pool._pool[0].pid
                try:
                    p = psutil.Process(pid)
                    rss = p.get_memory_info()[0] >> 20
                except psutil.error.Error, e:
                    logger.error("worker process %d of task %s is dead: %s", pid, tid, e)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                if p.status == psutil.STATUS_ZOMBIE or not p.is_running():
                    logger.error("worker process %d of task %s is zombie", pid, tid)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                offered = get_task_memory(task)
                if not offered:
                    continue
                if rss > offered * 2:
                    logger.error("task %s used too much memory: %dMB > %dMB * 2, kill it. "
                            + "use -M argument to request more memory.", tid, rss, offered)
                    reply_status(driver, task, mesos_pb2.TASK_KILLED)
                    self.busy_workers.pop(tid)
                    pool.terminate()
                elif rss > offered * mem_limit.get(tid, 1.0):
                    logger.warning("task %s used too much memory: %dMB > %dMB, "
                            + "use -M to request more memory", tid, rss, offered)
                    mem_limit[tid] = rss / offered + 0.2

            now = time.time()
            n = len([1 for t, p in self.idle_workers if t + MAX_IDLE_TIME < now])
            if n:
                for _, p in self.idle_workers[:n]:
                    p.terminate()
                self.idle_workers = self.idle_workers[n:]

            self.lock.release()

            if self.idle_workers or self.busy_workers:
                idle_cycle = 0
            else:
                idle_cycle += 1
                if idle_cycle > MAX_IDLE_CYCLE:
                    logger.warning("shutdown idle executor")
                    self.shutdown()

            time.sleep(1)

    @safe
    def registered(self, driver, executorInfo, frameworkInfo, slaveInfo):
        try:
            global Script
            Script, cwd, python_path, osenv, self.parallel, out_logger, err_logger, logLevel, args = marshal.loads(executorInfo.data)
            self.init_args = args
            try:
                os.chdir(cwd)
            except OSError:
                driver.sendFrameworkMessage("switch cwd failed: %s not exists!" % cwd)
            sys.path = python_path
            os.environ.update(osenv)
            prefix = '[%s] ' % socket.gethostname()
            if out_logger:
                self.outt, sys.stdout = start_forword(out_logger, prefix)
            if err_logger:
                self.errt, sys.stderr = start_forword(err_logger, prefix)
            logging.basicConfig(format='%(asctime)-15s [%(levelname)s] [%(name)-9s] %(message)s', level=logLevel)

            self.workdir = args['WORKDIR']
            root = os.path.dirname(self.workdir)
            if not os.path.exists(root):
                os.mkdir(root)
                os.chmod(root, 0777) # because umask
            if not os.path.exists(self.workdir):
                os.mkdir(self.workdir)
            args['SERVER_URI'] = startWebServer(args['WORKDIR'])

            spawn(self.check_memory, driver)

            # wait background threads to been initialized
            time.sleep(0.1)
            if out_logger:
                os.close(1)
                assert os.dup(sys.stdout.fileno()) == 1, 'redirect stdout failed'
            if err_logger:
                os.close(2)
                assert os.dup(sys.stderr.fileno()) == 2, 'redirect stderr failed'

            multiprocessing.Pool(1).apply_async(clean_work_dir, [self.workdir])

            logger.debug("executor started at %s", slaveInfo.hostname)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            logger.error("init executor failed: %s", msg)
            raise

    def reregitered(self, driver, slaveInfo):
        logger.info("executor is reregistered at %s", slaveInfo.hostname)

    def disconnected(self, driver, slaveInfo):
        logger.info("executor is disconnected at %s", slaveInfo.hostname)

    def get_idle_worker(self):
        try:
            return self.idle_workers.pop()[1]
        except IndexError:
            p = multiprocessing.Pool(1, init_env, [self.init_args, self.workdir])
            p.done = 0
            return p

    @safe
    def launchTask(self, driver, task):
        try:
            t, ntry = cPickle.loads(decompress(task.data))

            reply_status(driver, task, mesos_pb2.TASK_RUNNING)

            logging.debug("launch task %s", t.id)

            pool = self.get_idle_worker()
            self.busy_workers[task.task_id.value] = (task, pool)

            def callback((state, data)):
                with self.lock:
                    if task.task_id.value not in self.busy_workers:
                        return
                    _, pool = self.busy_workers.pop(task.task_id.value)
                    pool.done += 1
                    reply_status(driver, task, state, data)
                    if (len(self.idle_workers) + len(self.busy_workers) < self.parallel
                            and len(self.idle_workers) < MAX_IDLE_WORKERS
                            and pool.done < MAX_TASKS_PER_WORKER
                            and get_pool_memory(pool) < get_task_memory(task)): # maybe memory leak in executor
                        self.idle_workers.append((time.time(), pool))
                    else:
                        try: pool.terminate()
                        except: pass

            pool.apply_async(run_task, [t, ntry], callback=callback)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            reply_status(driver, task, mesos_pb2.TASK_LOST, msg)
            return

    @safe
    def killTask(self, driver, taskId):
        if taskId.value in self.busy_workers:
            task, pool = self.busy_workers.pop(taskId.value)
            reply_status(driver, task, mesos_pb2.TASK_KILLED)
            pool.terminate()

    @safe
    def shutdown(self, driver=None):
        def terminate(p):
            try:
                for pi in p._pool:
                    os.kill(pi.pid, signal.SIGKILL)
            except Exception, e:
                pass
        for _, p in self.idle_workers:
            terminate(p)
        for _, p in self.busy_workers.itervalues():
            terminate(p)

        # clean work files
        try: shutil.rmtree(self.workdir, True)
        except: pass

        os._exit(0)

    def error(self, driver, code, message):
        logger.error("error: %s, %s", code, message)

    def frameworkMessage(self, driver, data):
        pass

def run():
    executor = MyExecutor()
    driver = mesos.MesosExecutorDriver(executor)
    driver.run()

if __name__ == '__main__':
    run()
<commit_msg>clean thread should not been daemon<commit_after>#!/usr/bin/env python
import logging
import os, sys, time
import signal
import os.path
import marshal
import cPickle
import multiprocessing
import threading
import SocketServer
import SimpleHTTPServer
import shutil
import socket
import urllib2
import platform
import gc

import zmq

# ignore INFO and DEBUG log
os.environ['GLOG_logtostderr'] = '1'
os.environ['GLOG_minloglevel'] = '1'
import mesos
import mesos_pb2

sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))
from dpark.util import compress, decompress, getproctitle, setproctitle, spawn
from dpark.serialize import marshalable
from dpark.accumulator import Accumulator
from dpark.schedule import Success, OtherFailure
from dpark.env import env
from dpark.shuffle import LocalFileShuffle
from dpark.broadcast import Broadcast

logger = logging.getLogger("executor@%s" % socket.gethostname())

TASK_RESULT_LIMIT = 1024 * 256
DEFAULT_WEB_PORT = 5055
MAX_TASKS_PER_WORKER = 50
MAX_IDLE_TIME = 60
MAX_IDLE_WORKERS = 8
MAX_IDLE_CYCLE = 60 * 30

Script = ''

def reply_status(driver, task, status, data=None):
    update = mesos_pb2.TaskStatus()
    update.task_id.value = task.task_id.value
    update.state = status
    if data is not None:
        update.data = data
    driver.sendStatusUpdate(update)

def run_task(task, ntry):
    try:
        setproctitle('dpark worker %s: run task %s' % (Script, task))
        Accumulator.clear()
        gc.disable()
        result = task.run(ntry)
        accUpdate = Accumulator.values()

        if marshalable(result):
            flag, data = 0, marshal.dumps(result)
        else:
            flag, data = 1, cPickle.dumps(result, -1)
        data = compress(data)

        if len(data) > TASK_RESULT_LIMIT:
            workdir = env.get('WORKDIR')
            name = 'task_%s_%s.result' % (task.id, ntry)
            path = os.path.join(workdir, name)
            f = open(path, 'w')
            f.write(data)
            f.close()
            data = LocalFileShuffle.getServerUri() + '/' + name
            flag += 2

        return mesos_pb2.TASK_FINISHED, cPickle.dumps((task.id, Success(), (flag, data), accUpdate), -1)
    except Exception, e:
        import traceback
        msg = traceback.format_exc()
        return mesos_pb2.TASK_FAILED, cPickle.dumps((task.id, OtherFailure(msg), None, None), -1)
    finally:
        setproctitle('dpark worker: idle')
        gc.collect()
        gc.enable()

def cleanup(workdir):
    while os.getppid() > 1:
        time.sleep(1)

    while os.path.exists(workdir):
        try: shutil.rmtree(workdir, True)
        except: pass

    os._exit(0)

def init_env(args, workdir):
    setproctitle('dpark worker: idle')
    env.start(False, args)
    threading.Thread(target=cleanup, args=[workdir]).start()

class LocalizedHTTP(SimpleHTTPServer.SimpleHTTPRequestHandler):
    basedir = None
    def translate_path(self, path):
        out = SimpleHTTPServer.SimpleHTTPRequestHandler.translate_path(self, path)
        return self.basedir + '/' + os.path.relpath(out)

    def log_message(self, format, *args):
        pass

def startWebServer(path):
    # check the default web server
    testpath = os.path.join(path, 'test')
    with open(testpath, 'w') as f:
        f.write(path)
    default_uri = 'http://%s:%d/%s' % (socket.gethostname(), DEFAULT_WEB_PORT,
            os.path.basename(path))
    try:
        data = urllib2.urlopen(default_uri + '/' + 'test').read()
        if data == path:
            return default_uri
    except IOError, e:
        pass

    logger.warning("default webserver at %s not available", DEFAULT_WEB_PORT)
    LocalizedHTTP.basedir = os.path.dirname(path)
    ss = SocketServer.TCPServer(('0.0.0.0', 0), LocalizedHTTP)
    spawn(ss.serve_forever)
    uri = "http://%s:%d/%s" % (socket.gethostname(), ss.server_address[1],
            os.path.basename(path))
    return uri

def forword(fd, addr, prefix=''):
    f = os.fdopen(fd, 'r')
    ctx = zmq.Context()
    out = [None]
    buf = []
    def send(buf):
        if not out[0]:
            out[0] = ctx.socket(zmq.PUSH)
            out[0].connect(addr)
        out[0].send(prefix+''.join(buf))

    while True:
        try:
            line = f.readline()
            if not line: break
            buf.append(line)
            if line.endswith('\n'):
                send(buf)
                buf = []
        except IOError:
            break
    if buf:
        send(buf)
    if out[0]:
        out[0].close()
    f.close()
    ctx.shutdown()

def start_forword(addr, prefix=''):
    rfd, wfd = os.pipe()
    t = spawn(forword, rfd, addr, prefix)
    return t, os.fdopen(wfd, 'w', 0)

def get_pool_memory(pool):
    try:
        import psutil
        p = psutil.Process(pool._pool[0].pid)
        return p.get_memory_info()[0] >> 20
    except Exception:
        return 0

def get_task_memory(task):
    for r in task.resources:
        if r.name == 'mem':
            return r.scalar.value
    logger.error("no memory in resource: %s", task.resources)
    return 100 # 100M

def safe(f):
    def _(self, *a, **kw):
        with self.lock:
            r = f(self, *a, **kw)
        return r
    return _

# cleaner process
def clean_work_dir(path):
    setproctitle('dpark cleaner %s' % path)
    threading.Thread(target=cleanup, args=[path]).start()

class MyExecutor(mesos.Executor):
    def __init__(self):
        self.workdir = None
        self.idle_workers = []
        self.busy_workers = {}
        self.lock = threading.RLock()

    def check_memory(self, driver):
        try:
            import psutil
        except ImportError:
            logger.error("no psutil module")
            return

        mem_limit = {}
        idle_cycle = 0

        while True:
            self.lock.acquire()

            for tid, (task, pool) in self.busy_workers.items():
                pid = pool._pool[0].pid
                try:
                    p = psutil.Process(pid)
                    rss = p.get_memory_info()[0] >> 20
                except psutil.error.Error, e:
                    logger.error("worker process %d of task %s is dead: %s", pid, tid, e)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                if p.status == psutil.STATUS_ZOMBIE or not p.is_running():
                    logger.error("worker process %d of task %s is zombie", pid, tid)
                    reply_status(driver, task, mesos_pb2.TASK_LOST)
                    self.busy_workers.pop(tid)
                    continue

                offered = get_task_memory(task)
                if not offered:
                    continue
                if rss > offered * 2:
                    logger.error("task %s used too much memory: %dMB > %dMB * 2, kill it. "
                            + "use -M argument to request more memory.", tid, rss, offered)
                    reply_status(driver, task, mesos_pb2.TASK_KILLED)
                    self.busy_workers.pop(tid)
                    pool.terminate()
                elif rss > offered * mem_limit.get(tid, 1.0):
                    logger.warning("task %s used too much memory: %dMB > %dMB, "
                            + "use -M to request more memory", tid, rss, offered)
                    mem_limit[tid] = rss / offered + 0.2

            now = time.time()
            n = len([1 for t, p in self.idle_workers if t + MAX_IDLE_TIME < now])
            if n:
                for _, p in self.idle_workers[:n]:
                    p.terminate()
                self.idle_workers = self.idle_workers[n:]

            self.lock.release()

            if self.idle_workers or self.busy_workers:
                idle_cycle = 0
            else:
                idle_cycle += 1
                if idle_cycle > MAX_IDLE_CYCLE:
                    logger.warning("shutdown idle executor")
                    self.shutdown()

            time.sleep(1)

    @safe
    def registered(self, driver, executorInfo, frameworkInfo, slaveInfo):
        try:
            global Script
            Script, cwd, python_path, osenv, self.parallel, out_logger, err_logger, logLevel, args = marshal.loads(executorInfo.data)
            self.init_args = args
            try:
                os.chdir(cwd)
            except OSError:
                driver.sendFrameworkMessage("switch cwd failed: %s not exists!" % cwd)
            sys.path = python_path
            os.environ.update(osenv)
            prefix = '[%s] ' % socket.gethostname()
            if out_logger:
                self.outt, sys.stdout = start_forword(out_logger, prefix)
            if err_logger:
                self.errt, sys.stderr = start_forword(err_logger, prefix)
            logging.basicConfig(format='%(asctime)-15s [%(levelname)s] [%(name)-9s] %(message)s', level=logLevel)

            self.workdir = args['WORKDIR']
            root = os.path.dirname(self.workdir)
            if not os.path.exists(root):
                os.mkdir(root)
                os.chmod(root, 0777) # because umask
            if not os.path.exists(self.workdir):
                os.mkdir(self.workdir)
            args['SERVER_URI'] = startWebServer(args['WORKDIR'])

            spawn(self.check_memory, driver)

            # wait background threads to been initialized
            time.sleep(0.1)
            if out_logger:
                os.close(1)
                assert os.dup(sys.stdout.fileno()) == 1, 'redirect stdout failed'
            if err_logger:
                os.close(2)
                assert os.dup(sys.stderr.fileno()) == 2, 'redirect stderr failed'

            multiprocessing.Pool(1).apply_async(clean_work_dir, [self.workdir])

            logger.debug("executor started at %s", slaveInfo.hostname)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            logger.error("init executor failed: %s", msg)
            raise

    def reregitered(self, driver, slaveInfo):
        logger.info("executor is reregistered at %s", slaveInfo.hostname)

    def disconnected(self, driver, slaveInfo):
        logger.info("executor is disconnected at %s", slaveInfo.hostname)

    def get_idle_worker(self):
        try:
            return self.idle_workers.pop()[1]
        except IndexError:
            p = multiprocessing.Pool(1, init_env, [self.init_args, self.workdir])
            p.done = 0
            return p

    @safe
    def launchTask(self, driver, task):
        try:
            t, ntry = cPickle.loads(decompress(task.data))

            reply_status(driver, task, mesos_pb2.TASK_RUNNING)

            logging.debug("launch task %s", t.id)

            pool = self.get_idle_worker()
            self.busy_workers[task.task_id.value] = (task, pool)

            def callback((state, data)):
                with self.lock:
                    if task.task_id.value not in self.busy_workers:
                        return
                    _, pool = self.busy_workers.pop(task.task_id.value)
                    pool.done += 1
                    reply_status(driver, task, state, data)
                    if (len(self.idle_workers) + len(self.busy_workers) < self.parallel
                            and len(self.idle_workers) < MAX_IDLE_WORKERS
                            and pool.done < MAX_TASKS_PER_WORKER
                            and get_pool_memory(pool) < get_task_memory(task)): # maybe memory leak in executor
                        self.idle_workers.append((time.time(), pool))
                    else:
                        try: pool.terminate()
                        except: pass

            pool.apply_async(run_task, [t, ntry], callback=callback)

        except Exception, e:
            import traceback
            msg = traceback.format_exc()
            reply_status(driver, task, mesos_pb2.TASK_LOST, msg)
            return

    @safe
    def killTask(self, driver, taskId):
        if taskId.value in self.busy_workers:
            task, pool = self.busy_workers.pop(taskId.value)
            reply_status(driver, task, mesos_pb2.TASK_KILLED)
            pool.terminate()

    @safe
    def shutdown(self, driver=None):
        def terminate(p):
            try:
                for pi in p._pool:
                    os.kill(pi.pid, signal.SIGKILL)
            except Exception, e:
                pass
        for _, p in self.idle_workers:
            terminate(p)
        for _, p in self.busy_workers.itervalues():
            terminate(p)

        # clean work files
        try: shutil.rmtree(self.workdir, True)
        except: pass

        os._exit(0)

    def error(self, driver, code, message):
        logger.error("error: %s, %s", code, message)

    def frameworkMessage(self, driver, data):
        pass

def run():
    executor = MyExecutor()
    driver = mesos.MesosExecutorDriver(executor)
    driver.run()

if __name__ == '__main__':
    run()
<|endoftext|>