<commit_before>import copy
import six

from chainer.training import extension
from chainer import backend
from chainer.dataset import convert
from chainer import function
from chainer.utils import argument

class GenericMultiNodeEvaluator(extension.Extension):
    '''Generic multi-node evaluator for non-allreducable evaluation.

    This is to evaluate a Dataset that cannot evenly divided accross
    nodes, and for evaluation calculation that is not applicable to a
    simple add-and-devide style averaging among processes.

    Users are recommeneded to implement its own local calculation
    ``calc_local()`` (e.g.  at each distributed GPU) and aggregation
    ``aggregate()`` of its results. Although it has built-in
    implementaiton of those two methods.

    It has several drawbacks; 1) Additional implementation of
    aggregation required to users, and 2) no compatibility with
    :class:`~chainer.training.extensions.Evaluator`.

    .. note:: No automatic support of Reporter is provided; Set it up
       at ``initialize()`` method

    Args:
        comm:
            ChainerMN communicator object
        iterator:
            An iterator for test dataset. Must be non-repeated.
        target (callable):
            A model to evaluate with test dataset
        device (int or chainer.backend.Device):
            A device indicator to send data with converter. Not used
            when the converter is not using any devices.
        converter (callable):
            A converter. Default value is
            :func:`chainer.dataset.concat_examples` .
        root (int):
            Rank number of root process to run bcast and gather with.
        progress_hook (callable):
            A callable that receives single argument for indicators. The
            callable is only callled at root process.

    '''
    trigger = 1, 'epoch'
    default_name = 'validation'
    priority = extension.PRIORITY_WRITER

    name = None

    def __init__(self, comm, iterator, target, device=None,
                 converter=convert.concat_examples, root=0,
                 **kwargs):
        progress_hook, = argument.parse_kwargs(kwargs, ('progress_hook', None))

        self.comm = comm
        self.iterator = iterator
        self._targets = {"main": target}
        self.converter = converter

        if device is not None:
            device = backend.get_device(device)
        self.device = device

        self._progress_hook = progress_hook

        assert 0 <= root and root < self.comm.size
        self.root = root

    def __call__(self, trainer):
        if hasattr(self.iterator, 'reset'):
            self.iterator.reset()
            it = self.iterator
        else:
            it = copy.copy(self.iterator)

        if self.comm is not None:
            gen = self._evaluate_local(it)

            if self.comm.rank == self.root:
                total_result = self.aggregate([result for result in gen])
            else:
                for _ in gen:
                    pass
                total_result = None

        else:
            # Non-multinode environment
            gen = self._evaluate_local_single(self, it)
            total_result = self.aggregate([result for result in gen])

        return total_result

    def calc_local(self, *args, **kwargs):
        '''A generic method for local calculation.

        Override this method to run its local calculation.  Otherwise,
        results are calculated with original target and test dataset.

        Args:
            args:
                Result of converter when it is tuple.
            kwargs:
                Result of converter when it is dict.

        Returns:
            Arbrary value may be returned, but must not be ``None``.

        '''
        target = self._targets['main']
        return target(*args, **kwargs)

    def aggregate(self, results):
        '''A generic aggregation method.

        Override this method for original aggregation calculation. By
        default, it just does nothing but returns the input. This
        method is called once and only once accross the cluster, at
        root process. Reporting can be run here.

        Args:
            results (list):
                List of return value of ``calc_local()`` obtained from
                all nodes..

        '''
        return results

    def _evaluate_local_single(self, iterator):
        for batch in iterator:
            in_arrays = convert._call_converter(
                self.converter, batch, self.device)

            with function.no_backprop_mode():
                if isinstance(in_arrays, tuple):
                    results = self.calc_local(*in_arrays)
                elif isinstance(in_arrays, dict):
                    results = self.calc_local(**in_arrays)
                else:
                    results = self.calc_local(in_arrays)

            if self._progress_hook:
                self._progress_hook(batch)
            yield results

    def _evaluate_local(self, iterator):
        rounds = 8  # Checks whether local eval is all done every 8 rounds
        all_done = None

        while not all_done:
            all_done = None
            results = None
            for i in range(rounds):
                try:
                    batch = iterator.next()
                    in_arrays = convert._call_converter(
                        self.converter, batch, self.device)

                    with function.no_backprop_mode():
                        if isinstance(in_arrays, tuple):
                            results = self.calc_local(*in_arrays)
                        elif isinstance(in_arrays, dict):
                            results = self.calc_local(**in_arrays)
                        else:
                            results = self.calc_local(in_arrays)

                    if self.comm.rank == self.root and self._progress_hook:
                        self._progress_hook(batch)

                except StopIteration:
                    batch = None
                    results = None

                results = self.comm.gather_obj(results, root=self.root)

                if self.comm.rank == self.root:
                    valid_results = [r for r in results if r is not None]
                    for result in valid_results:
                        yield result

                    all_done = len(valid_results) == 0

            all_done = self.comm.bcast_obj(all_done, root=self.root)
        return

def create_multi_node_evaluator(actual_evaluator, communicator):
    """Create a multi node evaluator from a normal evaluator.

    Actually this method patches the evaluator to work in multi node
    environment. This method adds several hidden attributes starting
    with `_mn_` prefix.

    Args:
        actual_evaluator: evaluator to be patched
            (e.g., ``chainer.training.extensions.Evaluator``)
        communicator: ChainerMN communicator

    Returns:
        The multi-node patched ``actual_evaluator``.

    .. note:: After patched, original evaluator does not work
              correctly in non-MPI environment.

    """

    actual_evaluator._mn_original_evaluate = actual_evaluator.evaluate
    actual_evaluator._mn_communicator = communicator

    def new_evaluate(self):
        local_mean_dict = self._mn_original_evaluate()
        global_mean_dict = {
            name:
            self._mn_communicator.allreduce_obj(
                value) / self._mn_communicator.size
            for name, value in sorted(local_mean_dict.items())
        }
        return global_mean_dict

    actual_evaluator.evaluate = six.create_bound_method(
        new_evaluate, actual_evaluator)
    return actual_evaluator
<commit_msg>Typo and names fix<commit_after>import copy
import six

from chainer.training import extension
from chainer import backend
from chainer.dataset import convert
from chainer import function
from chainer.utils import argument

class GenericMultiNodeEvaluator(extension.Extension):
    '''Generic multi-node evaluator for non-allreducable evaluation.

    This is to evaluate a Dataset that cannot evenly divided across
    all processes in the communicator, for evaluation calculation that
    is not applicable to a simple add-and-devide style averaging among
    processes.

    Users are recommeneded to implement its own local calculation
    ``calc_local()`` (e.g.  at each distributed GPU) and aggregation
    ``aggregate()`` of its results. Although it has built-in
    implementaiton of those two methods.

    It has several drawbacks; 1) Additional implementation of
    aggregation required to users, and 2) no compatibility with
    :class:`~chainer.training.extensions.Evaluator`.

    .. note:: No automatic support of Reporter is provided; Set it up
       at ``initialize()`` method

    Args:
        comm:
            ChainerMN communicator object
        iterator:
            An iterator for test dataset. Must be non-repeated.
        target (callable):
            A model to evaluate with test dataset
        device (int or chainer.backend.Device):
            A device indicator to send data with converter. Not used
            when the converter is not using any devices.
        converter (callable):
            A converter. Default value is
            :func:`chainer.dataset.concat_examples` .
        root (int):
            Rank number of root process to run bcast and gather with.
        progress_hook (callable):
            A callable that receives single argument for indicators. The
            callable is only callled at root process.

    '''
    trigger = 1, 'epoch'
    default_name = 'validation'
    priority = extension.PRIORITY_WRITER

    name = None

    def __init__(self, comm, iterator, target, device=None,
                 converter=convert.concat_examples, root=0,
                 **kwargs):
        progress_hook, = argument.parse_kwargs(kwargs, ('progress_hook', None))

        self.comm = comm
        self.iterator = iterator
        self._targets = {"main": target}
        self.converter = converter

        if device is not None:
            device = backend.get_device(device)
        self.device = device

        self._progress_hook = progress_hook

        assert 0 <= root and root < self.comm.size
        self.root = root

    def __call__(self, trainer):
        if hasattr(self.iterator, 'reset'):
            self.iterator.reset()
            it = self.iterator
        else:
            it = copy.copy(self.iterator)

        if self.comm is not None:
            gen = self._evaluate_local(it)

            if self.comm.rank == self.root:
                total_result = self.aggregate([result for result in gen])
            else:
                for _ in gen:
                    pass
                total_result = None

        else:
            # Non-multinode environment
            gen = self._evaluate_local_single(self, it)
            total_result = self.aggregate([result for result in gen])

        return total_result

    def calc_local(self, *args, **kwargs):
        '''A generic method for local calculation.

        Override this method to run its local calculation.  Otherwise,
        results are calculated with original target and test dataset.

        Args:
            args:
                Result of converter when it is tuple.
            kwargs:
                Result of converter when it is dict.

        Returns:
            Arbrary value may be returned, but must not be ``None``.

        '''
        target = self._targets['main']
        return target(*args, **kwargs)

    def aggregate(self, results):
        '''A generic aggregation method.

        Override this method for original aggregation calculation. By
        default, it just does nothing but returns the input. This
        method is called once and only once accross the cluster, at
        root process. Reporting can be run here.

        Args:
            results (list):
                List of return value of ``calc_local()`` obtained from
                all nodes..

        '''
        return results

    def _evaluate_local_single(self, iterator):
        for batch in iterator:
            in_arrays = convert._call_converter(
                self.converter, batch, self.device)

            with function.no_backprop_mode():
                if isinstance(in_arrays, tuple):
                    results = self.calc_local(*in_arrays)
                elif isinstance(in_arrays, dict):
                    results = self.calc_local(**in_arrays)
                else:
                    results = self.calc_local(in_arrays)

            if self._progress_hook:
                self._progress_hook(batch)
            yield results

    def _evaluate_local(self, iterator):
        # Check whether local eval is all done every 8 rounds
        gather_interval = 8

        all_done = None
        while not all_done:
            all_done = None
            results = None
            for _ in range(gather_interval):
                try:
                    batch = iterator.next()
                    in_arrays = convert._call_converter(
                        self.converter, batch, self.device)

                    with function.no_backprop_mode():
                        if isinstance(in_arrays, tuple):
                            results = self.calc_local(*in_arrays)
                        elif isinstance(in_arrays, dict):
                            results = self.calc_local(**in_arrays)
                        else:
                            results = self.calc_local(in_arrays)

                    if self.comm.rank == self.root and self._progress_hook:
                        self._progress_hook(batch)

                except StopIteration:
                    batch = None
                    results = None

                results = self.comm.gather_obj(results, root=self.root)

                if self.comm.rank == self.root:
                    valid_results = [r for r in results if r is not None]
                    for result in valid_results:
                        yield result

                    all_done = len(valid_results) == 0

            all_done = self.comm.bcast_obj(all_done, root=self.root)
        return

def create_multi_node_evaluator(actual_evaluator, communicator):
    """Create a multi node evaluator from a normal evaluator.

    Actually this method patches the evaluator to work in multi node
    environment. This method adds several hidden attributes starting
    with `_mn_` prefix.

    Args:
        actual_evaluator: evaluator to be patched
            (e.g., ``chainer.training.extensions.Evaluator``)
        communicator: ChainerMN communicator

    Returns:
        The multi-node patched ``actual_evaluator``.

    .. note:: After patched, original evaluator does not work
              correctly in non-MPI environment.

    """

    actual_evaluator._mn_original_evaluate = actual_evaluator.evaluate
    actual_evaluator._mn_communicator = communicator

    def new_evaluate(self):
        local_mean_dict = self._mn_original_evaluate()
        global_mean_dict = {
            name:
            self._mn_communicator.allreduce_obj(
                value) / self._mn_communicator.size
            for name, value in sorted(local_mean_dict.items())
        }
        return global_mean_dict

    actual_evaluator.evaluate = six.create_bound_method(
        new_evaluate, actual_evaluator)
    return actual_evaluator
<|endoftext|>


<commit_before># Copyright (c) 2013-2014 Parallels, Inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         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.

import mock

from oslo.config import cfg

from nova.compute import power_state
from nova.openstack.common import uuidutils
from nova import test
from nova.virt import fake

from pcsnovadriver.pcs import prlsdkapi_proxy
from pcsnovadriver.tests.pcs import fakeprlsdkapi

prlsdkapi_proxy.prlsdkapi = fakeprlsdkapi

from pcsnovadriver.pcs import driver

CONF = cfg.CONF
CONF.import_opt('compute_manager', 'nova.service')
CONF.import_opt('firewall_driver', 'nova.virt.firewall')
CONF.import_opt('host', 'nova.netconf')
CONF.import_opt('my_ip', 'nova.netconf')
CONF.import_opt('image_cache_subdirectory_name', 'nova.virt.imagecache')
CONF.import_opt('instances_path', 'nova.compute.manager')

CONF.firewall_driver = "nova.virt.firewall.NoopFirewallDriver"

vm_stopped = {
            'name': 'instance001',
            'uuid': '{19be06cb-a6f2-47a7-a53e-11bc6d4c3b98}',
            'ram_size': 1024,
            'cpu_count': 1,
            'state': fakeprlsdkapi.consts.VMS_STOPPED,
        }

vm_running = {
            'name': 'instance002',
            'uuid': '{d58fe074-ce99-46b7-8ce1-83620ba26426}',
            'ram_size': 2048,
            'cpu_count': 4,
            'state': fakeprlsdkapi.consts.VMS_RUNNING,
        }

vms = [vm_stopped, vm_running]

instance_stopped = {
            'name': vm_stopped['name'],
            'power_state': power_state.SHUTDOWN
        }

instance_running = {
            'name': vm_running['name'],
            'power_state': power_state.RUNNING
        }

network_info_1vif = [
                {'id': uuidutils.generate_uuid()},
            ]

network_info_3vif = [
                {'id': uuidutils.generate_uuid()},
                {'id': uuidutils.generate_uuid()},
                {'id': uuidutils.generate_uuid()},
            ]

class PCSDriverTestCase(test.TestCase):

    def setUp(self):
        super(PCSDriverTestCase, self).setUp()
        self.conn = driver.PCSDriver(fake.FakeVirtAPI(), True)
        self.conn.init_host(host='localhost')
        self.conn.psrv.test_add_vms(vms)

    def test_list_instances(self):
        instances = self.conn.list_instances()
        self.assertEqual(instances, map(lambda x: x['name'], vms))

    def test_list_instance_uuids(self):
        instances = self.conn.list_instance_uuids()
        self.assertEqual(instances, map(lambda x: x['uuid'][1:-1], vms))

    def test_instance_exists_exists(self):
        self.assertTrue(self.conn.instance_exists(vms[0]['name']))

    def test_instance_exists_notexists(self):
        self.assertFalse(self.conn.instance_exists('x' + vms[0]['name']))

    def test_get_host_ip_addr(self):
        self.assertEqual(self.conn.get_host_ip_addr(), CONF.my_ip)

    def test_get_info(self):
        vm = vms[0]
        info = self.conn.get_info({'id': vm['uuid'], 'name': vm['name']})
        self.assertEqual(info['state'], power_state.SHUTDOWN)
        self.assertEqual(info['max_mem'], vm['ram_size'])
        self.assertEqual(info['mem'], vm['ram_size'])
        self.assertEqual(info['num_cpu'], vm['cpu_count'])
        self.assertEqual(info['cpu_time'], 1000)

    def _prep_plug_vifs(self, vm):
        self.conn._get_ve_by_name = mock.MagicMock()
        sdk_ve = fakeprlsdkapi.Vm(vm)
        self.conn._get_ve_by_name.return_value = sdk_ve

        self.conn.vif_driver = mock.MagicMock()

        return sdk_ve

    def test_plug_vifs_running(self):
        vm = vm_running
        sdk_ve = self._prep_plug_vifs(vm)

        self.conn.plug_vifs(instance_running, network_info_1vif)

        args = (self.conn, instance_running, sdk_ve, network_info_1vif[0],)
        self.conn.vif_driver.plug.assert_called_once_with(*args)
        self.assertEqual(sdk_ve.state, fakeprlsdkapi.consts.VMS_RUNNING)

    def test_plug_vifs_stopped(self):
        vm = vm_stopped
        sdk_ve = self._prep_plug_vifs(vm)

        self.conn.plug_vifs(instance_stopped, network_info_1vif)

        self.assertEqual(self.conn.vif_driver.plug.call_count, 0)
        self.assertEqual(sdk_ve.state, fakeprlsdkapi.consts.VMS_STOPPED)

    def test_plug_vifs_unexistent(self):
        instance = {
                    'name': 'this_vm_does_not_exist',
                    'power_state': power_state.RUNNING
                }

        self.conn.vif_driver = mock.MagicMock()
        self.conn.plug_vifs(instance, network_info_1vif)

        self.assertEqual(self.conn.vif_driver.plug.call_count, 0)

    def test_plug_vifs_stopped_started(self):
        vm = vm_stopped
        sdk_ve = self._prep_plug_vifs(vm)

        instance = instance_stopped.copy()
        instance['power_state'] = power_state.RUNNING
        self.conn.plug_vifs(instance, network_info_1vif)

        args = (self.conn, instance, sdk_ve, network_info_1vif[0],)
        self.conn.vif_driver.plug.assert_called_once_with(*args)
        self.assertEqual(sdk_ve.state, fakeprlsdkapi.consts.VMS_RUNNING)

    def test_plug_vifs_started_stopped(self):
        vm = vm_running
        sdk_ve = self._prep_plug_vifs(vm)

        instance = instance_running.copy()
        instance['power_state'] = power_state.SHUTDOWN
        self.conn.plug_vifs(instance, network_info_1vif)

        self.assertEqual(self.conn.vif_driver.plug.call_count, 0)
        self.assertEqual(sdk_ve.state, fakeprlsdkapi.consts.VMS_STOPPED)
<commit_msg>make variable 'pc' with consts<commit_after># Copyright (c) 2013-2014 Parallels, Inc.
# All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         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.

import mock

from oslo.config import cfg

from nova.compute import power_state
from nova.openstack.common import uuidutils
from nova import test
from nova.virt import fake

from pcsnovadriver.pcs import prlsdkapi_proxy
from pcsnovadriver.tests.pcs import fakeprlsdkapi

prlsdkapi_proxy.prlsdkapi = fakeprlsdkapi
pc = fakeprlsdkapi.consts

from pcsnovadriver.pcs import driver

CONF = cfg.CONF
CONF.import_opt('compute_manager', 'nova.service')
CONF.import_opt('firewall_driver', 'nova.virt.firewall')
CONF.import_opt('host', 'nova.netconf')
CONF.import_opt('my_ip', 'nova.netconf')
CONF.import_opt('image_cache_subdirectory_name', 'nova.virt.imagecache')
CONF.import_opt('instances_path', 'nova.compute.manager')

CONF.firewall_driver = "nova.virt.firewall.NoopFirewallDriver"

vm_stopped = {
            'name': 'instance001',
            'uuid': '{19be06cb-a6f2-47a7-a53e-11bc6d4c3b98}',
            'ram_size': 1024,
            'cpu_count': 1,
            'state': pc.VMS_STOPPED,
        }

vm_running = {
            'name': 'instance002',
            'uuid': '{d58fe074-ce99-46b7-8ce1-83620ba26426}',
            'ram_size': 2048,
            'cpu_count': 4,
            'state': pc.VMS_RUNNING,
        }

vms = [vm_stopped, vm_running]

instance_stopped = {
            'name': vm_stopped['name'],
            'power_state': power_state.SHUTDOWN
        }

instance_running = {
            'name': vm_running['name'],
            'power_state': power_state.RUNNING
        }

network_info_1vif = [
                {'id': uuidutils.generate_uuid()},
            ]

network_info_3vif = [
                {'id': uuidutils.generate_uuid()},
                {'id': uuidutils.generate_uuid()},
                {'id': uuidutils.generate_uuid()},
            ]

class PCSDriverTestCase(test.TestCase):

    def setUp(self):
        super(PCSDriverTestCase, self).setUp()
        self.conn = driver.PCSDriver(fake.FakeVirtAPI(), True)
        self.conn.init_host(host='localhost')
        self.conn.psrv.test_add_vms(vms)

    def test_list_instances(self):
        instances = self.conn.list_instances()
        self.assertEqual(instances, map(lambda x: x['name'], vms))

    def test_list_instance_uuids(self):
        instances = self.conn.list_instance_uuids()
        self.assertEqual(instances, map(lambda x: x['uuid'][1:-1], vms))

    def test_instance_exists_exists(self):
        self.assertTrue(self.conn.instance_exists(vms[0]['name']))

    def test_instance_exists_notexists(self):
        self.assertFalse(self.conn.instance_exists('x' + vms[0]['name']))

    def test_get_host_ip_addr(self):
        self.assertEqual(self.conn.get_host_ip_addr(), CONF.my_ip)

    def test_get_info(self):
        vm = vms[0]
        info = self.conn.get_info({'id': vm['uuid'], 'name': vm['name']})
        self.assertEqual(info['state'], power_state.SHUTDOWN)
        self.assertEqual(info['max_mem'], vm['ram_size'])
        self.assertEqual(info['mem'], vm['ram_size'])
        self.assertEqual(info['num_cpu'], vm['cpu_count'])
        self.assertEqual(info['cpu_time'], 1000)

    def _prep_plug_vifs(self, vm):
        self.conn._get_ve_by_name = mock.MagicMock()
        sdk_ve = fakeprlsdkapi.Vm(vm)
        self.conn._get_ve_by_name.return_value = sdk_ve

        self.conn.vif_driver = mock.MagicMock()

        return sdk_ve

    def test_plug_vifs_running(self):
        vm = vm_running
        sdk_ve = self._prep_plug_vifs(vm)

        self.conn.plug_vifs(instance_running, network_info_1vif)

        args = (self.conn, instance_running, sdk_ve, network_info_1vif[0],)
        self.conn.vif_driver.plug.assert_called_once_with(*args)
        self.assertEqual(sdk_ve.state, pc.VMS_RUNNING)

    def test_plug_vifs_stopped(self):
        vm = vm_stopped
        sdk_ve = self._prep_plug_vifs(vm)

        self.conn.plug_vifs(instance_stopped, network_info_1vif)

        self.assertEqual(self.conn.vif_driver.plug.call_count, 0)
        self.assertEqual(sdk_ve.state, pc.VMS_STOPPED)

    def test_plug_vifs_unexistent(self):
        instance = {
                    'name': 'this_vm_does_not_exist',
                    'power_state': power_state.RUNNING
                }

        self.conn.vif_driver = mock.MagicMock()
        self.conn.plug_vifs(instance, network_info_1vif)

        self.assertEqual(self.conn.vif_driver.plug.call_count, 0)

    def test_plug_vifs_stopped_started(self):
        vm = vm_stopped
        sdk_ve = self._prep_plug_vifs(vm)

        instance = instance_stopped.copy()
        instance['power_state'] = power_state.RUNNING
        self.conn.plug_vifs(instance, network_info_1vif)

        args = (self.conn, instance, sdk_ve, network_info_1vif[0],)
        self.conn.vif_driver.plug.assert_called_once_with(*args)
        self.assertEqual(sdk_ve.state, pc.VMS_RUNNING)

    def test_plug_vifs_started_stopped(self):
        vm = vm_running
        sdk_ve = self._prep_plug_vifs(vm)

        instance = instance_running.copy()
        instance['power_state'] = power_state.SHUTDOWN
        self.conn.plug_vifs(instance, network_info_1vif)

        self.assertEqual(self.conn.vif_driver.plug.call_count, 0)
        self.assertEqual(sdk_ve.state, pc.VMS_STOPPED)
<|endoftext|>


<commit_before>"""
Provides a build rule for building the AdvancedExample from Blaze.
"""

load("//java/com/google/ads/interactivemedia/v3/samples:build_defs.bzl", "build_sample_package")

COMMON_DEPS = [
    "//third_party/java/android_libs/exoplayer:exoplayer2_core",
    "//third_party/java/android_libs/exoplayer:exoplayer2_dash",
    "//third_party/java/android_libs/exoplayer:exoplayer2_ext_mediasession",
    "//third_party/java/android_libs/exoplayer:exoplayer2_hls",
    "//third_party/java/android_libs/exoplayer:exoplayer2_ui",
    "//third_party/java/android_libs/guava_jdk5:collect",
    "//third_party/java/androidx/annotation",
    "//third_party/java/android/android_sdk_linux/extras/android/compatibility/v4",
    "//third_party/java/android/android_sdk_linux/extras/android/compatibility/v7/appcompat",
    "//third_party/java/androidx/constraintlayout",
    "//java/com/google/android/gmscore/integ/client/base",
    "//java/com/google/android/gmscore/integ/client/cast/framework",
    "//java/com/google/android/gmscore/integ/client/cast",
]

DEBUG_DEPS = COMMON_DEPS + [
    "//java/com/google/ads/interactivemedia/v3:sdk_lib_debug",
]

COMPILED_DEPS = COMMON_DEPS + [
    "//java/com/google/ads/interactivemedia/v3:sdk_lib",
    "//java/com/google/android/gmscore/integ/client/ads_identifier",
]

def advanced_example_package():
    build_sample_package(
        name_prefix = "advanced",
        package_name = "com.google.ads.interactivemedia.v3.samples.videoplayerapp",
        srcs = native.glob([
            "AdvancedExample/app/src/main/java/com/google/ads/interactivemedia/v3/samples/samplevideoplayer/*.java",
            "AdvancedExample/app/src/main/java/com/google/ads/interactivemedia/v3/samples/videoplayerapp/*.java",
        ]),
        compiled_deps = COMPILED_DEPS,
        debug_deps = COMPILED_DEPS,  # We can't use real debug deps here due to collisions between AdShield and GMS Core dependencies.
        manifest = "AdvancedExample/app/src/main/AndroidManifest.xml",
        resources = native.glob(["AdvancedExample/app/src/main/res/**"]),
    )
<commit_msg>Internal change<commit_after>"""
Provides a build rule for building the AdvancedExample from Blaze.
"""

load("//java/com/google/ads/interactivemedia/v3/samples:build_defs.bzl", "build_sample_package")

COMMON_DEPS = [
    "//third_party/java/android_libs/exoplayer:exoplayer2_core",
    "//third_party/java/android_libs/exoplayer:exoplayer2_dash",
    "//third_party/java/android_libs/exoplayer:exoplayer2_ext_mediasession",
    "//third_party/java/android_libs/exoplayer:exoplayer2_hls",
    "//third_party/java/android_libs/exoplayer:exoplayer2_ui",
    "//third_party/java/android_libs/guava_jdk5:collect",
    "//third_party/java/androidx/annotation",
    "//third_party/java/android/android_sdk_linux/extras/android/compatibility/v4",
    "//third_party/java/androidx/appcompat",
    "//third_party/java/androidx/constraintlayout",
    "//java/com/google/android/gmscore/integ/client/base",
    "//java/com/google/android/gmscore/integ/client/cast/framework",
    "//java/com/google/android/gmscore/integ/client/cast",
]

DEBUG_DEPS = COMMON_DEPS + [
    "//java/com/google/ads/interactivemedia/v3:sdk_lib_debug",
]

COMPILED_DEPS = COMMON_DEPS + [
    "//java/com/google/ads/interactivemedia/v3:sdk_lib",
    "//java/com/google/android/gmscore/integ/client/ads_identifier",
]

def advanced_example_package():
    build_sample_package(
        name_prefix = "advanced",
        package_name = "com.google.ads.interactivemedia.v3.samples.videoplayerapp",
        srcs = native.glob([
            "AdvancedExample/app/src/main/java/com/google/ads/interactivemedia/v3/samples/samplevideoplayer/*.java",
            "AdvancedExample/app/src/main/java/com/google/ads/interactivemedia/v3/samples/videoplayerapp/*.java",
        ]),
        compiled_deps = COMPILED_DEPS,
        debug_deps = COMPILED_DEPS,  # We can't use real debug deps here due to collisions between AdShield and GMS Core dependencies.
        manifest = "AdvancedExample/app/src/main/AndroidManifest.xml",
        resources = native.glob(["AdvancedExample/app/src/main/res/**"]),
    )
<|endoftext|>


<commit_before># This file is part of beets.
# Copyright 2013, Jakob Schnitzer.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Converts tracks or albums to external directory
"""
import logging
import os
import threading
from subprocess import Popen

from beets.plugins import BeetsPlugin
from beets import ui, util
from beetsplug.embedart import _embed
from beets import config

log = logging.getLogger('beets')
DEVNULL = open(os.devnull, 'wb')
_fs_lock = threading.Lock()

def _destination(lib, dest_dir, item, keep_new):
    """Return the path under `dest_dir` where the file should be placed
    (possibly after conversion).
    """
    dest = lib.destination(item, basedir=dest_dir)
    if keep_new:
        # When we're keeping the converted file, no extension munging
        # occurs.
        return dest
    else:
        # Otherwise, replace the extension with .mp3.
        return os.path.splitext(dest)[0] + '.mp3'

def encode(source, dest):
    log.info(u'Started encoding {0}'.format(util.displayable_path(source)))

    opts = config['convert']['opts'].get(unicode).split(u' ')
    encode = Popen([config['convert']['ffmpeg'].get(unicode), '-i', source] +
                   opts + [dest],
                   close_fds=True, stderr=DEVNULL)
    encode.wait()
    if encode.returncode != 0:
        # Something went wrong (probably Ctrl+C), remove temporary files
        log.info(u'Encoding {0} failed. Cleaning up...'
                 .format(util.displayable_path(source)))
        util.remove(dest)
        util.prune_dirs(os.path.dirname(dest))
        return
    log.info(u'Finished encoding {0}'.format(util.displayable_path(source)))

def convert_item(lib, dest_dir, keep_new):
    while True:
        item = yield
        dest = _destination(lib, dest_dir, item, keep_new)

        if os.path.exists(util.syspath(dest)):
            log.info(u'Skipping {0} (target file exists)'.format(
                util.displayable_path(item.path)
            ))
            continue

        # Ensure that only one thread tries to create directories at a
        # time. (The existence check is not atomic with the directory
        # creation inside this function.)
        with _fs_lock:
            util.mkdirall(dest)

        # When keeping the new file in the library, we first move the
        # current (pristine) file to the destination. We'll then copy it
        # back to its old path or transcode it to a new path.
        if keep_new:
            log.info(u'Moving to {0}'.
                     format(util.displayable_path(dest)))
            util.move(item.path, dest)

        maxbr = config['convert']['max_bitrate'].get(int)
        if item.format == 'MP3' and item.bitrate < 1000 * maxbr:
            # No transcoding necessary.
            log.info(u'Copying {0}'.format(util.displayable_path(item.path)))
            if keep_new:
                util.copy(dest, item.path)
            else:
                util.copy(item.path, dest)

        else:
            if keep_new:
                item.path = os.path.splitext(item.path)[0] + '.mp3'
                encode(dest, item.path)
            else:
                encode(item.path, dest)

        # Write tags from the database to the converted file.
        if not keep_new:
            item.path = dest
        item.write()

        # If we're keeping the transcoded file, read it again (after
        # writing) to get new bitrate, duration, etc.
        if keep_new:
            item.read()
            lib.store(item)  # Store new path and audio data.

        if config['convert']['embed']:
            album = lib.get_album(item)
            if album:
                artpath = album.artpath
                if artpath:
                    _embed(artpath, [item])

def convert_func(lib, opts, args):
    dest = opts.dest if opts.dest is not None else \
            config['convert']['dest'].get()
    if not dest:
        raise ui.UserError('no convert destination set')
    threads = opts.threads if opts.threads is not None else \
            config['convert']['threads'].get(int)
    keep_new = opts.keep_new

    ui.commands.list_items(lib, ui.decargs(args), opts.album, None)

    if not ui.input_yn("Convert? (Y/n)"):
        return

    if opts.album:
        items = (i for a in lib.albums(ui.decargs(args)) for i in a.items())
    else:
        items = lib.items(ui.decargs(args))
    convert = [convert_item(lib, dest, keep_new) for i in range(threads)]
    pipe = util.pipeline.Pipeline([items, convert])
    pipe.run_parallel()

class ConvertPlugin(BeetsPlugin):
    def __init__(self):
        super(ConvertPlugin, self).__init__()
        self.config.add({
            u'dest': None,
            u'threads': util.cpu_count(),
            u'ffmpeg': u'ffmpeg',
            u'opts': u'-aq 2',
            u'max_bitrate': 500,
            u'embed': True,
        })

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-a', '--album', action='store_true',
                              help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                              help='change the number of threads, \
                              defaults to maximum availble processors ')
        cmd.parser.add_option('-k', '--keep-new', action='store_true',
                              dest='keep_new', help='keep only the converted \
                              and move the old files')
        cmd.parser.add_option('-d', '--dest', action='store',
                              help='set the destination directory')
        cmd.func = convert_func
        return [cmd]
<commit_msg>convert: fix unicode error in path construction<commit_after># This file is part of beets.
# Copyright 2013, Jakob Schnitzer.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Converts tracks or albums to external directory
"""
import logging
import os
import threading
from subprocess import Popen

from beets.plugins import BeetsPlugin
from beets import ui, util
from beetsplug.embedart import _embed
from beets import config

log = logging.getLogger('beets')
DEVNULL = open(os.devnull, 'wb')
_fs_lock = threading.Lock()

def _destination(lib, dest_dir, item, keep_new):
    """Return the path under `dest_dir` where the file should be placed
    (possibly after conversion).
    """
    dest = lib.destination(item, basedir=dest_dir)
    if keep_new:
        # When we're keeping the converted file, no extension munging
        # occurs.
        return dest
    else:
        # Otherwise, replace the extension with .mp3.
        return os.path.splitext(dest)[0] + '.mp3'

def encode(source, dest):
    log.info(u'Started encoding {0}'.format(util.displayable_path(source)))

    opts = config['convert']['opts'].get(unicode).split(u' ')
    encode = Popen([config['convert']['ffmpeg'].get(unicode), '-i', source] +
                   opts + [dest],
                   close_fds=True, stderr=DEVNULL)
    encode.wait()
    if encode.returncode != 0:
        # Something went wrong (probably Ctrl+C), remove temporary files
        log.info(u'Encoding {0} failed. Cleaning up...'
                 .format(util.displayable_path(source)))
        util.remove(dest)
        util.prune_dirs(os.path.dirname(dest))
        return
    log.info(u'Finished encoding {0}'.format(util.displayable_path(source)))

def convert_item(lib, dest_dir, keep_new):
    while True:
        item = yield
        dest = _destination(lib, dest_dir, item, keep_new)

        if os.path.exists(util.syspath(dest)):
            log.info(u'Skipping {0} (target file exists)'.format(
                util.displayable_path(item.path)
            ))
            continue

        # Ensure that only one thread tries to create directories at a
        # time. (The existence check is not atomic with the directory
        # creation inside this function.)
        with _fs_lock:
            util.mkdirall(dest)

        # When keeping the new file in the library, we first move the
        # current (pristine) file to the destination. We'll then copy it
        # back to its old path or transcode it to a new path.
        if keep_new:
            log.info(u'Moving to {0}'.
                     format(util.displayable_path(dest)))
            util.move(item.path, dest)

        maxbr = config['convert']['max_bitrate'].get(int)
        if item.format == 'MP3' and item.bitrate < 1000 * maxbr:
            # No transcoding necessary.
            log.info(u'Copying {0}'.format(util.displayable_path(item.path)))
            if keep_new:
                util.copy(dest, item.path)
            else:
                util.copy(item.path, dest)

        else:
            if keep_new:
                item.path = os.path.splitext(item.path)[0] + '.mp3'
                encode(dest, item.path)
            else:
                encode(item.path, dest)

        # Write tags from the database to the converted file.
        if not keep_new:
            item.path = dest
        item.write()

        # If we're keeping the transcoded file, read it again (after
        # writing) to get new bitrate, duration, etc.
        if keep_new:
            item.read()
            lib.store(item)  # Store new path and audio data.

        if config['convert']['embed']:
            album = lib.get_album(item)
            if album:
                artpath = album.artpath
                if artpath:
                    _embed(artpath, [item])

def convert_func(lib, opts, args):
    dest = opts.dest if opts.dest is not None else \
            config['convert']['dest'].get()
    if not dest:
        raise ui.UserError('no convert destination set')
    dest = util.bytestring_path(dest)
    threads = opts.threads if opts.threads is not None else \
            config['convert']['threads'].get(int)
    keep_new = opts.keep_new

    ui.commands.list_items(lib, ui.decargs(args), opts.album, None)

    if not ui.input_yn("Convert? (Y/n)"):
        return

    if opts.album:
        items = (i for a in lib.albums(ui.decargs(args)) for i in a.items())
    else:
        items = lib.items(ui.decargs(args))
    convert = [convert_item(lib, dest, keep_new) for i in range(threads)]
    pipe = util.pipeline.Pipeline([items, convert])
    pipe.run_parallel()

class ConvertPlugin(BeetsPlugin):
    def __init__(self):
        super(ConvertPlugin, self).__init__()
        self.config.add({
            u'dest': None,
            u'threads': util.cpu_count(),
            u'ffmpeg': u'ffmpeg',
            u'opts': u'-aq 2',
            u'max_bitrate': 500,
            u'embed': True,
        })

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-a', '--album', action='store_true',
                              help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                              help='change the number of threads, \
                              defaults to maximum availble processors ')
        cmd.parser.add_option('-k', '--keep-new', action='store_true',
                              dest='keep_new', help='keep only the converted \
                              and move the old files')
        cmd.parser.add_option('-d', '--dest', action='store',
                              help='set the destination directory')
        cmd.func = convert_func
        return [cmd]
<|endoftext|>


<commit_before># Copyright (c) 2016 Duke Center for Genomic and Computational Biology
#
# 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
import logging
import os
from pipes import quote
import subprocess
import time
import math
from Queue import Queue, Empty
from threading import Thread

from toil.batchSystems import MemoryString
from toil.batchSystems.abstractBatchSystem import BatchSystemSupport

logger = logging.getLogger(__name__)

sleepSeconds = 1

class Worker(Thread):
    def __init__(self, newJobsQueue, updatedJobsQueue, killQueue, killedJobsQueue, boss):
        Thread.__init__(self)
        self.newJobsQueue = newJobsQueue
        self.updatedJobsQueue = updatedJobsQueue
        self.killQueue = killQueue
        self.killedJobsQueue = killedJobsQueue
        self.waitingJobs = list()
        self.runningJobs = set()
        self.boss = boss
        self.allocatedCpus = dict()
        self.slurmJobIDs = dict()

    def parse_elapsed(self, elapsed):
        # slurm returns elapsed time in days-hours:minutes:seconds format
        # Sometimes it will only return minutes:seconds, so days may be omitted
        # For ease of calculating, we'll make sure all the delimeters are ':'
        # Then reverse the list so that we're always counting up from seconds -> minutes -> hours -> days
        total_seconds = 0
        try:
            elapsed = elapsed.replace('-', ':').split(':')
            elapsed.reverse()
            seconds_per_unit = [1, 60, 3600, 86400]
            for index, multiplier in enumerate(seconds_per_unit):
                if index < len(elapsed):
                    total_seconds += multiplier * int(elapsed[index])
        except ValueError:
            pass  # slurm may return INVALID instead of a time
        return total_seconds

    def getRunningJobIDs(self):
        # Should return a dictionary of Job IDs and number of seconds
        times = {}
        currentjobs = dict((str(self.slurmJobIDs[x]), x) for x in self.runningJobs)
        # currentjobs is a dictionary that maps a slurm job id (string) to our own internal job id
        # squeue arguments:
        # -h for no header
        # --format to get jobid i, state %t and time days-hours:minutes:seconds

        lines = subprocess.check_output(['squeue', '-h', '--format', '%i %t %M']).split('\n')
        for line in lines:
            values = line.split()
            if len(values) < 3:
                continue
            slurm_jobid, state, elapsed_time = values
            if slurm_jobid in currentjobs and state == 'R':
                seconds_running = self.parse_elapsed(elapsed_time)
                times[currentjobs[slurm_jobid]] = seconds_running

        return times

    def getSlurmID(self, jobID):
        if not jobID in self.slurmJobIDs:
            RuntimeError("Unknown jobID, could not be converted")

        job = self.slurmJobIDs[jobID]
        return str(job)

    def forgetJob(self, jobID):
        self.runningJobs.remove(jobID)
        del self.allocatedCpus[jobID]
        del self.slurmJobIDs[jobID]

    def killJobs(self):
        # Load hit list:
        killList = list()
        while True:
            try:
                jobId = self.killQueue.get(block=False)
            except Empty:
                break
            else:
                killList.append(jobId)

        if not killList:
            return False

        # Do the dirty job
        for jobID in list(killList):
            if jobID in self.runningJobs:
                logger.debug('Killing job: %s', jobID)
                subprocess.check_call(['scancel', self.getSlurmID(jobID)])
            else:
                if jobID in self.waitingJobs:
                    self.waitingJobs.remove(jobID)
                self.killedJobsQueue.put(jobID)
                killList.remove(jobID)

        # Wait to confirm the kill
        while killList:
            for jobID in list(killList):
                if self.getJobExitCode(self.slurmJobIDs[jobID]) is not None:
                    logger.debug('Adding jobID %s to killedJobsQueue', jobID)
                    self.killedJobsQueue.put(jobID)
                    killList.remove(jobID)
                    self.forgetJob(jobID)
            if len(killList) > 0:
                logger.warn("Some jobs weren't killed, trying again in %is.", sleepSeconds)
                time.sleep(sleepSeconds)

        return True

    def createJobs(self, newJob):
        activity = False
        # Load new job id if present:
        if newJob is not None:
            self.waitingJobs.append(newJob)
        # Launch jobs as necessary:
        while (len(self.waitingJobs) > 0
               and sum(self.allocatedCpus.values()) < int(self.boss.maxCores)):
            activity = True
            jobID, cpu, memory, command = self.waitingJobs.pop(0)
            sbatch_line = self.prepareSbatch(cpu, memory, jobID) + ['--wrap={}'.format(command)]
            slurmJobID = self.sbatch(sbatch_line)
            self.slurmJobIDs[jobID] = slurmJobID
            self.runningJobs.add(jobID)
            self.allocatedCpus[jobID] = cpu
        return activity

    def checkOnJobs(self):
        activity = False
        logger.debug('List of running jobs: %r', self.runningJobs)
        for jobID in list(self.runningJobs):
            logger.debug("Checking status of internal job id %d", jobID)
            status = self.getJobExitCode(self.slurmJobIDs[jobID])
            if status is not None:
                activity = True
                self.updatedJobsQueue.put((jobID, status))
                self.forgetJob(jobID)
        return activity

    def run(self):
        while True:
            activity = False
            newJob = None
            if not self.newJobsQueue.empty():
                activity = True
                newJob = self.newJobsQueue.get()
                if newJob is None:
                    logger.debug('Received queue sentinel.')
                    break
            activity |= self.killJobs()
            activity |= self.createJobs(newJob)
            activity |= self.checkOnJobs()
            if not activity:
                logger.debug('No activity, sleeping for %is', sleepSeconds)
                time.sleep(sleepSeconds)

    def prepareSbatch(self, cpu, mem, jobID):
        #  Returns the sbatch command line before the script to run
        sbatch_line = ['sbatch', '-Q', '-J', 'toil_job_{}'.format(jobID)]

        if self.boss.environment:
            for k, v in self.boss.environment.iteritems():
                quoted_value = quote(os.environ[k] if v is None else v)
                sbatch_line.append('--export={}={}'.format(k, quoted_value))

        if mem is not None:
            # memory passed in is in bytes, but slurm expects megabytes
            sbatch_line.append('--mem={}'.format(int(mem) / 2 ** 20))
        if cpu is not None:
            sbatch_line.append('--cpus-per-task={}'.format(int(math.ceil(cpu))))

        return sbatch_line

    def sbatch(self, sbatch_line):
        logger.debug("Running %r", sbatch_line)
        try:
            output = subprocess.check_output(sbatch_line, stderr=subprocess.STDOUT)
            # sbatch prints a line like 'Submitted batch job 2954103'
            result = int(output.strip().split()[-1])
            logger.debug("sbatch submitted job %d", result)
            return result
        except subprocess.CalledProcessError as e:
            logger.error("sbatch command failed with code %d: %s", e.returncode, e.output)
            raise e
        except OSError as e:
            logger.error("sbatch command failed")
            raise e

    def getJobExitCode(self, slurmJobID):
        logger.debug("Getting exit code for slurm job %d", slurmJobID)
        # SLURM job exit codes are obtained by running sacct.
        # sacct returns
        # -n : no header
        # -j : job
        # --format : specify output columns
        # -P : separate columns with pipes
        # -S 1970-01-01 override start time limit
        args = ['sacct',
                '-n',
                '-j', str(slurmJobID),
                '--format', 'State,ExitCode',
                '-P',
                '-S', '1970-01-01']
        process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        for line in process.stdout:
            values = line.strip().split('|')
            if len(values) < 2:
                continue
            state, exitcode = values
            logger.debug("sacct job state is %s", state)
            # If Job is in a running state, return None to indicate we don't have an update
            if state in ('PENDING', 'RUNNING', 'CONFIGURING', 'COMPLETING', 'RESIZING', 'SUSPENDED'):
                return None
            status, _ = exitcode.split(':')
            logger.debug("sacct exit code is %s, returning status %s", exitcode, status)
            return int(status)
        logger.debug("Did not find exit code for job in sacct output")
        return None

class SlurmBatchSystem(BatchSystemSupport):
    """
    The interface for SLURM
    """

    @classmethod
    def supportsWorkerCleanup(cls):
        return False

    @classmethod
    def supportsHotDeployment(cls):
        return False

    def __init__(self, config, maxCores, maxMemory, maxDisk):
        super(SlurmBatchSystem, self).__init__(config, maxCores, maxMemory, maxDisk)
        self.slurmResultsFile = self._getResultsFileName(config.jobStore)
        # Reset the job queue and results (initially, we do this again once we've killed the jobs)
        self.slurmResultsFileHandle = open(self.slurmResultsFile, 'w')
        # We lose any previous state in this file, and ensure the files existence
        self.slurmResultsFileHandle.close()
        self.currentJobs = set()
        self.maxCPU, self.maxMEM = self.obtainSystemConstants()
        self.nextJobID = 0
        self.newJobsQueue = Queue()
        self.updatedJobsQueue = Queue()
        self.killQueue = Queue()
        self.killedJobsQueue = Queue()
        self.worker = Worker(self.newJobsQueue, self.updatedJobsQueue, self.killQueue,
                             self.killedJobsQueue, self)
        self.worker.start()

    def __des__(self):
        # Closes the file handle associated with the results file.
        self.slurmResultsFileHandle.close()

    def issueBatchJob(self, command, memory, cores, disk, preemptable):
        self.checkResourceRequest(memory, cores, disk)
        jobID = self.nextJobID
        self.nextJobID += 1
        self.currentJobs.add(jobID)
        self.newJobsQueue.put((jobID, cores, memory, command))
        logger.debug("Issued the job command: %s with job id: %s ", command, str(jobID))
        return jobID

    def killBatchJobs(self, jobIDs):
        """
        Kills the given jobs, represented as Job ids, then checks they are dead by checking
        they are not in the list of issued jobs.
        """
        jobIDs = set(jobIDs)
        logger.debug('Jobs to be killed: %r', jobIDs)
        for jobID in jobIDs:
            self.killQueue.put(jobID)
        while jobIDs:
            killedJobId = self.killedJobsQueue.get()
            if killedJobId is None:
                break
            jobIDs.remove(killedJobId)
            if killedJobId in self.currentJobs:
                self.currentJobs.remove(killedJobId)
            if jobIDs:
                logger.debug('Some kills (%s) still pending, sleeping %is', len(jobIDs),
                             sleepSeconds)
                time.sleep(sleepSeconds)

    def getIssuedBatchJobIDs(self):
        """
        Gets the list of jobs issued to SLURM.
        """
        return list(self.currentJobs)

    def getRunningBatchJobIDs(self):
        return self.worker.getRunningJobIDs()

    def getUpdatedBatchJob(self, maxWait):
        try:
            i = self.updatedJobsQueue.get(timeout=maxWait)
        except Empty:
            return None
        logger.debug('UpdatedJobsQueue Item: %s', i)
        jobID, retcode = i
        self.updatedJobsQueue.task_done()
        self.currentJobs.remove(jobID)
        return jobID, retcode, None

    def shutdown(self):
        """
        Signals worker to shutdown (via sentinel) then cleanly joins the thread
        """
        newJobsQueue = self.newJobsQueue
        self.newJobsQueue = None

        newJobsQueue.put(None)
        self.worker.join()

    def getWaitDuration(self):
        return 1.0

    @classmethod
    def getRescueBatchJobFrequency(cls):
        return 1800  # Half an hour

    @staticmethod
    def obtainSystemConstants():
        # sinfo -Ne --format '%m,%c'
        # sinfo arguments:
        # -N for node-oriented
        # -h for no header
        # -e for exact values (e.g. don't return 32+)
        # --format to get memory, cpu
        max_cpu = 0
        max_mem = MemoryString('0')
        lines = subprocess.check_output(['sinfo', '-Nhe', '--format', '%m %c']).split('\n')
        for line in lines:
            values = line.split()
            if len(values) < 2:
                continue
            mem, cpu = values
            max_cpu = max(max_cpu, int(cpu))
            max_mem = max(max_mem, MemoryString(mem + 'M'))
        if max_cpu == 0 or max_mem.byteVal() == 0:
            RuntimeError('sinfo did not return memory or cpu info')
        return max_cpu, max_mem
<commit_msg>SLURM: Eliminate unnecessary queue method call<commit_after># Copyright (c) 2016 Duke Center for Genomic and Computational Biology
#
# 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
import logging
import os
from pipes import quote
import subprocess
import time
import math
from Queue import Queue, Empty
from threading import Thread

from toil.batchSystems import MemoryString
from toil.batchSystems.abstractBatchSystem import BatchSystemSupport

logger = logging.getLogger(__name__)

sleepSeconds = 1

class Worker(Thread):
    def __init__(self, newJobsQueue, updatedJobsQueue, killQueue, killedJobsQueue, boss):
        Thread.__init__(self)
        self.newJobsQueue = newJobsQueue
        self.updatedJobsQueue = updatedJobsQueue
        self.killQueue = killQueue
        self.killedJobsQueue = killedJobsQueue
        self.waitingJobs = list()
        self.runningJobs = set()
        self.boss = boss
        self.allocatedCpus = dict()
        self.slurmJobIDs = dict()

    def parse_elapsed(self, elapsed):
        # slurm returns elapsed time in days-hours:minutes:seconds format
        # Sometimes it will only return minutes:seconds, so days may be omitted
        # For ease of calculating, we'll make sure all the delimeters are ':'
        # Then reverse the list so that we're always counting up from seconds -> minutes -> hours -> days
        total_seconds = 0
        try:
            elapsed = elapsed.replace('-', ':').split(':')
            elapsed.reverse()
            seconds_per_unit = [1, 60, 3600, 86400]
            for index, multiplier in enumerate(seconds_per_unit):
                if index < len(elapsed):
                    total_seconds += multiplier * int(elapsed[index])
        except ValueError:
            pass  # slurm may return INVALID instead of a time
        return total_seconds

    def getRunningJobIDs(self):
        # Should return a dictionary of Job IDs and number of seconds
        times = {}
        currentjobs = dict((str(self.slurmJobIDs[x]), x) for x in self.runningJobs)
        # currentjobs is a dictionary that maps a slurm job id (string) to our own internal job id
        # squeue arguments:
        # -h for no header
        # --format to get jobid i, state %t and time days-hours:minutes:seconds

        lines = subprocess.check_output(['squeue', '-h', '--format', '%i %t %M']).split('\n')
        for line in lines:
            values = line.split()
            if len(values) < 3:
                continue
            slurm_jobid, state, elapsed_time = values
            if slurm_jobid in currentjobs and state == 'R':
                seconds_running = self.parse_elapsed(elapsed_time)
                times[currentjobs[slurm_jobid]] = seconds_running

        return times

    def getSlurmID(self, jobID):
        if not jobID in self.slurmJobIDs:
            RuntimeError("Unknown jobID, could not be converted")

        job = self.slurmJobIDs[jobID]
        return str(job)

    def forgetJob(self, jobID):
        self.runningJobs.remove(jobID)
        del self.allocatedCpus[jobID]
        del self.slurmJobIDs[jobID]

    def killJobs(self):
        # Load hit list:
        killList = list()
        while True:
            try:
                jobId = self.killQueue.get(block=False)
            except Empty:
                break
            else:
                killList.append(jobId)

        if not killList:
            return False

        # Do the dirty job
        for jobID in list(killList):
            if jobID in self.runningJobs:
                logger.debug('Killing job: %s', jobID)
                subprocess.check_call(['scancel', self.getSlurmID(jobID)])
            else:
                if jobID in self.waitingJobs:
                    self.waitingJobs.remove(jobID)
                self.killedJobsQueue.put(jobID)
                killList.remove(jobID)

        # Wait to confirm the kill
        while killList:
            for jobID in list(killList):
                if self.getJobExitCode(self.slurmJobIDs[jobID]) is not None:
                    logger.debug('Adding jobID %s to killedJobsQueue', jobID)
                    self.killedJobsQueue.put(jobID)
                    killList.remove(jobID)
                    self.forgetJob(jobID)
            if len(killList) > 0:
                logger.warn("Some jobs weren't killed, trying again in %is.", sleepSeconds)
                time.sleep(sleepSeconds)

        return True

    def createJobs(self, newJob):
        activity = False
        # Load new job id if present:
        if newJob is not None:
            self.waitingJobs.append(newJob)
        # Launch jobs as necessary:
        while (len(self.waitingJobs) > 0
               and sum(self.allocatedCpus.values()) < int(self.boss.maxCores)):
            activity = True
            jobID, cpu, memory, command = self.waitingJobs.pop(0)
            sbatch_line = self.prepareSbatch(cpu, memory, jobID) + ['--wrap={}'.format(command)]
            slurmJobID = self.sbatch(sbatch_line)
            self.slurmJobIDs[jobID] = slurmJobID
            self.runningJobs.add(jobID)
            self.allocatedCpus[jobID] = cpu
        return activity

    def checkOnJobs(self):
        activity = False
        logger.debug('List of running jobs: %r', self.runningJobs)
        for jobID in list(self.runningJobs):
            logger.debug("Checking status of internal job id %d", jobID)
            status = self.getJobExitCode(self.slurmJobIDs[jobID])
            if status is not None:
                activity = True
                self.updatedJobsQueue.put((jobID, status))
                self.forgetJob(jobID)
        return activity

    def run(self):
        while True:
            activity = False
            newJob = None
            if not self.newJobsQueue.empty():
                activity = True
                newJob = self.newJobsQueue.get()
                if newJob is None:
                    logger.debug('Received queue sentinel.')
                    break
            activity |= self.killJobs()
            activity |= self.createJobs(newJob)
            activity |= self.checkOnJobs()
            if not activity:
                logger.debug('No activity, sleeping for %is', sleepSeconds)
                time.sleep(sleepSeconds)

    def prepareSbatch(self, cpu, mem, jobID):
        #  Returns the sbatch command line before the script to run
        sbatch_line = ['sbatch', '-Q', '-J', 'toil_job_{}'.format(jobID)]

        if self.boss.environment:
            for k, v in self.boss.environment.iteritems():
                quoted_value = quote(os.environ[k] if v is None else v)
                sbatch_line.append('--export={}={}'.format(k, quoted_value))

        if mem is not None:
            # memory passed in is in bytes, but slurm expects megabytes
            sbatch_line.append('--mem={}'.format(int(mem) / 2 ** 20))
        if cpu is not None:
            sbatch_line.append('--cpus-per-task={}'.format(int(math.ceil(cpu))))

        return sbatch_line

    def sbatch(self, sbatch_line):
        logger.debug("Running %r", sbatch_line)
        try:
            output = subprocess.check_output(sbatch_line, stderr=subprocess.STDOUT)
            # sbatch prints a line like 'Submitted batch job 2954103'
            result = int(output.strip().split()[-1])
            logger.debug("sbatch submitted job %d", result)
            return result
        except subprocess.CalledProcessError as e:
            logger.error("sbatch command failed with code %d: %s", e.returncode, e.output)
            raise e
        except OSError as e:
            logger.error("sbatch command failed")
            raise e

    def getJobExitCode(self, slurmJobID):
        logger.debug("Getting exit code for slurm job %d", slurmJobID)
        # SLURM job exit codes are obtained by running sacct.
        # sacct returns
        # -n : no header
        # -j : job
        # --format : specify output columns
        # -P : separate columns with pipes
        # -S 1970-01-01 override start time limit
        args = ['sacct',
                '-n',
                '-j', str(slurmJobID),
                '--format', 'State,ExitCode',
                '-P',
                '-S', '1970-01-01']
        process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        for line in process.stdout:
            values = line.strip().split('|')
            if len(values) < 2:
                continue
            state, exitcode = values
            logger.debug("sacct job state is %s", state)
            # If Job is in a running state, return None to indicate we don't have an update
            if state in ('PENDING', 'RUNNING', 'CONFIGURING', 'COMPLETING', 'RESIZING', 'SUSPENDED'):
                return None
            status, _ = exitcode.split(':')
            logger.debug("sacct exit code is %s, returning status %s", exitcode, status)
            return int(status)
        logger.debug("Did not find exit code for job in sacct output")
        return None

class SlurmBatchSystem(BatchSystemSupport):
    """
    The interface for SLURM
    """

    @classmethod
    def supportsWorkerCleanup(cls):
        return False

    @classmethod
    def supportsHotDeployment(cls):
        return False

    def __init__(self, config, maxCores, maxMemory, maxDisk):
        super(SlurmBatchSystem, self).__init__(config, maxCores, maxMemory, maxDisk)
        self.slurmResultsFile = self._getResultsFileName(config.jobStore)
        # Reset the job queue and results (initially, we do this again once we've killed the jobs)
        self.slurmResultsFileHandle = open(self.slurmResultsFile, 'w')
        # We lose any previous state in this file, and ensure the files existence
        self.slurmResultsFileHandle.close()
        self.currentJobs = set()
        self.maxCPU, self.maxMEM = self.obtainSystemConstants()
        self.nextJobID = 0
        self.newJobsQueue = Queue()
        self.updatedJobsQueue = Queue()
        self.killQueue = Queue()
        self.killedJobsQueue = Queue()
        self.worker = Worker(self.newJobsQueue, self.updatedJobsQueue, self.killQueue,
                             self.killedJobsQueue, self)
        self.worker.start()

    def __des__(self):
        # Closes the file handle associated with the results file.
        self.slurmResultsFileHandle.close()

    def issueBatchJob(self, command, memory, cores, disk, preemptable):
        self.checkResourceRequest(memory, cores, disk)
        jobID = self.nextJobID
        self.nextJobID += 1
        self.currentJobs.add(jobID)
        self.newJobsQueue.put((jobID, cores, memory, command))
        logger.debug("Issued the job command: %s with job id: %s ", command, str(jobID))
        return jobID

    def killBatchJobs(self, jobIDs):
        """
        Kills the given jobs, represented as Job ids, then checks they are dead by checking
        they are not in the list of issued jobs.
        """
        jobIDs = set(jobIDs)
        logger.debug('Jobs to be killed: %r', jobIDs)
        for jobID in jobIDs:
            self.killQueue.put(jobID)
        while jobIDs:
            killedJobId = self.killedJobsQueue.get()
            if killedJobId is None:
                break
            jobIDs.remove(killedJobId)
            if killedJobId in self.currentJobs:
                self.currentJobs.remove(killedJobId)
            if jobIDs:
                logger.debug('Some kills (%s) still pending, sleeping %is', len(jobIDs),
                             sleepSeconds)
                time.sleep(sleepSeconds)

    def getIssuedBatchJobIDs(self):
        """
        Gets the list of jobs issued to SLURM.
        """
        return list(self.currentJobs)

    def getRunningBatchJobIDs(self):
        return self.worker.getRunningJobIDs()

    def getUpdatedBatchJob(self, maxWait):
        try:
            i = self.updatedJobsQueue.get(timeout=maxWait)
        except Empty:
            return None
        logger.debug('UpdatedJobsQueue Item: %s', i)
        jobID, retcode = i
        self.currentJobs.remove(jobID)
        return jobID, retcode, None

    def shutdown(self):
        """
        Signals worker to shutdown (via sentinel) then cleanly joins the thread
        """
        newJobsQueue = self.newJobsQueue
        self.newJobsQueue = None

        newJobsQueue.put(None)
        self.worker.join()

    def getWaitDuration(self):
        return 1.0

    @classmethod
    def getRescueBatchJobFrequency(cls):
        return 1800  # Half an hour

    @staticmethod
    def obtainSystemConstants():
        # sinfo -Ne --format '%m,%c'
        # sinfo arguments:
        # -N for node-oriented
        # -h for no header
        # -e for exact values (e.g. don't return 32+)
        # --format to get memory, cpu
        max_cpu = 0
        max_mem = MemoryString('0')
        lines = subprocess.check_output(['sinfo', '-Nhe', '--format', '%m %c']).split('\n')
        for line in lines:
            values = line.split()
            if len(values) < 2:
                continue
            mem, cpu = values
            max_cpu = max(max_cpu, int(cpu))
            max_mem = max(max_mem, MemoryString(mem + 'M'))
        if max_cpu == 0 or max_mem.byteVal() == 0:
            RuntimeError('sinfo did not return memory or cpu info')
        return max_cpu, max_mem
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2016 GEM Foundation
#
# OpenQuake 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.
#
# OpenQuake 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 OpenQuake. If not, see <http://www.gnu.org/licenses/>.

import os
import getpass
import subprocess
import unittest
import tempfile
import mock

from django.db import connection
from django.core import exceptions
from openquake.server.db import models, actions, upgrade_manager
from openquake.server.settings import DATABASE
from openquake.server.tests import helpers

def setup_module():
    global tmpfile
    fh, tmpfile = tempfile.mkstemp()
    os.close(fh)
    DATABASE['name'] = tmpfile
    connection.cursor()  # connect to the db
    print 'Using the database', tmpfile
    upgrade_manager.upgrade_db(connection.connection)

def teardown_module():
    connection.close()

def get_job(cfg, username, hazard_calculation_id=None):
    job_id, oq = actions.job_from_file(cfg, username, hazard_calculation_id)
    return models.OqJob.objects.get(pk=job_id)

class FakeJob(object):
    def __init__(self, job_type, calculation_mode):
        self.job_type = job_type
        self.calculation_mode = calculation_mode

    def get_oqparam(self):
        m = mock.Mock()
        m.calculation_mode = self.calculation_mode
        return m

class CheckHazardRiskConsistencyTestCase(unittest.TestCase):
    def test_ok(self):
        haz_job = FakeJob('hazard', 'scenario')
        actions.check_hazard_risk_consistency(
            haz_job, 'scenario_risk')

    def test_obsolete_mode(self):
        haz_job = FakeJob('hazard', 'scenario')
        with self.assertRaises(ValueError) as ctx:
            actions.check_hazard_risk_consistency(
                haz_job, 'scenario')
        msg = str(ctx.exception)
        self.assertEqual(msg, 'Please change calculation_mode=scenario into '
                         'scenario_risk in the .ini file')

    def test_inconsistent_mode(self):
        haz_job = FakeJob('hazard', 'scenario')
        with self.assertRaises(actions.InvalidCalculationID) as ctx:
            actions.check_hazard_risk_consistency(
                haz_job, 'classical_risk')
        msg = str(ctx.exception)
        self.assertEqual(
            msg, "In order to run a risk calculation of kind "
            "'classical_risk', you need to provide a "
            "calculation of kind ['classical', 'classical_risk'], "
            "but you provided a 'scenario' instead")

class JobFromFileTestCase(unittest.TestCase):

    def test_create_job_default_user(self):
        job = actions.create_job('classical', 'test_create_job_default_user')

        self.assertEqual('openquake', job.user_name)
        self.assertEqual('executing', job.status)

        # Check the make sure it's in the database.
        try:
            models.OqJob.objects.get(id=job.id)
        except exceptions.ObjectDoesNotExist:
            self.fail('Job was not found in the database')

    def test_create_job_specified_user(self):
        user_name = helpers.random_string()
        job = actions.create_job(
            'classical', 'test_create_job_specified_user',
            user_name=user_name)

        self.assertEqual(user_name, job.user_name)
        self.assertEqual('executing', job.status)

        try:
            models.OqJob.objects.get(id=job.id)
        except exceptions.ObjectDoesNotExist:
            self.fail('Job was not found in the database')

class OpenquakeCliTestCase(unittest.TestCase):
    """
    Run "oq-engine --version" as a separate process using `subprocess`.
    """

    def test_run_version(self):
        args = [helpers.RUNNER, "--version"]

        print 'Running:', ' '.join(args)  # this is useful for debugging
        return subprocess.check_call(args)

class DeleteHazCalcTestCase(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.hazard_cfg = helpers.get_data_path(
            'simple_fault_demo_hazard/job.ini')
        cls.risk_cfg = helpers.get_data_path(
            'classical_psha_based_risk/job.ini')
        cls.job = get_job(cls.hazard_cfg, getpass.getuser())

    def test_del_calc(self):
        hazard_job = get_job(self.hazard_cfg, getpass.getuser())

        models.Output.objects.create_output(
            hazard_job, 'test_curves_1', ds_key='hcurve'
        )
        models.Output.objects.create_output(
            hazard_job, 'test_curves_2', ds_key='hcurve'
        )

        # Sanity check: make sure the hazard calculation and outputs exist in
        # the database:
        hazard_jobs = models.OqJob.objects.filter(id=hazard_job.id)
        self.assertEqual(1, hazard_jobs.count())

        outputs = models.Output.objects.filter(oq_job=hazard_job.id)
        self.assertEqual(2, outputs.count())

        # Delete the calculation
        actions.del_calc(hazard_job.id)

        # Check that the hazard calculation and its outputs were deleted:
        outputs = models.Output.objects.filter(oq_job=hazard_job.id)
        self.assertEqual(0, outputs.count())

        hazard_jobs = models.OqJob.objects.filter(id=hazard_job.id)
        self.assertEqual(0, hazard_jobs.count())

    def test_del_calc_does_not_exist(self):
        self.assertRaises(RuntimeError, actions.del_calc, -1)

    def test_del_calc_no_access(self):
        # Test the case where we try to delete a hazard calculation which does
        # not belong to current user.
        # In this case, deletion is now allowed and should raise an exception.
        hazard_job = get_job(self.hazard_cfg, helpers.random_string())
        self.assertRaises(RuntimeError, actions.del_calc, hazard_job.id)

    def test_del_calc_referenced_by_risk_calc(self):
        # Test the case where a risk calculation is referencing the hazard
        # calculation we want to delete.
        # In this case, deletion is not allowed and should raise an exception.
        risk_job = get_job(self.risk_cfg, getpass.getuser(),
                           hazard_calculation_id=self.job.id)
        hc = risk_job.hazard_calculation
        self.assertRaises(RuntimeError, actions.del_calc, hc.id)

    def test_del_calc_output_referenced_by_risk_calc(self):
        # Test the case where a risk calculation is referencing one of the
        # belonging to the hazard calculation we want to delete.
        # In this case, deletion is not allowed and should raise an exception.
        risk_job = get_job(self.risk_cfg, getpass.getuser(),
                           hazard_calculation_id=self.job.id)
        hc = risk_job.hazard_calculation
        self.assertRaises(RuntimeError, actions.del_calc, hc.id)

class DeleteRiskCalcTestCase(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.hazard_cfg = helpers.get_data_path(
            'simple_fault_demo_hazard/job.ini')
        cls.risk_cfg = helpers.get_data_path(
            'classical_psha_based_risk/job.ini')
        cls.job = get_job(cls.hazard_cfg, getpass.getuser())

    def test_del_calc(self):
        risk_job = get_job(self.risk_cfg, getpass.getuser(),
                           hazard_calculation_id=self.job.id)
        models.Output.objects.create_output(
            risk_job, 'test_curves_1', ds_key='rcurves-rlzs'
        )
        models.Output.objects.create_output(
            risk_job, 'test_curves_2', ds_key='rcurves-rlzs'
        )

        # Sanity check: make sure the risk calculation and outputs exist in
        # the database:
        risk_calcs = models.OqJob.objects.filter(id=risk_job.id)
        self.assertEqual(1, risk_calcs.count())

        outputs = models.Output.objects.filter(oq_job=risk_job.id)
        self.assertEqual(2, outputs.count())

        # Delete the calculation
        actions.del_calc(risk_job.id)

        # Check that the risk calculation and its outputs were deleted:
        outputs = models.Output.objects.filter(oq_job=risk_job.id)
        self.assertEqual(0, outputs.count())

        risk_calcs = models.OqJob.objects.filter(id=risk_job.id)
        self.assertEqual(0, risk_calcs.count())

    def test_del_calc_does_not_exist(self):
        self.assertRaises(RuntimeError, actions.del_calc, -1)

    def test_del_calc_no_access(self):
        # Test the case where we try to delete a risk calculation which does
        # not belong to current user.
        # In this case, deletion is now allowed and should raise an exception.
        risk_job = get_job(self.risk_cfg, helpers.random_string(),
                           hazard_calculation_id=self.job.id)
        self.assertRaises(RuntimeError, actions.del_calc, risk_job.id)

class FakeOutput(object):
    def __init__(self, id, display_name):
        self.id = id
        self.display_name = display_name

    def get_output_type_display(self):
        return self.display_name + str(self.id)

class PrintSummaryTestCase(unittest.TestCase):
    outputs = [FakeOutput(i, 'gmf') for i in range(1, 12)]

    def print_outputs_summary(self, full):
        got = actions.print_outputs_summary(self.outputs, full)
        return '\n'.join(got)

    def test_print_outputs_summary_full(self):
        self.assertEqual(self.print_outputs_summary(full=True), '''\
  id | name
   1 | gmf
   2 | gmf
   3 | gmf
   4 | gmf
   5 | gmf
   6 | gmf
   7 | gmf
   8 | gmf
   9 | gmf
  10 | gmf
  11 | gmf''')

    def test_print_outputs_summary_short(self):
        self.assertEqual(
            self.print_outputs_summary(full=False), '''\
  id | name
   1 | gmf
   2 | gmf
   3 | gmf
   4 | gmf
   5 | gmf
   6 | gmf
   7 | gmf
   8 | gmf
   9 | gmf
  10 | gmf
 ... | 1 additional output(s)
Some outputs where not shown. You can see the full list with the command
`oq-engine --list-outputs`''')
<commit_msg>Fixed engine_test.py<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-2016 GEM Foundation
#
# OpenQuake 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.
#
# OpenQuake 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 OpenQuake. If not, see <http://www.gnu.org/licenses/>.

import os
import getpass
import subprocess
import unittest
import tempfile
import mock

from django.db import connection
from django.core import exceptions
from openquake.server.db import models, actions, upgrade_manager
from openquake.server.settings import DATABASE
from openquake.server.tests import helpers

def setup_module():
    global tmpfile
    fh, tmpfile = tempfile.mkstemp()
    os.close(fh)
    DATABASE['name'] = tmpfile
    connection.cursor()  # connect to the db
    print 'Using the database', tmpfile
    upgrade_manager.upgrade_db(connection.connection)

def teardown_module():
    connection.close()

def get_job(cfg, username, hazard_calculation_id=None):
    job_id, oq = actions.job_from_file(cfg, username, hazard_calculation_id)
    return models.OqJob.objects.get(pk=job_id)

class FakeJob(object):
    def __init__(self, job_type, calculation_mode):
        self.job_type = job_type
        self.calculation_mode = calculation_mode

    def get_oqparam(self):
        m = mock.Mock()
        m.calculation_mode = self.calculation_mode
        return m

class CheckHazardRiskConsistencyTestCase(unittest.TestCase):
    def test_ok(self):
        haz_job = FakeJob('hazard', 'scenario')
        actions.check_hazard_risk_consistency(
            haz_job, 'scenario_risk')

    def test_obsolete_mode(self):
        haz_job = FakeJob('hazard', 'scenario')
        with self.assertRaises(ValueError) as ctx:
            actions.check_hazard_risk_consistency(
                haz_job, 'scenario')
        msg = str(ctx.exception)
        self.assertEqual(msg, 'Please change calculation_mode=scenario into '
                         'scenario_risk in the .ini file')

    def test_inconsistent_mode(self):
        haz_job = FakeJob('hazard', 'scenario')
        with self.assertRaises(actions.InvalidCalculationID) as ctx:
            actions.check_hazard_risk_consistency(
                haz_job, 'classical_risk')
        msg = str(ctx.exception)
        self.assertEqual(
            msg, "In order to run a risk calculation of kind "
            "'classical_risk', you need to provide a "
            "calculation of kind ['classical', 'classical_risk'], "
            "but you provided a 'scenario' instead")

class JobFromFileTestCase(unittest.TestCase):

    def test_create_job_default_user(self):
        jid = actions.create_job('classical', 'test_create_job_default_user')
        job = models.OqJob.objects.get(pk=jid)
        self.assertEqual('openquake', job.user_name)
        self.assertEqual('executing', job.status)

    def test_create_job_specified_user(self):
        user_name = helpers.random_string()
        jid = actions.create_job(
            'classical', 'test_create_job_specified_user',
            user_name=user_name)
        job = models.OqJob.objects.get(pk=jid)
        self.assertEqual(user_name, job.user_name)
        self.assertEqual('executing', job.status)

class OpenquakeCliTestCase(unittest.TestCase):
    """
    Run "oq-engine --version" as a separate process using `subprocess`.
    """

    def test_run_version(self):
        args = [helpers.RUNNER, "--version"]

        print 'Running:', ' '.join(args)  # this is useful for debugging
        return subprocess.check_call(args)

class DeleteHazCalcTestCase(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.hazard_cfg = helpers.get_data_path(
            'simple_fault_demo_hazard/job.ini')
        cls.risk_cfg = helpers.get_data_path(
            'classical_psha_based_risk/job.ini')
        cls.job = get_job(cls.hazard_cfg, getpass.getuser())

    def test_del_calc(self):
        hazard_job = get_job(self.hazard_cfg, getpass.getuser())

        models.Output.objects.create_output(
            hazard_job, 'test_curves_1', ds_key='hcurve'
        )
        models.Output.objects.create_output(
            hazard_job, 'test_curves_2', ds_key='hcurve'
        )

        # Sanity check: make sure the hazard calculation and outputs exist in
        # the database:
        hazard_jobs = models.OqJob.objects.filter(id=hazard_job.id)
        self.assertEqual(1, hazard_jobs.count())

        outputs = models.Output.objects.filter(oq_job=hazard_job.id)
        self.assertEqual(2, outputs.count())

        # Delete the calculation
        actions.del_calc(hazard_job.id)

        # Check that the hazard calculation and its outputs were deleted:
        outputs = models.Output.objects.filter(oq_job=hazard_job.id)
        self.assertEqual(0, outputs.count())

        hazard_jobs = models.OqJob.objects.filter(id=hazard_job.id)
        self.assertEqual(0, hazard_jobs.count())

    def test_del_calc_does_not_exist(self):
        self.assertRaises(RuntimeError, actions.del_calc, -1)

    def test_del_calc_no_access(self):
        # Test the case where we try to delete a hazard calculation which does
        # not belong to current user.
        # In this case, deletion is now allowed and should raise an exception.
        hazard_job = get_job(self.hazard_cfg, helpers.random_string())
        self.assertRaises(RuntimeError, actions.del_calc, hazard_job.id)

    def test_del_calc_referenced_by_risk_calc(self):
        # Test the case where a risk calculation is referencing the hazard
        # calculation we want to delete.
        # In this case, deletion is not allowed and should raise an exception.
        risk_job = get_job(self.risk_cfg, getpass.getuser(),
                           hazard_calculation_id=self.job.id)
        hc = risk_job.hazard_calculation
        self.assertRaises(RuntimeError, actions.del_calc, hc.id)

    def test_del_calc_output_referenced_by_risk_calc(self):
        # Test the case where a risk calculation is referencing one of the
        # belonging to the hazard calculation we want to delete.
        # In this case, deletion is not allowed and should raise an exception.
        risk_job = get_job(self.risk_cfg, getpass.getuser(),
                           hazard_calculation_id=self.job.id)
        hc = risk_job.hazard_calculation
        self.assertRaises(RuntimeError, actions.del_calc, hc.id)

class DeleteRiskCalcTestCase(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.hazard_cfg = helpers.get_data_path(
            'simple_fault_demo_hazard/job.ini')
        cls.risk_cfg = helpers.get_data_path(
            'classical_psha_based_risk/job.ini')
        cls.job = get_job(cls.hazard_cfg, getpass.getuser())

    def test_del_calc(self):
        risk_job = get_job(self.risk_cfg, getpass.getuser(),
                           hazard_calculation_id=self.job.id)
        models.Output.objects.create_output(
            risk_job, 'test_curves_1', ds_key='rcurves-rlzs'
        )
        models.Output.objects.create_output(
            risk_job, 'test_curves_2', ds_key='rcurves-rlzs'
        )

        # Sanity check: make sure the risk calculation and outputs exist in
        # the database:
        risk_calcs = models.OqJob.objects.filter(id=risk_job.id)
        self.assertEqual(1, risk_calcs.count())

        outputs = models.Output.objects.filter(oq_job=risk_job.id)
        self.assertEqual(2, outputs.count())

        # Delete the calculation
        actions.del_calc(risk_job.id)

        # Check that the risk calculation and its outputs were deleted:
        outputs = models.Output.objects.filter(oq_job=risk_job.id)
        self.assertEqual(0, outputs.count())

        risk_calcs = models.OqJob.objects.filter(id=risk_job.id)
        self.assertEqual(0, risk_calcs.count())

    def test_del_calc_does_not_exist(self):
        self.assertRaises(RuntimeError, actions.del_calc, -1)

    def test_del_calc_no_access(self):
        # Test the case where we try to delete a risk calculation which does
        # not belong to current user.
        # In this case, deletion is now allowed and should raise an exception.
        risk_job = get_job(self.risk_cfg, helpers.random_string(),
                           hazard_calculation_id=self.job.id)
        self.assertRaises(RuntimeError, actions.del_calc, risk_job.id)

class FakeOutput(object):
    def __init__(self, id, display_name):
        self.id = id
        self.display_name = display_name

    def get_output_type_display(self):
        return self.display_name + str(self.id)

class PrintSummaryTestCase(unittest.TestCase):
    outputs = [FakeOutput(i, 'gmf') for i in range(1, 12)]

    def print_outputs_summary(self, full):
        got = actions.print_outputs_summary(self.outputs, full)
        return '\n'.join(got)

    def test_print_outputs_summary_full(self):
        self.assertEqual(self.print_outputs_summary(full=True), '''\
  id | name
   1 | gmf
   2 | gmf
   3 | gmf
   4 | gmf
   5 | gmf
   6 | gmf
   7 | gmf
   8 | gmf
   9 | gmf
  10 | gmf
  11 | gmf''')

    def test_print_outputs_summary_short(self):
        self.assertEqual(
            self.print_outputs_summary(full=False), '''\
  id | name
   1 | gmf
   2 | gmf
   3 | gmf
   4 | gmf
   5 | gmf
   6 | gmf
   7 | gmf
   8 | gmf
   9 | gmf
  10 | gmf
 ... | 1 additional output(s)
Some outputs where not shown. You can see the full list with the command
`oq-engine --list-outputs`''')
<|endoftext|>