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

import mock

from six.moves import http_client
import st2common.bootstrap.actionsregistrar as actions_registrar
from st2common.persistence.pack import Pack
from st2tests.api import FunctionalTest

@mock.patch("st2common.bootstrap.base.REGISTERED_PACKS_CACHE", {})
class PacksViewsControllerTestCase(FunctionalTest):
    @classmethod
    def setUpClass(cls):
        super(PacksViewsControllerTestCase, cls).setUpClass()

        # Register local action and pack fixtures
        actions_registrar.register_actions(use_pack_cache=False)

    def test_get_pack_files_success(self):
        resp = self.app.get("/v1/packs/views/files/dummy_pack_1")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertTrue(len(resp.json) > 1)
        item = [_item for _item in resp.json if _item["file_path"] == "pack.yaml"][0]
        self.assertEqual(item["file_path"], "pack.yaml")
        item = [
            _item for _item in resp.json if _item["file_path"] == "actions/my_action.py"
        ][0]
        self.assertEqual(item["file_path"], "actions/my_action.py")

    def test_get_pack_files_pack_doesnt_exist(self):
        resp = self.app.get("/v1/packs/views/files/doesntexist", expect_errors=True)
        self.assertEqual(resp.status_int, http_client.NOT_FOUND)

    def test_get_pack_files_binary_files_are_excluded(self):
        binary_files = [
            "icon.png",
            "etc/permissions.png",
            "etc/travisci.png",
            "etc/generate_new_token.png",
        ]

        pack_db = Pack.get_by_ref("dummy_pack_1")

        all_files_count = len(pack_db.files)
        non_binary_files_count = all_files_count - len(binary_files)

        resp = self.app.get("/v1/packs/views/files/dummy_pack_1")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertEqual(len(resp.json), non_binary_files_count)

        for file_path in binary_files:
            self.assertIn(file_path, pack_db.files)

        # But not in files controller response
        for file_path in binary_files:
            item = [item for item in resp.json if item["file_path"] == file_path]
            self.assertFalse(item)

    def test_get_pack_file_success(self):
        resp = self.app.get("/v1/packs/views/file/dummy_pack_1/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

    def test_get_pack_file_pack_doesnt_exist(self):
        resp = self.app.get(
            "/v1/packs/views/files/doesntexist/pack.yaml", expect_errors=True
        )
        self.assertEqual(resp.status_int, http_client.NOT_FOUND)

    @mock.patch("st2api.controllers.v1.pack_views.MAX_FILE_SIZE", 1)
    def test_pack_file_file_larger_then_maximum_size(self):
        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml", expect_errors=True
        )
        self.assertEqual(resp.status_int, http_client.BAD_REQUEST)
        self.assertIn("File pack.yaml exceeds maximum allowed file size", resp)

    def test_headers_get_pack_file(self):
        resp = self.app.get("/v1/packs/views/file/dummy_pack_1/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)
        self.assertIsNotNone(resp.headers["ETag"])
        self.assertIsNotNone(resp.headers["Last-Modified"])

    def test_no_change_get_pack_file(self):
        resp = self.app.get("/v1/packs/views/file/dummy_pack_1/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

        # Confirm NOT_MODIFIED
        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-None-Match": resp.headers["ETag"]},
        )
        self.assertEqual(resp.status_code, http_client.NOT_MODIFIED)

        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-Modified-Since": resp.headers["Last-Modified"]},
        )
        self.assertEqual(resp.status_code, http_client.NOT_MODIFIED)

        # Confirm value is returned if header do not match
        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-None-Match": "ETAG"},
        )
        self.assertEqual(resp.status_code, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-Modified-Since": "Last-Modified"},
        )
        self.assertEqual(resp.status_code, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

    def test_get_pack_files_and_pack_file_ref_doesnt_equal_pack_name(self):
        # Ref is not equal to the name, controller should still work
        resp = self.app.get("/v1/packs/views/files/dummy_pack_16")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertEqual(len(resp.json), 1)
        self.assertEqual(resp.json[0]["file_path"], "pack.yaml")

        resp = self.app.get("/v1/packs/views/file/dummy_pack_16/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"ref: dummy_pack_16", resp.body)
<commit_msg>adjust file count for test that uses dummy_pack_16<commit_after># Copyright 2020 The StackStorm Authors.
# Copyright 2019 Extreme Networks, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import mock

from six.moves import http_client
import st2common.bootstrap.actionsregistrar as actions_registrar
from st2common.persistence.pack import Pack
from st2tests.api import FunctionalTest

@mock.patch("st2common.bootstrap.base.REGISTERED_PACKS_CACHE", {})
class PacksViewsControllerTestCase(FunctionalTest):
    @classmethod
    def setUpClass(cls):
        super(PacksViewsControllerTestCase, cls).setUpClass()

        # Register local action and pack fixtures
        actions_registrar.register_actions(use_pack_cache=False)

    def test_get_pack_files_success(self):
        resp = self.app.get("/v1/packs/views/files/dummy_pack_1")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertTrue(len(resp.json) > 1)
        item = [_item for _item in resp.json if _item["file_path"] == "pack.yaml"][0]
        self.assertEqual(item["file_path"], "pack.yaml")
        item = [
            _item for _item in resp.json if _item["file_path"] == "actions/my_action.py"
        ][0]
        self.assertEqual(item["file_path"], "actions/my_action.py")

    def test_get_pack_files_pack_doesnt_exist(self):
        resp = self.app.get("/v1/packs/views/files/doesntexist", expect_errors=True)
        self.assertEqual(resp.status_int, http_client.NOT_FOUND)

    def test_get_pack_files_binary_files_are_excluded(self):
        binary_files = [
            "icon.png",
            "etc/permissions.png",
            "etc/travisci.png",
            "etc/generate_new_token.png",
        ]

        pack_db = Pack.get_by_ref("dummy_pack_1")

        all_files_count = len(pack_db.files)
        non_binary_files_count = all_files_count - len(binary_files)

        resp = self.app.get("/v1/packs/views/files/dummy_pack_1")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertEqual(len(resp.json), non_binary_files_count)

        for file_path in binary_files:
            self.assertIn(file_path, pack_db.files)

        # But not in files controller response
        for file_path in binary_files:
            item = [item for item in resp.json if item["file_path"] == file_path]
            self.assertFalse(item)

    def test_get_pack_file_success(self):
        resp = self.app.get("/v1/packs/views/file/dummy_pack_1/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

    def test_get_pack_file_pack_doesnt_exist(self):
        resp = self.app.get(
            "/v1/packs/views/files/doesntexist/pack.yaml", expect_errors=True
        )
        self.assertEqual(resp.status_int, http_client.NOT_FOUND)

    @mock.patch("st2api.controllers.v1.pack_views.MAX_FILE_SIZE", 1)
    def test_pack_file_file_larger_then_maximum_size(self):
        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml", expect_errors=True
        )
        self.assertEqual(resp.status_int, http_client.BAD_REQUEST)
        self.assertIn("File pack.yaml exceeds maximum allowed file size", resp)

    def test_headers_get_pack_file(self):
        resp = self.app.get("/v1/packs/views/file/dummy_pack_1/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)
        self.assertIsNotNone(resp.headers["ETag"])
        self.assertIsNotNone(resp.headers["Last-Modified"])

    def test_no_change_get_pack_file(self):
        resp = self.app.get("/v1/packs/views/file/dummy_pack_1/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

        # Confirm NOT_MODIFIED
        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-None-Match": resp.headers["ETag"]},
        )
        self.assertEqual(resp.status_code, http_client.NOT_MODIFIED)

        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-Modified-Since": resp.headers["Last-Modified"]},
        )
        self.assertEqual(resp.status_code, http_client.NOT_MODIFIED)

        # Confirm value is returned if header do not match
        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-None-Match": "ETAG"},
        )
        self.assertEqual(resp.status_code, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

        resp = self.app.get(
            "/v1/packs/views/file/dummy_pack_1/pack.yaml",
            headers={"If-Modified-Since": "Last-Modified"},
        )
        self.assertEqual(resp.status_code, http_client.OK)
        self.assertIn(b"name : dummy_pack_1", resp.body)

    def test_get_pack_files_and_pack_file_ref_doesnt_equal_pack_name(self):
        # Ref is not equal to the name, controller should still work
        resp = self.app.get("/v1/packs/views/files/dummy_pack_16")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertEqual(len(resp.json), 3)
        self.assertIn("pack.yaml", [f["file_path"] for f in resp.json])

        resp = self.app.get("/v1/packs/views/file/dummy_pack_16/pack.yaml")
        self.assertEqual(resp.status_int, http_client.OK)
        self.assertIn(b"ref: dummy_pack_16", resp.body)
<|endoftext|>


<commit_before># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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 os
import time
import subprocess

from st2common import log as logging
from st2common.util.sandboxing import get_sandbox_python_path
from st2common.util.sandboxing import get_sandbox_python_binary_path

__all__ = [
    'ProcessSensorContainer'
]

LOG = logging.getLogger('st2reactor.process_sensor_container')

PROCESS_EXIT_TIMEOUT = 5  # how long to wait for process to exit after sending SIGKILL (in seconds)
SUCCESS_EXIT_CODE = 0

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
WRAPPER_SCRIPT_NAME = 'sensor_wrapper.py'
WRAPPER_SCRIPT_PATH = os.path.join(BASE_DIR, WRAPPER_SCRIPT_NAME)

# TODO: Allow multiple instances of the same sensor with different configuration
# options - we need to update sensors for that and add "get_id" or similar
# method to the sensor class

class ProcessSensorContainer(object):
    """
    Sensor container which runs sensors in a separate process.
    """

    def __init__(self, sensors):
        """
        :param sensors: A list of sensor dicts.
        :type sensors: ``list`` of ``dict``
        """
        self._sensors = {}  # maps sensor_id -> sensor object
        self._processes = {}  # maps sensor_id -> sensor process

        for sensor_obj in sensors:
            sensor_id = self._get_sensor_id(sensor=sensor_obj)
            self._sensors[sensor_id] = sensor_obj

    def run(self):
        self._run_all_sensors()

        # Poll for all running processes
        sensor_ids = self._sensors.keys()
        poll_interval = 5

        while len(sensor_ids) >= 1:
            for sensor_id in sensor_ids:
                process = self._processes[sensor_id]
                status = process.poll()

                if status is not None:
                    # Dead process detected
                    LOG.info('Process for sensor %s has exited with code %s',
                             sensor_id, status)

                    del self._processes[sensor_id]
                    del self._sensors[sensor_id]

            sensor_ids = self._sensors.keys()
            time.sleep(poll_interval)

        LOG.info('Container has no active sensors running.')
        return SUCCESS_EXIT_CODE

    def running(self):
        return len(self._processes)

    def shutdown(self):
        LOG.info('Container shutting down. Invoking cleanup on sensors.')

        for sensor_id in self._sensors:
            self._stop_sensor_process(sensor_id=sensor_id)

        LOG.info('All sensors are shut down.')

        self._sensors = {}
        self._processes = {}

    def add_sensor(self, sensor):
        """
        Add a new sensor to the container.

        :type sensor: ``dict``
        """
        sensor_id = self._get_sensor_id(sensor=sensor)

        if sensor_id in self._sensors:
            LOG.warning('Sensor %s already exists and running.', sensor_id)
            return False

        self._spawn_sensor_process(sensor=sensor)
        self._sensors[sensor_id] = sensor
        return True

    def remove_sensor(self, sensor):
        """
        Remove an existing sensor from the container.

        :type sensor: ``dict``
        """
        sensor_id = self._get_sensor_id(sensor=sensor)

        if sensor_id not in self._sensors:
            LOG.warning('Sensor %s isn\'t running in this container.', sensor_id)
            return False

        self._stop_sensor_process(sensor_id=sensor_id)
        return True

    def _run_all_sensors(self):
        sensor_ids = self._sensors.keys()

        for sensor_id in sensor_ids:
            sensor_obj = self._sensors[sensor_id]
            LOG.info('Running sensor %s' % sensor_id)

            try:
                self._spawn_sensor_process(sensor=sensor_obj)
            except Exception as e:
                LOG.warn(e.message)

                # Disable sensor which we are unable to start
                del self._sensors[sensor_id]
                continue

            LOG.info('Sensor %s started' % sensor_id)

    def _spawn_sensor_process(self, sensor):
        """
        Spawn a new process for the provided sensor.

        New process uses isolated Python binary from a virtual environment
        belonging to the sensor pack.
        """
        sensor_id = self._get_sensor_id(sensor=sensor)
        python_path = get_sandbox_python_binary_path(pack=sensor['pack'])

        trigger_type_refs = sensor['trigger_types'] or []
        trigger_type_refs = ','.join(trigger_type_refs)

        args = [
            python_path,
            WRAPPER_SCRIPT_PATH,
            '--pack=%s' % (sensor['pack']),
            '--file-path=%s' % (sensor['file_path']),
            '--class-name=%s' % (sensor['class_name']),
            '--trigger-type-refs=%s' % (trigger_type_refs)
        ]

        if sensor['poll_interval']:
            args.append('--poll-interval=%s' % (sensor['poll_interval']))

        env = os.environ.copy()
        env['PYTHONPATH'] = get_sandbox_python_path(inherit_from_parent=True,
                                                    inherit_parent_virtualenv=True)

        LOG.debug('Running sensor subprocess (cmd="%s")', ' '.join(args))

        # TODO: Intercept stdout and stderr for aggregated logging purposes
        try:
            process = subprocess.Popen(args=args, stdin=None, stdout=None,
                                       stderr=None, shell=False, env=env)
        except Exception as e:
            cmd = ' '.join(args)
            message = ('Failed to spawn process for sensor %s ("%s"): %s' %
                       (sensor_id, cmd, str(e)))
            raise Exception(message)

        self._processes[sensor_id] = process
        return process

    def _stop_sensor_process(self, sensor_id, exit_timeout=5):
        """
        Stop a sensor process for the provided sensor.

        :param sensor_id: Sensor ID.
        :type sensor_id: ``str``

        :param exit_timeout: How long to wait for process to exit after
                             sending SIGTERM signal. If the process doesn't
                             exit in this amount of seconds, SIGTERM signal
                             will be sent to the process.
        :type exit__timeout: ``int``
        """
        process = self._processes[sensor_id]

        # Terminate the process and wait for up to stop_timeout seconds for the
        # process to exit
        process.terminate()

        timeout = 0
        sleep_delay = 1
        while timeout < exit_timeout:
            status = process.poll()

            if status is not None:
                # Process has exited
                break

            timeout += sleep_delay
            time.sleeps(sleep_delay)

        if status is None:
            # Process hasn't exited yet, forcefully kill it
            process.kill()

        del self._processes[sensor_id]
        del self._sensors[sensor_id]

    def _get_sensor_id(self, sensor):
        """
        Return unique identifier for the provider sensor dict.

        :type sensor: ``dict``
        """
        sensor_id = sensor['class_name']
        return sensor_id
<commit_msg>Fix a typo.<commit_after># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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 os
import time
import subprocess

from st2common import log as logging
from st2common.util.sandboxing import get_sandbox_python_path
from st2common.util.sandboxing import get_sandbox_python_binary_path

__all__ = [
    'ProcessSensorContainer'
]

LOG = logging.getLogger('st2reactor.process_sensor_container')

PROCESS_EXIT_TIMEOUT = 5  # how long to wait for process to exit after sending SIGKILL (in seconds)
SUCCESS_EXIT_CODE = 0

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
WRAPPER_SCRIPT_NAME = 'sensor_wrapper.py'
WRAPPER_SCRIPT_PATH = os.path.join(BASE_DIR, WRAPPER_SCRIPT_NAME)

# TODO: Allow multiple instances of the same sensor with different configuration
# options - we need to update sensors for that and add "get_id" or similar
# method to the sensor class

class ProcessSensorContainer(object):
    """
    Sensor container which runs sensors in a separate process.
    """

    def __init__(self, sensors):
        """
        :param sensors: A list of sensor dicts.
        :type sensors: ``list`` of ``dict``
        """
        self._sensors = {}  # maps sensor_id -> sensor object
        self._processes = {}  # maps sensor_id -> sensor process

        for sensor_obj in sensors:
            sensor_id = self._get_sensor_id(sensor=sensor_obj)
            self._sensors[sensor_id] = sensor_obj

    def run(self):
        self._run_all_sensors()

        # Poll for all running processes
        sensor_ids = self._sensors.keys()
        poll_interval = 5

        while len(sensor_ids) >= 1:
            for sensor_id in sensor_ids:
                process = self._processes[sensor_id]
                status = process.poll()

                if status is not None:
                    # Dead process detected
                    LOG.info('Process for sensor %s has exited with code %s',
                             sensor_id, status)

                    del self._processes[sensor_id]
                    del self._sensors[sensor_id]

            sensor_ids = self._sensors.keys()
            time.sleep(poll_interval)

        LOG.info('Container has no active sensors running.')
        return SUCCESS_EXIT_CODE

    def running(self):
        return len(self._processes)

    def shutdown(self):
        LOG.info('Container shutting down. Invoking cleanup on sensors.')

        for sensor_id in self._sensors:
            self._stop_sensor_process(sensor_id=sensor_id)

        LOG.info('All sensors are shut down.')

        self._sensors = {}
        self._processes = {}

    def add_sensor(self, sensor):
        """
        Add a new sensor to the container.

        :type sensor: ``dict``
        """
        sensor_id = self._get_sensor_id(sensor=sensor)

        if sensor_id in self._sensors:
            LOG.warning('Sensor %s already exists and running.', sensor_id)
            return False

        self._spawn_sensor_process(sensor=sensor)
        self._sensors[sensor_id] = sensor
        return True

    def remove_sensor(self, sensor):
        """
        Remove an existing sensor from the container.

        :type sensor: ``dict``
        """
        sensor_id = self._get_sensor_id(sensor=sensor)

        if sensor_id not in self._sensors:
            LOG.warning('Sensor %s isn\'t running in this container.', sensor_id)
            return False

        self._stop_sensor_process(sensor_id=sensor_id)
        return True

    def _run_all_sensors(self):
        sensor_ids = self._sensors.keys()

        for sensor_id in sensor_ids:
            sensor_obj = self._sensors[sensor_id]
            LOG.info('Running sensor %s' % sensor_id)

            try:
                self._spawn_sensor_process(sensor=sensor_obj)
            except Exception as e:
                LOG.warn(e.message)

                # Disable sensor which we are unable to start
                del self._sensors[sensor_id]
                continue

            LOG.info('Sensor %s started' % sensor_id)

    def _spawn_sensor_process(self, sensor):
        """
        Spawn a new process for the provided sensor.

        New process uses isolated Python binary from a virtual environment
        belonging to the sensor pack.
        """
        sensor_id = self._get_sensor_id(sensor=sensor)
        python_path = get_sandbox_python_binary_path(pack=sensor['pack'])

        trigger_type_refs = sensor['trigger_types'] or []
        trigger_type_refs = ','.join(trigger_type_refs)

        args = [
            python_path,
            WRAPPER_SCRIPT_PATH,
            '--pack=%s' % (sensor['pack']),
            '--file-path=%s' % (sensor['file_path']),
            '--class-name=%s' % (sensor['class_name']),
            '--trigger-type-refs=%s' % (trigger_type_refs)
        ]

        if sensor['poll_interval']:
            args.append('--poll-interval=%s' % (sensor['poll_interval']))

        env = os.environ.copy()
        env['PYTHONPATH'] = get_sandbox_python_path(inherit_from_parent=True,
                                                    inherit_parent_virtualenv=True)

        LOG.debug('Running sensor subprocess (cmd="%s")', ' '.join(args))

        # TODO: Intercept stdout and stderr for aggregated logging purposes
        try:
            process = subprocess.Popen(args=args, stdin=None, stdout=None,
                                       stderr=None, shell=False, env=env)
        except Exception as e:
            cmd = ' '.join(args)
            message = ('Failed to spawn process for sensor %s ("%s"): %s' %
                       (sensor_id, cmd, str(e)))
            raise Exception(message)

        self._processes[sensor_id] = process
        return process

    def _stop_sensor_process(self, sensor_id, exit_timeout=5):
        """
        Stop a sensor process for the provided sensor.

        :param sensor_id: Sensor ID.
        :type sensor_id: ``str``

        :param exit_timeout: How long to wait for process to exit after
                             sending SIGTERM signal. If the process doesn't
                             exit in this amount of seconds, SIGTERM signal
                             will be sent to the process.
        :type exit__timeout: ``int``
        """
        process = self._processes[sensor_id]

        # Terminate the process and wait for up to stop_timeout seconds for the
        # process to exit
        process.terminate()

        timeout = 0
        sleep_delay = 1
        while timeout < exit_timeout:
            status = process.poll()

            if status is not None:
                # Process has exited
                break

            timeout += sleep_delay
            time.sleep(sleep_delay)

        if status is None:
            # Process hasn't exited yet, forcefully kill it
            process.kill()

        del self._processes[sensor_id]
        del self._sensors[sensor_id]

    def _get_sensor_id(self, sensor):
        """
        Return unique identifier for the provider sensor dict.

        :type sensor: ``dict``
        """
        sensor_id = sensor['class_name']
        return sensor_id
<|endoftext|>


<commit_before># Copyright (c) 2011 Openstack, LLC.
# 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.

"""
The Zone Aware Scheduler is a base class Scheduler for creating instances
across zones. There are two expansion points to this class for:
1. Assigning Weights to hosts for requested instances
2. Filtering Hosts based on required instance capabilities
"""

import operator
import M2Crypto

from nova import db
from nova import rpc
from nova import log as logging
from nova.scheduler import api
from nova.scheduler import driver

LOG = logging.getLogger('nova.scheduler.zone_aware_scheduler')

class ZoneAwareScheduler(driver.Scheduler):
    """Base class for creating Zone Aware Schedulers."""

    def _call_zone_method(self, context, method, specs):
        """Call novaclient zone method. Broken out for testing."""
        return api.call_zone_method(context, method, specs=specs)

    def schedule_run_instance(self, context, instance_id, request_spec,
                                        *args, **kwargs):
        """This method is called from nova.compute.api to provision
        an instance. However we need to look at the parameters being
        passed in to see if this is a request to:
        1. Create a Build Plan and then provision, or
        2. Use the Build Plan information in the request parameters
           to simply create the instance (either in this zone or
           a child zone)."""

        # TODO(sandy): We'll have to look for richer specs at some point.

        if 'blob' in request_spec:
            self.provision_resource(context, request_spec, instance_id,
                                    request_spec, kwargs)
            return None

        # Create build plan and provision ...
        build_plan = self.select(context, request_spec)
        if not build_plan:
            raise driver.NoValidHost(_('No hosts were available'))

        for item in build_plan:
            self.provision_resource(context, item, instance_id, request_spec,
                                    kwargs)

        # Returning None short-circuits the routing to Compute (since
        # we've already done it here)
        return None

    def provision_resource(self, context, item, instance_id, request_spec,
                           kwargs):
        """Create the requested resource in this Zone or a child zone."""
        if "hostname" in item:
            self._provision_resource_locally(context, item, instance_id,
                            kwargs)
           return

        self._provision_resource_in_child_zone(context, item, instance_id,
                                               request_spec, kwargs)

    def _provision_resource_locally(self, context, item, instance_id, kwargs):
        """Create the requested resource in this Zone."""
        host = item['hostname']
        kwargs['instance_id'] = instance_id
        rpc.cast(context,
                 db.queue_get_for(context, "compute", host),
                 {"method": "run_instance",
                  "args": kwargs})
        LOG.debug(_("Casted to compute %(host)s for run_instance")
                            % locals())

    def _provision_resource_in_child_zone(self, context, item, instance_id,
                                          request_spec, kwargs):
        """Create the requested resource in a child zone."""
        # Start by attempting to decrypt the blob to see if this
        # request is:
        # 1. valid,
        # 2. intended for this zone or a child zone.
        # if 2 ... forward call to child zone.
        blob = item['blob']
        decryptor = crypto.decryptor(FLAGS.build_plan_encryption_key)
        host_info = None
        try:
            json_entry = decryptor(blob)
            host_info = json.dumps(entry)
        except M2Crypto.EVP.EVPError:
            pass

        if not host_info:
            raise exception.Invalid(_("Ill-formed or incorrectly "
                            "routed 'blob' data sent "
                            "to instance create request.") % locals())

        # Valid data ... is it for us?
        if 'child_zone' in host_info and 'child_blob' in host_info:
            self._ask_child_zone_to_create_instance(context, host_info,
                                                    request_spec, kwargs)
        else:
            self._provision_resource_locally(context, host_info,
                                             instance_id, kwargs)

    def _ask_child_zone_to_create_instance(self, zone_info, request_spec,
                                           kwargs):

        # Note: we have to reverse engineer from our args to get back the
        # image, flavor, ipgroup, etc. since the original call could have
        # come in from EC2 (which doesn't use these things).
        instance_type = request_spec['instance_type']
        instance_properties = request_spec['instance_properties']

        name = instance_properties['display_name']
        image_id = instance_properties['image_id'])
        flavor_id = instance_type['flavor_id']
        meta = instance_type['metadata']

        files = kwargs['injected_files']
        ipgroup = None  # Not supported in OS API ... yet

        child_zone = zone_info['child_zone']
        child_blob = zone_info['child_blob']
        zone = db.zone_get(child_zone)
        url = zone.api_url
        nova = None
        try:
            nova = novaclient.OpenStack(zone.username, zone.password, url)
            nova.authenticate()
        except novaclient.exceptions.BadRequest, e:
             raise exception.NotAuthorized(_("Bad credentials attempting "
                            "to talk to zone at %(url)s.") % locals())

        nova.servers.create(name, image, flavor, ipgroup, meta, files)

    def select(self, context, request_spec, *args, **kwargs):
        """Select returns a list of weights and zone/host information
        corresponding to the best hosts to service the request. Any
        child zone information has been encrypted so as not to reveal
        anything about the children."""
        return self._schedule(context, "compute", request_spec,
                              *args, **kwargs)

    # TODO(sandy): We're only focused on compute instances right now,
    # so we don't implement the default "schedule()" method required
    # of Schedulers.
    def schedule(self, context, topic, request_spec, *args, **kwargs):
        """The schedule() contract requires we return the one
        best-suited host for this request.
        """
        raise driver.NoValidHost(_('No hosts were available'))

    def _schedule(self, context, topic, request_spec, *args, **kwargs):
        """Returns a list of hosts that meet the required specs,
        ordered by their fitness.
        """

        if topic != "compute":
            raise NotImplemented(_("Zone Aware Scheduler only understands "
                                   "Compute nodes (for now)"))

        #TODO(sandy): how to infer this from OS API params?
        num_instances = 1

        # Filter local hosts based on requirements ...
        host_list = self.filter_hosts(num_instances, request_spec)

        # then weigh the selected hosts.
        # weighted = [{weight=weight, name=hostname}, ...]
        weighted = self.weigh_hosts(num_instances, request_spec, host_list)

        # Next, tack on the best weights from the child zones ...
        child_results = self._call_zone_method(context, "select",
                specs=request_spec)
        for child_zone, result in child_results:
            for weighting in result:
                # Remember the child_zone so we can get back to
                # it later if needed. This implicitly builds a zone
                # path structure.
                host_dict = {
                        "weight": weighting["weight"],
                        "child_zone": child_zone,
                        "child_blob": weighting["blob"]}
                weighted.append(host_dict)

        weighted.sort(key=operator.itemgetter('weight'))
        return weighted

    def filter_hosts(self, num, request_spec):
        """Derived classes must override this method and return
           a list of hosts in [(hostname, capability_dict)] format."""
        raise NotImplemented()

    def weigh_hosts(self, num, request_spec, hosts):
        """Derived classes must override this method and return
           a lists of hosts in [{weight, hostname}] format."""
        raise NotImplemented()
<commit_msg>temp fixes<commit_after># Copyright (c) 2011 Openstack, LLC.
# 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.

"""
The Zone Aware Scheduler is a base class Scheduler for creating instances
across zones. There are two expansion points to this class for:
1. Assigning Weights to hosts for requested instances
2. Filtering Hosts based on required instance capabilities
"""

import operator
import M2Crypto

from nova import db
from nova import rpc
from nova import log as logging
from nova.scheduler import api
from nova.scheduler import driver

LOG = logging.getLogger('nova.scheduler.zone_aware_scheduler')

class ZoneAwareScheduler(driver.Scheduler):
    """Base class for creating Zone Aware Schedulers."""

    def _call_zone_method(self, context, method, specs):
        """Call novaclient zone method. Broken out for testing."""
        return api.call_zone_method(context, method, specs=specs)

    def schedule_run_instance(self, context, instance_id, request_spec,
                                        *args, **kwargs):
        """This method is called from nova.compute.api to provision
        an instance. However we need to look at the parameters being
        passed in to see if this is a request to:
        1. Create a Build Plan and then provision, or
        2. Use the Build Plan information in the request parameters
           to simply create the instance (either in this zone or
           a child zone)."""

        # TODO(sandy): We'll have to look for richer specs at some point.

        if 'blob' in request_spec:
            self.provision_resource(context, request_spec, instance_id,
                                    request_spec, kwargs)
            return None

        # Create build plan and provision ...
        build_plan = self.select(context, request_spec)
        if not build_plan:
            raise driver.NoValidHost(_('No hosts were available'))

        for item in build_plan:
            self.provision_resource(context, item, instance_id, request_spec,
                                    kwargs)

        # Returning None short-circuits the routing to Compute (since
        # we've already done it here)
        return None

    def provision_resource(self, context, item, instance_id, request_spec,
                           kwargs):
        """Create the requested resource in this Zone or a child zone."""
        if "hostname" in item:
            self._provision_resource_locally(context, item, instance_id,
                            kwargs)
           return

        self._provision_resource_in_child_zone(context, item, instance_id,
                                               request_spec, kwargs)

    def _provision_resource_locally(self, context, item, instance_id, kwargs):
        """Create the requested resource in this Zone."""
        host = item['hostname']
        kwargs['instance_id'] = instance_id
        rpc.cast(context,
                 db.queue_get_for(context, "compute", host),
                 {"method": "run_instance",
                  "args": kwargs})
        LOG.debug(_("Casted to compute %(host)s for run_instance")
                            % locals())

    def _provision_resource_in_child_zone(self, context, item, instance_id,
                                          request_spec, kwargs):
        """Create the requested resource in a child zone."""
        # Start by attempting to decrypt the blob to see if this
        # request is:
        # 1. valid,
        # 2. intended for this zone or a child zone.
        # if 2 ... forward call to child zone.
        blob = item['blob']
        decryptor = crypto.decryptor(FLAGS.build_plan_encryption_key)
        host_info = None
        try:
            json_entry = decryptor(blob)
            host_info = json.dumps(entry)
        except M2Crypto.EVP.EVPError:
            pass

        if not host_info:
            raise exception.Invalid(_("Ill-formed or incorrectly "
                            "routed 'blob' data sent "
                            "to instance create request.") % locals())

        # Valid data ... is it for us?
        if 'child_zone' in host_info and 'child_blob' in host_info:
            self._ask_child_zone_to_create_instance(context, host_info,
                                                    request_spec, kwargs)
        else:
            self._provision_resource_locally(context, host_info,
                                             instance_id, kwargs)

    def _ask_child_zone_to_create_instance(self, zone_info, request_spec,
                                           kwargs):

        # Note: we have to reverse engineer from our args to get back the
        # image, flavor, ipgroup, etc. since the original call could have
        # come in from EC2 (which doesn't use these things).
        instance_type = request_spec['instance_type']
        instance_properties = request_spec['instance_properties']

        name = instance_properties['display_name']
        image_id = instance_properties['image_id'])
        flavor_id = instance_type['flavor_id']
        meta = instance_type['metadata']

        files = kwargs['injected_files']
        ipgroup = None  # Not supported in OS API ... yet

        child_zone = zone_info['child_zone']
        child_blob = zone_info['child_blob']
        zone = db.zone_get(child_zone)
        url = zone.api_url
        nova = None
        try:
            nova = novaclient.OpenStack(zone.username, zone.password, url)
            nova.authenticate()
        except novaclient.exceptions.BadRequest, e:
             raise exception.NotAuthorized(_("Bad credentials attempting "
                            "to talk to zone at %(url)s.") % locals())

        nova.servers.create(name, image, flavor, ipgroup, meta, files,
                            child_blob)

    def select(self, context, request_spec, *args, **kwargs):
        """Select returns a list of weights and zone/host information
        corresponding to the best hosts to service the request. Any
        child zone information has been encrypted so as not to reveal
        anything about the children."""
        return self._schedule(context, "compute", request_spec,
                              *args, **kwargs)

    # TODO(sandy): We're only focused on compute instances right now,
    # so we don't implement the default "schedule()" method required
    # of Schedulers.
    def schedule(self, context, topic, request_spec, *args, **kwargs):
        """The schedule() contract requires we return the one
        best-suited host for this request.
        """
        raise driver.NoValidHost(_('No hosts were available'))

    def _schedule(self, context, topic, request_spec, *args, **kwargs):
        """Returns a list of hosts that meet the required specs,
        ordered by their fitness.
        """

        if topic != "compute":
            raise NotImplemented(_("Zone Aware Scheduler only understands "
                                   "Compute nodes (for now)"))

        #TODO(sandy): how to infer this from OS API params?
        num_instances = 1

        # Filter local hosts based on requirements ...
        host_list = self.filter_hosts(num_instances, request_spec)

        # then weigh the selected hosts.
        # weighted = [{weight=weight, name=hostname}, ...]
        weighted = self.weigh_hosts(num_instances, request_spec, host_list)

        # Next, tack on the best weights from the child zones ...
        child_results = self._call_zone_method(context, "select",
                specs=request_spec)
        for child_zone, result in child_results:
            for weighting in result:
                # Remember the child_zone so we can get back to
                # it later if needed. This implicitly builds a zone
                # path structure.
                host_dict = {
                        "weight": weighting["weight"],
                        "child_zone": child_zone,
                        "child_blob": weighting["blob"]}
                weighted.append(host_dict)

        weighted.sort(key=operator.itemgetter('weight'))
        return weighted

    def filter_hosts(self, num, request_spec):
        """Derived classes must override this method and return
           a list of hosts in [(hostname, capability_dict)] format."""
        raise NotImplemented()

    def weigh_hosts(self, num, request_spec, hosts):
        """Derived classes must override this method and return
           a lists of hosts in [{weight, hostname}] format."""
        raise NotImplemented()
<|endoftext|>


<commit_before>from __future__ import print_function, absolute_import

import tempfile
import os
import ctypes
from time import time

import numba.unittest_support as unittest
from numba import hsa
from numba import types
from numba.hsa import compiler
from numba.hsa.hsadrv.driver import hsa as hsart, BrigModule
from numba.targets.arrayobj import make_array_ctype

def copy_kernel(out, inp):
    out[0] = inp[0]

def copy_kernel_1d(out, inp):
    i = hsa.get_global_id(0)
    if i < out.size:
        out[i] = inp[i]

def assign_value(out, inp):
    i = hsa.get_global_id(0)
    if i < out.size:
        out[i] = inp

class TestCodeGeneration(unittest.TestCase):
    def test_copy_kernel(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel, [arytype] * 2)
        self.assertIn("prog kernel &{0}".format(kernel.entry_name),
                      kernel.assembly)

    def test_copy_kernel_1d(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)
        self.assertIn("prog kernel &{0}".format(kernel.entry_name),
                      kernel.assembly)

class _TestBase(unittest.TestCase):
    def setUp(self):
        self.gpu = [a for a in hsart.agents if a.is_component][0]
        self.cpu = [a for a in hsart.agents if not a.is_component][0]
        self.queue = self.gpu.create_queue_multi(self.gpu.queue_max_size)

    def tearDown(self):
        del self.queue
        del self.gpu
        del self.cpu

class TestCodeLoading(_TestBase):
    def test_loading_from_file(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)

        # Write the brig file out
        brig_file = tempfile.NamedTemporaryFile(delete=False)
        with brig_file:
            brig_file.write(kernel.binary)

        # Load BRIG file
        symbol = '&{0}'.format(kernel.entry_name)
        brig_module = BrigModule.from_file(brig_file.name)
        symbol_offset = brig_module.find_symbol_offset(symbol)
        self.assertTrue(symbol_offset)
        program = hsart.create_program([self.gpu])
        module = program.add_module(brig_module)
        code_descriptor = program.finalize(self.gpu, module, symbol_offset)
        self.assertGreater(code_descriptor._id.kernarg_segment_byte_size, 0)

        # Cleanup
        os.unlink(brig_file.name)

    def test_loading_from_memory(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)

        # Load BRIG memory
        symbol = '&{0}'.format(kernel.entry_name)
        brig_module = BrigModule.from_memory(kernel.binary)
        symbol_offset = brig_module.find_symbol_offset(symbol)
        self.assertTrue(symbol_offset)
        program = hsart.create_program([self.gpu])
        module = program.add_module(brig_module)
        code_descriptor = program.finalize(self.gpu, module, symbol_offset)
        self.assertGreater(code_descriptor._id.kernarg_segment_byte_size, 0)

import numpy as np
from numba.hsa.hsadrv import enums

class TestExecution(unittest.TestCase):
    def test_execute(self):
        src = np.arange(1024 * 1024 * 20, dtype=np.float32)
        dst = np.zeros_like(src)

        components = [a for a in hsart.agents if a.is_component]

        gpu = components[0]
        print("Using agent: {0} with queue size: {1}".format(gpu.name,
                                                             gpu.queue_max_size))
        q = gpu.create_queue_multi(gpu.queue_max_size)

        # Compiler kernel
        arytype = types.float32[::1]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)
        # print(kernel.assembly)

        # Load BRIG memory
        symbol = '&{0}'.format(kernel.entry_name)
        brig_module = BrigModule.from_memory(kernel.binary)
        symbol_offset = brig_module.find_symbol_offset(symbol)
        program = hsart.create_program([gpu])
        module = program.add_module(brig_module)
        code_descriptor = program.finalize(gpu, module, symbol_offset)

        print(program)

        kernarg_regions = [r for r in gpu.regions if r.supports_kernargs]
        assert kernarg_regions
        kernarg_region = kernarg_regions[0]
        # Specify the argument types required by the kernel and allocate them
        # note: in an ideal world this should come from kernel metadata
        kernarg_types = ctypes.c_void_p * (6 + 2)
        kernargs = kernarg_region.allocate(kernarg_types)

        cstruct = make_array_ctype(ndim=1)

        src_cstruct = cstruct(parent=None,
                              data=src.ctypes.data,
                              shape=src.ctypes.shape,
                              strides=src.ctypes.strides)

        dst_cstruct = cstruct(parent=None,
                              data=dst.ctypes.data,
                              shape=dst.ctypes.shape,
                              strides=dst.ctypes.strides)

        kernargs[0] = 0
        kernargs[1] = 0
        kernargs[2] = 0
        kernargs[3] = 0
        kernargs[4] = 0
        kernargs[5] = 0
        kernargs[6] = ctypes.addressof(dst_cstruct)
        kernargs[7] = ctypes.addressof(src_cstruct)

        hsart.hsa_memory_register(kernargs[6], ctypes.sizeof(cstruct))
        hsart.hsa_memory_register(kernargs[7], ctypes.sizeof(cstruct))
        hsart.hsa_memory_register(src.ctypes.data, src.nbytes)
        hsart.hsa_memory_register(dst.ctypes.data, dst.nbytes)
        hsart.hsa_memory_register(ctypes.byref(kernargs),
                                  ctypes.sizeof(kernargs))

        # sync (in fact, dispatch will create a dummy signal for the dispatch and
        # wait for it before returning)
        print("dispatch synchronous... ", end="")
        t_start = time()
        q.dispatch(code_descriptor, kernargs, workgroup_size=(256,),
                   grid_size=(src.size,))
        t_end = time()
        print("ellapsed: {0:10.9f} s.".format(t_end - t_start))

        # async: handle the signal by hand
        print("dispatch asynchronous... ", end="")
        t_start = time()
        s = hsart.create_signal(1)
        q.dispatch(code_descriptor, kernargs,
                   workgroup_size=(256,), grid_size=(src.size,), signal=s)
        t_launched = time()
        hsart.hsa_signal_wait_acquire(s._id, enums.HSA_LT, 1, -1,
                                      enums.HSA_WAIT_EXPECTANCY_UNKNOWN)
        t_end = time()
        print("launch: {0:10.9f} s. total: {1:10.9g} s.".format(
            t_launched - t_start, t_end - t_start))

        # this is placed in the kernarg_region for symmetry, but shouldn't be required.
        kernarg_region.free(kernargs)

        np.testing.assert_equal(src, dst)

    def test_hsa_kernel(self):
        src = np.arange(1024, dtype=np.float32)
        dst = np.zeros_like(src)

        # Compiler kernel
        arytype = types.float32[::1]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)

        # Run kernel
        kernel[src.size // 256, 256](dst, src)

        np.testing.assert_equal(src, dst)

class TestKernelArgument(unittest.TestCase):
    def _test_template(self, nbtype, src):
        dtype = np.dtype(str(nbtype))
        dst = np.zeros(1, dtype=dtype)
        src = dtype.type(src)
        arytype = nbtype[::1]
        kernel = compiler.compile_kernel(assign_value, [arytype, nbtype])
        kernel[1, 1](dst, src)
        self.assertEqual(dst[0], src)

    def test_float64(self):
        self._test_template(nbtype=types.float64, src=1. / 3.)

    def test_float32(self):
        self._test_template(nbtype=types.float32, src=1. / 3.)

    def test_int32(self):
        self._test_template(nbtype=types.int32, src=123)

    def test_int16(self):
        self._test_template(nbtype=types.int16, src=123)

    def test_complex64(self):
        self._test_template(nbtype=types.complex64, src=12 + 34j)

    def test_complex128(self):
        self._test_template(nbtype=types.complex128, src=12 + 34j)

def udt_devfunc(a, i):
    return a[i]

class TestDeviceFunction(unittest.TestCase):
    def test_device_function(self):
        src = np.arange(10, dtype=np.int32)
        dst = np.zeros_like(src)

        arytype = types.int32[::1]
        devfn = compiler.compile_device(udt_devfunc, arytype.dtype,
                                        [arytype, types.intp])

        def udt_devfunc_caller(dst, src):
            i = hsa.get_global_id(0)
            if i < dst.size:
                dst[i] = devfn(src, i)

        kernel = compiler.compile_kernel(udt_devfunc_caller,
                                         [arytype, arytype])

        kernel[src.size, 1](dst, src)
        np.testing.assert_equal(dst, src)

if __name__ == '__main__':
    unittest.main()

<commit_msg>Retire an execution test - it duplicates code - it is sensitive to backend modifications<commit_after>from __future__ import print_function, absolute_import

import tempfile
import os
import numpy as np

import numba.unittest_support as unittest
from numba import hsa
from numba import types
from numba.hsa import compiler
from numba.hsa.hsadrv.driver import hsa as hsart, BrigModule

def copy_kernel(out, inp):
    out[0] = inp[0]

def copy_kernel_1d(out, inp):
    i = hsa.get_global_id(0)
    if i < out.size:
        out[i] = inp[i]

def assign_value(out, inp):
    i = hsa.get_global_id(0)
    if i < out.size:
        out[i] = inp

class TestCodeGeneration(unittest.TestCase):
    def test_copy_kernel(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel, [arytype] * 2)
        self.assertIn("prog kernel &{0}".format(kernel.entry_name),
                      kernel.assembly)

    def test_copy_kernel_1d(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)
        self.assertIn("prog kernel &{0}".format(kernel.entry_name),
                      kernel.assembly)

class _TestBase(unittest.TestCase):
    def setUp(self):
        self.gpu = [a for a in hsart.agents if a.is_component][0]
        self.cpu = [a for a in hsart.agents if not a.is_component][0]
        self.queue = self.gpu.create_queue_multi(self.gpu.queue_max_size)

    def tearDown(self):
        del self.queue
        del self.gpu
        del self.cpu

class TestCodeLoading(_TestBase):
    def test_loading_from_file(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)

        # Write the brig file out
        brig_file = tempfile.NamedTemporaryFile(delete=False)
        with brig_file:
            brig_file.write(kernel.binary)

        # Load BRIG file
        symbol = '&{0}'.format(kernel.entry_name)
        brig_module = BrigModule.from_file(brig_file.name)
        symbol_offset = brig_module.find_symbol_offset(symbol)
        self.assertTrue(symbol_offset)
        program = hsart.create_program([self.gpu])
        module = program.add_module(brig_module)
        code_descriptor = program.finalize(self.gpu, module, symbol_offset)
        self.assertGreater(code_descriptor._id.kernarg_segment_byte_size, 0)

        # Cleanup
        os.unlink(brig_file.name)

    def test_loading_from_memory(self):
        arytype = types.float32[:]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)

        # Load BRIG memory
        symbol = '&{0}'.format(kernel.entry_name)
        brig_module = BrigModule.from_memory(kernel.binary)
        symbol_offset = brig_module.find_symbol_offset(symbol)
        self.assertTrue(symbol_offset)
        program = hsart.create_program([self.gpu])
        module = program.add_module(brig_module)
        code_descriptor = program.finalize(self.gpu, module, symbol_offset)
        self.assertGreater(code_descriptor._id.kernarg_segment_byte_size, 0)

class TestExecution(unittest.TestCase):
    def test_hsa_kernel(self):
        src = np.arange(1024, dtype=np.float32)
        dst = np.zeros_like(src)

        # Compiler kernel
        arytype = types.float32[::1]
        kernel = compiler.compile_kernel(copy_kernel_1d, [arytype] * 2)

        # Run kernel
        kernel[src.size // 256, 256](dst, src)

        np.testing.assert_equal(src, dst)

class TestKernelArgument(unittest.TestCase):
    def _test_template(self, nbtype, src):
        dtype = np.dtype(str(nbtype))
        dst = np.zeros(1, dtype=dtype)
        src = dtype.type(src)
        arytype = nbtype[::1]
        kernel = compiler.compile_kernel(assign_value, [arytype, nbtype])
        kernel[1, 1](dst, src)
        self.assertEqual(dst[0], src)

    def test_float64(self):
        self._test_template(nbtype=types.float64, src=1. / 3.)

    def test_float32(self):
        self._test_template(nbtype=types.float32, src=1. / 3.)

    def test_int32(self):
        self._test_template(nbtype=types.int32, src=123)

    def test_int16(self):
        self._test_template(nbtype=types.int16, src=123)

    def test_complex64(self):
        self._test_template(nbtype=types.complex64, src=12 + 34j)

    def test_complex128(self):
        self._test_template(nbtype=types.complex128, src=12 + 34j)

def udt_devfunc(a, i):
    return a[i]

class TestDeviceFunction(unittest.TestCase):
    def test_device_function(self):
        src = np.arange(10, dtype=np.int32)
        dst = np.zeros_like(src)

        arytype = types.int32[::1]
        devfn = compiler.compile_device(udt_devfunc, arytype.dtype,
                                        [arytype, types.intp])

        def udt_devfunc_caller(dst, src):
            i = hsa.get_global_id(0)
            if i < dst.size:
                dst[i] = devfn(src, i)

        kernel = compiler.compile_kernel(udt_devfunc_caller,
                                         [arytype, arytype])

        kernel[src.size, 1](dst, src)
        np.testing.assert_equal(dst, src)

if __name__ == '__main__':
    unittest.main()

<|endoftext|>


<commit_before>from mpi4py import MPI
import mpiunittest as unittest

try:
    _reduce = reduce
except NameError:
    from functools import reduce as _reduce
cumsum  = lambda seq: _reduce(lambda x, y: x+y, seq, 0)
cumprod = lambda seq: _reduce(lambda x, y: x*y, seq, 1)

_basic = [None,
          True, False,
          -7, 0, 7, 2**31,
          -2**63, 2**63-1,
          -2.17, 0.0, 3.14,
          1+2j, 2-3j,
          'mpi4py',
          ]
messages = _basic
messages += [ list(_basic),
              tuple(_basic),
              dict([('k%d' % key, val)
                    for key, val in enumerate(_basic)])
              ]

class BaseTestCCOObjInter(object):

    BASECOMM  = MPI.COMM_NULL
    INTRACOMM = MPI.COMM_NULL
    INTERCOMM = MPI.COMM_NULL

    def setUp(self):
        BASE_SIZE = self.BASECOMM.Get_size()
        BASE_RANK = self.BASECOMM.Get_rank()
        if BASE_SIZE < 2:
            return
        if BASE_RANK < BASE_SIZE // 2 :
            self.COLOR = 0
            self.LOCAL_LEADER = 0
            self.REMOTE_LEADER = BASE_SIZE // 2
        else:
            self.COLOR = 1
            self.LOCAL_LEADER = 0
            self.REMOTE_LEADER = 0
        self.INTRACOMM = self.BASECOMM.Split(self.COLOR, key=0)
        self.INTERCOMM = self.INTRACOMM.Create_intercomm(self.LOCAL_LEADER,
                                                         self.BASECOMM,
                                                         self.REMOTE_LEADER)

    def tearDown(self):
        if self.INTRACOMM != MPI.COMM_NULL:
            self.INTRACOMM.Free()
            del self.INTRACOMM
        if self.INTERCOMM != MPI.COMM_NULL:
            self.INTERCOMM.Free()
            del self.INTERCOMM

    def testBarrier(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        self.INTERCOMM.Barrier()

    def testBcast(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            rmess = self.INTERCOMM.bcast(smess, root=MPI.ROOT)
                        else:
                            rmess = self.INTERCOMM.bcast(None, root=MPI.PROC_NULL)
                        self.assertEqual(rmess, None)
                else:
                    for root in range(rsize):
                        rmess = self.INTERCOMM.bcast(None, root=root)
                        self.assertEqual(rmess, smess)

    def testGather(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            rmess = self.INTERCOMM.gather(smess, root=MPI.ROOT)
                            self.assertEqual(rmess, [smess] * rsize)
                        else:
                            rmess = self.INTERCOMM.gather(None, root=MPI.PROC_NULL)
                            self.assertEqual(rmess, None)
                else:
                    for root in range(rsize):
                        rmess = self.INTERCOMM.gather(smess, root=root)
                        self.assertEqual(rmess, None)

    def testScatter(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            rmess = self.INTERCOMM.scatter([smess] * rsize, root=MPI.ROOT)
                        else:
                            rmess = self.INTERCOMM.scatter(None, root=MPI.PROC_NULL)
                        self.assertEqual(rmess, None)
                else:
                    for root in range(rsize):
                        rmess = self.INTERCOMM.scatter(None, root=root)
                        self.assertEqual(rmess, smess)

    def testAllgather(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            rmess = self.INTERCOMM.allgather(smess)
            self.assertEqual(rmess, [smess] * rsize)

    def testAlltoall(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            rmess = self.INTERCOMM.alltoall([smess] * rsize)
            self.assertEqual(rmess, [smess] * rsize)

    def testReduce(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for op in (MPI.SUM, MPI.MAX, MPI.MIN, MPI.PROD):
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            value = self.INTERCOMM.reduce(None, op=op, root=MPI.ROOT)
                            if op == MPI.SUM:
                                self.assertEqual(value, cumsum(range(rsize)))
                            elif op == MPI.PROD:
                                self.assertEqual(value, cumprod(range(rsize)))
                            elif op == MPI.MAX:
                                self.assertEqual(value, rsize-1)
                            elif op == MPI.MIN:
                                self.assertEqual(value, 0)
                        else:
                            value = self.INTERCOMM.reduce(None, op=op, root=MPI.PROC_NULL)
                            self.assertEqual(value, None)
                else:
                    for root in range(rsize):
                        value = self.INTERCOMM.reduce(rank, op=op, root=root)
                        self.assertEqual(value, None)

    def testAllreduce(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for op in (MPI.SUM, MPI.MAX, MPI.MIN, MPI.PROD):
            value = self.INTERCOMM.allreduce(rank, None, op)
            if op == MPI.SUM:
                self.assertEqual(value, cumsum(range(rsize)))
            elif op == MPI.PROD:
                self.assertEqual(value, cumprod(range(rsize)))
            elif op == MPI.MAX:
                self.assertEqual(value, rsize-1)
            elif op == MPI.MIN:
                self.assertEqual(value, 0)

class TestCCOObjInter(BaseTestCCOObjInter, unittest.TestCase):
    BASECOMM = MPI.COMM_WORLD

class TestCCOObjInterDup(TestCCOObjInter):
    def setUp(self):
        self.BASECOMM = self.BASECOMM.Dup()
        super(TestCCOObjInterDup, self).setUp()
    def tearDown(self):
        self.BASECOMM.Free()
        del self.BASECOMM
        super(TestCCOObjInterDup, self).tearDown()

class TestCCOObjInterDupDup(TestCCOObjInterDup):
    BASECOMM = MPI.COMM_WORLD
    INTERCOMM_ORIG = MPI.COMM_NULL
    def setUp(self):
        super(TestCCOObjInterDupDup, self).setUp()
        if self.INTERCOMM == MPI.COMM_NULL: return
        self.INTERCOMM_ORIG = self.INTERCOMM
        self.INTERCOMM = self.INTERCOMM.Dup()
    def tearDown(self):
        super(TestCCOObjInterDupDup, self).tearDown()
        if self.INTERCOMM_ORIG == MPI.COMM_NULL: return
        self.INTERCOMM_ORIG.Free()

_name, _version = MPI.get_vendor()
if _name == 'Open MPI':
    if _version < (1, 3, 0):
        del TestCCOObjInter
        del TestCCOObjInterDup
        del TestCCOObjInterDupDup
    elif _version < (1, 4, 0):
        if MPI.Query_thread() > MPI.THREAD_SINGLE:
            del TestCCOObjInter
            del TestCCOObjInterDup
            del TestCCOObjInterDupDup
elif _name == "MPICH2" or _name == "Microsoft MPI":
    if _version <= (1, 0, 7):
        def _SKIPPED(*args, **kwargs): pass
        TestCCOObjInterDupDup.testBarrier   = _SKIPPED
        TestCCOObjInterDupDup.testAllgather = _SKIPPED
        TestCCOObjInterDupDup.testAllreduce = _SKIPPED
elif _name == "DeinoMPI":
    def _SKIPPED(*args, **kwargs): pass
    TestCCOObjInterDupDup.testBarrier   = _SKIPPED
    TestCCOObjInterDupDup.testAllgather = _SKIPPED
    TestCCOObjInterDupDup.testAllreduce = _SKIPPED
elif _name == "MPICH1":
    del BaseTestCCOObjInter
    del TestCCOObjInter
    del TestCCOObjInterDup
    del TestCCOObjInterDupDup
elif MPI.ROOT == MPI.PROC_NULL:
    del BaseTestCCOObjInter
    del TestCCOObjInter
    del TestCCOObjInterDup
    del TestCCOObjInterDupDup

if __name__ == '__main__':
    unittest.main()
<commit_msg>Open MPI: Disable failing tests with 1.4.x<commit_after>from mpi4py import MPI
import mpiunittest as unittest

try:
    _reduce = reduce
except NameError:
    from functools import reduce as _reduce
cumsum  = lambda seq: _reduce(lambda x, y: x+y, seq, 0)
cumprod = lambda seq: _reduce(lambda x, y: x*y, seq, 1)

_basic = [None,
          True, False,
          -7, 0, 7, 2**31,
          -2**63, 2**63-1,
          -2.17, 0.0, 3.14,
          1+2j, 2-3j,
          'mpi4py',
          ]
messages = _basic
messages += [ list(_basic),
              tuple(_basic),
              dict([('k%d' % key, val)
                    for key, val in enumerate(_basic)])
              ]

class BaseTestCCOObjInter(object):

    BASECOMM  = MPI.COMM_NULL
    INTRACOMM = MPI.COMM_NULL
    INTERCOMM = MPI.COMM_NULL

    def setUp(self):
        BASE_SIZE = self.BASECOMM.Get_size()
        BASE_RANK = self.BASECOMM.Get_rank()
        if BASE_SIZE < 2:
            return
        if BASE_RANK < BASE_SIZE // 2 :
            self.COLOR = 0
            self.LOCAL_LEADER = 0
            self.REMOTE_LEADER = BASE_SIZE // 2
        else:
            self.COLOR = 1
            self.LOCAL_LEADER = 0
            self.REMOTE_LEADER = 0
        self.INTRACOMM = self.BASECOMM.Split(self.COLOR, key=0)
        self.INTERCOMM = self.INTRACOMM.Create_intercomm(self.LOCAL_LEADER,
                                                         self.BASECOMM,
                                                         self.REMOTE_LEADER)

    def tearDown(self):
        if self.INTRACOMM != MPI.COMM_NULL:
            self.INTRACOMM.Free()
            del self.INTRACOMM
        if self.INTERCOMM != MPI.COMM_NULL:
            self.INTERCOMM.Free()
            del self.INTERCOMM

    def testBarrier(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        self.INTERCOMM.Barrier()

    def testBcast(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            rmess = self.INTERCOMM.bcast(smess, root=MPI.ROOT)
                        else:
                            rmess = self.INTERCOMM.bcast(None, root=MPI.PROC_NULL)
                        self.assertEqual(rmess, None)
                else:
                    for root in range(rsize):
                        rmess = self.INTERCOMM.bcast(None, root=root)
                        self.assertEqual(rmess, smess)

    def testGather(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            rmess = self.INTERCOMM.gather(smess, root=MPI.ROOT)
                            self.assertEqual(rmess, [smess] * rsize)
                        else:
                            rmess = self.INTERCOMM.gather(None, root=MPI.PROC_NULL)
                            self.assertEqual(rmess, None)
                else:
                    for root in range(rsize):
                        rmess = self.INTERCOMM.gather(smess, root=root)
                        self.assertEqual(rmess, None)

    def testScatter(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            rmess = self.INTERCOMM.scatter([smess] * rsize, root=MPI.ROOT)
                        else:
                            rmess = self.INTERCOMM.scatter(None, root=MPI.PROC_NULL)
                        self.assertEqual(rmess, None)
                else:
                    for root in range(rsize):
                        rmess = self.INTERCOMM.scatter(None, root=root)
                        self.assertEqual(rmess, smess)

    def testAllgather(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            rmess = self.INTERCOMM.allgather(smess)
            self.assertEqual(rmess, [smess] * rsize)

    def testAlltoall(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for smess in messages + [messages]:
            rmess = self.INTERCOMM.alltoall([smess] * rsize)
            self.assertEqual(rmess, [smess] * rsize)

    def testReduce(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for op in (MPI.SUM, MPI.MAX, MPI.MIN, MPI.PROD):
            for color in [0, 1]:
                if self.COLOR == color:
                    for root in range(size):
                        if root == rank:
                            value = self.INTERCOMM.reduce(None, op=op, root=MPI.ROOT)
                            if op == MPI.SUM:
                                self.assertEqual(value, cumsum(range(rsize)))
                            elif op == MPI.PROD:
                                self.assertEqual(value, cumprod(range(rsize)))
                            elif op == MPI.MAX:
                                self.assertEqual(value, rsize-1)
                            elif op == MPI.MIN:
                                self.assertEqual(value, 0)
                        else:
                            value = self.INTERCOMM.reduce(None, op=op, root=MPI.PROC_NULL)
                            self.assertEqual(value, None)
                else:
                    for root in range(rsize):
                        value = self.INTERCOMM.reduce(rank, op=op, root=root)
                        self.assertEqual(value, None)

    def testAllreduce(self):
        if self.INTRACOMM == MPI.COMM_NULL: return
        if self.INTERCOMM == MPI.COMM_NULL: return
        rank = self.INTERCOMM.Get_rank()
        size = self.INTERCOMM.Get_size()
        rsize = self.INTERCOMM.Get_remote_size()
        for op in (MPI.SUM, MPI.MAX, MPI.MIN, MPI.PROD):
            value = self.INTERCOMM.allreduce(rank, None, op)
            if op == MPI.SUM:
                self.assertEqual(value, cumsum(range(rsize)))
            elif op == MPI.PROD:
                self.assertEqual(value, cumprod(range(rsize)))
            elif op == MPI.MAX:
                self.assertEqual(value, rsize-1)
            elif op == MPI.MIN:
                self.assertEqual(value, 0)

class TestCCOObjInter(BaseTestCCOObjInter, unittest.TestCase):
    BASECOMM = MPI.COMM_WORLD

class TestCCOObjInterDup(TestCCOObjInter):
    def setUp(self):
        self.BASECOMM = self.BASECOMM.Dup()
        super(TestCCOObjInterDup, self).setUp()
    def tearDown(self):
        self.BASECOMM.Free()
        del self.BASECOMM
        super(TestCCOObjInterDup, self).tearDown()

class TestCCOObjInterDupDup(TestCCOObjInterDup):
    BASECOMM = MPI.COMM_WORLD
    INTERCOMM_ORIG = MPI.COMM_NULL
    def setUp(self):
        super(TestCCOObjInterDupDup, self).setUp()
        if self.INTERCOMM == MPI.COMM_NULL: return
        self.INTERCOMM_ORIG = self.INTERCOMM
        self.INTERCOMM = self.INTERCOMM.Dup()
    def tearDown(self):
        super(TestCCOObjInterDupDup, self).tearDown()
        if self.INTERCOMM_ORIG == MPI.COMM_NULL: return
        self.INTERCOMM_ORIG.Free()

_name, _version = MPI.get_vendor()
if _name == 'Open MPI':
    if _version < (1, 6, 0):
        del BaseTestCCOObjInter
        del TestCCOObjInter
        del TestCCOObjInterDup
        del TestCCOObjInterDupDup
elif _name == "MPICH2" or _name == "Microsoft MPI":
    if _version <= (1, 0, 7):
        def _SKIPPED(*args, **kwargs): pass
        TestCCOObjInterDupDup.testBarrier   = _SKIPPED
        TestCCOObjInterDupDup.testAllgather = _SKIPPED
        TestCCOObjInterDupDup.testAllreduce = _SKIPPED
elif _name == "DeinoMPI":
    def _SKIPPED(*args, **kwargs): pass
    TestCCOObjInterDupDup.testBarrier   = _SKIPPED
    TestCCOObjInterDupDup.testAllgather = _SKIPPED
    TestCCOObjInterDupDup.testAllreduce = _SKIPPED
elif _name == "MPICH1":
    del BaseTestCCOObjInter
    del TestCCOObjInter
    del TestCCOObjInterDup
    del TestCCOObjInterDupDup
elif MPI.ROOT == MPI.PROC_NULL:
    del BaseTestCCOObjInter
    del TestCCOObjInter
    del TestCCOObjInterDup
    del TestCCOObjInterDupDup

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>"""store write permissions as meta information

Revision ID: ada15f85c3d0
Revises: 07fb4c284883
Create Date: 2022-01-25 10:30:07.890360

"""
# third party
from alembic import op
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = "ada15f85c3d0"
down_revision = "07fb4c284883"
branch_labels = None
depends_on = None

def upgrade() -> None:
    op.add_column(
        "obj_metadata",
        sa.Column("write_permissions", sa.JSON(), nullable=True),
    )

def downgrade() -> None:
    op.drop_column("obj_metadata", "write_permissions")
<commit_msg>grid: ignore op import from alembic<commit_after>"""store write permissions as meta information

Revision ID: ada15f85c3d0
Revises: 07fb4c284883
Create Date: 2022-01-25 10:30:07.890360

"""
# third party
from alembic import op  # type: ignore
import sqlalchemy as sa

# revision identifiers, used by Alembic.
revision = "ada15f85c3d0"
down_revision = "07fb4c284883"
branch_labels = None
depends_on = None

def upgrade() -> None:
    op.add_column(
        "obj_metadata",
        sa.Column("write_permissions", sa.JSON(), nullable=True),
    )

def downgrade() -> None:
    op.drop_column("obj_metadata", "write_permissions")
<|endoftext|>


<commit_before>import time
import numpy as np
from PyQt4 import QtGui, QtCore

from acq4.modules.Module import Module
import acq4.pyqtgraph as pg

class NoiseMonitor(Module):
    """ Used to monitor electrical noise over long time periods.
    """
    def __init__(self, manager, name, config):
        Module.__init__(self, manager, name, config)

        self.recordDir = None

        self.win = pg.LayoutWidget()

        self.startBtn = QtGui.QPushButton('Start')
        self.stopBtn = QtGui.QPushButton('Stop')
        self.win.addWidget(self.startBtn, 0, 0)
        self.win.addWidget(self.stopBtn, 0, 1)
        self.startBtn.clicked.connect(self.start)
        self.stopBtn.clicked.connect(self.stop)

        self.plot1 = pg.PlotWidget(labels={'left': ('Channel 1', 'V'), 'bottom': ('Time', 's')})
        self.plot2 = pg.PlotWidget(labels={'left': ('Channel 2', 'V'), 'bottom': ('Time', 's')})
        self.win.addWidget(self.plot1, 1, 0, 1, 2)
        self.win.addWidget(self.plot2, 2, 0, 1, 2)

        self.spectrogram = pg.ImageView()
        self.envelopePlot = pg.PlotWidget()
        self.win.addWidget(self.spectrogram, 0, 2, 2, 1)
        self.win.addWidget(self.envelopePlot, 2, 2, 1, 1)
        self.win.show()

        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self.runOnce)

    def start(self):
        self.recordDir = self.manager.getCurrentDir().mkdir('NoiseMonitor', autoIncrement=True)

        self.rawDataFile = None
        self.envelopeFile = None
        self.spectrogramFile = None
        self.startTime = time.time()

        self.timer.start(1000)

    def stop(self):
        self.timer.stop()

    def runOnce(self):
        dur = 0.02
        rate = 10000
        npts = int(dur * rate)
        cmd = {
            'protocol': {'duration': dur},
            'DAQ': {'rate': rate, 'numPts': npts},
            'Clamp1': {'mode': 'vc', 'holding': 0, 'command': np.zeros(npts)},
        }
        task = self.manager.createTask(cmd)
        task.execute()
        result = task.getResult()

        data = result['Clamp1']['Channel': 'primary']

        self.plot1.plot(data.asarray())

        ma = pg.metaarray.MetaArray(data.asarray()[np.newaxis, :], info=[
            {'name': 'Trial', 'units': 's', 'values': np.array([time.time()])}] + data._info)

        if self.rawDataFile is None:
            self.rawDataFile = self.recordDir.writeFile(ma, 'rawData', appendAxis='Trial', newFile=True)
        else:
            ma.write(self.rawDataFile.name(), appendAxis='Trial')

    def quit(self):
        self.stop()
        Module.quit(self)

<commit_msg>Add envelope and spectrum analysis to noise monitor<commit_after>import time
import numpy as np
from PyQt4 import QtGui, QtCore

from acq4.modules.Module import Module
import acq4.pyqtgraph as pg

class NoiseMonitor(Module):
    """ Used to monitor electrical noise over long time periods.
    """
    def __init__(self, manager, name, config):
        Module.__init__(self, manager, name, config)

        self.recordDir = None

        self.win = pg.LayoutWidget()

        self.startBtn = QtGui.QPushButton('Start')
        self.stopBtn = QtGui.QPushButton('Stop')
        self.win.addWidget(self.startBtn, 0, 0)
        self.win.addWidget(self.stopBtn, 0, 1)
        self.startBtn.clicked.connect(self.start)
        self.stopBtn.clicked.connect(self.stop)

        self.plot1 = pg.PlotWidget(labels={'left': ('Channel 1', 'V'), 'bottom': ('Time', 's')})
        self.plot2 = pg.PlotWidget(labels={'left': ('Channel 2', 'V'), 'bottom': ('Time', 's')})
        self.win.addWidget(self.plot1, 1, 0, 1, 2)
        self.win.addWidget(self.plot2, 2, 0, 1, 2)

        self.spectrogram = pg.ImageView()
        self.envelopePlot = pg.PlotWidget()
        self.win.addWidget(self.spectrogram, 0, 2, 2, 1)
        self.win.addWidget(self.envelopePlot, 2, 2, 1, 1)
        self.win.show()

        self.timer = QtCore.QTimer()
        self.timer.timeout.connect(self.runOnce)

    def start(self):
        self.recordDir = self.manager.getCurrentDir().mkdir('NoiseMonitor', autoIncrement=True)

        self.rawDataFile = None
        self.envelopeFile = None
        self.spectrogramFile = None
        self.startTime = time.time()

        self.timer.start(1000)

    def stop(self):
        self.timer.stop()

    def runOnce(self):
        dur = 0.02
        rate = 10000
        npts = int(dur * rate)
        cmd = {
            'protocol': {'duration': dur},
            'DAQ': {'rate': rate, 'numPts': npts},
            'Clamp1': {'mode': 'vc', 'holding': 0, 'command': np.zeros(npts)},
        }
        task = self.manager.createTask(cmd)
        task.execute()
        result = task.getResult()

        # Raw data plot/storage
        data = result['Clamp1']['Channel': 'primary']
        dataArr = data.asarray() + np.random.normal(size=data.shape[0], scale=1e-9)
        self.plot1.plot(data.xvals('Time'), dataArr, clear=True)
        trialArr = np.array([time.time() - self.startTime])
        ma = pg.metaarray.MetaArray(dataArr[np.newaxis, :], info=[
            {'name': 'Trial', 'units': 's', 'values': trialArr}] + data._info)

        if self.rawDataFile is None:
            self.rawDataFile = self.recordDir.writeFile(ma, 'rawData', appendAxis='Trial', newFile=True)
        else:
            ma.write(self.rawDataFile.name(), appendAxis='Trial')

        # Envelope analysis
        envData = np.array([[dataArr.min(), dataArr.mean(), dataArr.max(), dataArr.std()]])
        env = pg.metaarray.MetaArray(envData, info=[
            {'name': 'Trial', 'units': 's', 'values': trialArr},
            {'name': 'Metric', 'units': 'V'}])
        if self.envelopeFile is None:
            self.envelopeFile = self.recordDir.writeFile(env, 'envelope', appendAxis='Trial', newFile=True)
        else:
            env.write(self.envelopeFile.name(), appendAxis='Trial')

        envelope = self.envelopeFile.read()
        envArr = envelope.asarray()
        self.envelopePlot.clear()
        grey = (255, 255, 255, 100)
        c1 = self.envelopePlot.plot(envelope.xvals('Trial'), envArr[:,0], pen=grey)  # min
        c2 = self.envelopePlot.plot(envelope.xvals('Trial'), envArr[:,2], pen=grey)  # max
        fill = pg.FillBetweenItem(c1, c2, grey)
        self.envelopePlot.addItem(fill)
        c1 = self.envelopePlot.plot(envelope.xvals('Trial'), envArr[:,1] + envArr[:,3], pen=grey)  # +std
        c2 = self.envelopePlot.plot(envelope.xvals('Trial'), envArr[:,1] - envArr[:,3], pen=grey)  # -std
        fill = pg.FillBetweenItem(c1, c2, grey)
        self.envelopePlot.addItem(fill)

        self.envelopePlot.plot(envelope.xvals('Trial'), envArr[:,1])  # mean

        # Spectrum analysis
        fft = np.abs(np.fft.fft(dataArr))
        fft = fft[:len(fft)/2]
        freqArr = np.linspace(0, rate/2., len(fft))
        spec = pg.metaarray.MetaArray(fft[np.newaxis, :], info=[
            {'name': 'Trial', 'units': 's', 'values': trialArr},
            {'name': 'Frequency', 'units': 'Hz', 'values': freqArr}])
        if self.spectrogramFile is None:
            self.spectrogramFile = self.recordDir.writeFile(spec, 'spectrogram', appendAxis='Trial', newFile=True)
            auto = True
        else:
            spec.write(self.spectrogramFile.name(), appendAxis='Trial')
            auto = False

        specArr = self.spectrogramFile.read().asarray()
        self.spectrogram.setImage(specArr, autoLevels=auto)

    def quit(self):
        self.stop()
        Module.quit(self)

<|endoftext|>