<commit_before>from collections import defaultdict
from string import Formatter

from i3pystatus import IntervalModule

try:
    from natsort import natsorted as sorted
except ImportError:
    pass

class CpuUsage(IntervalModule):
    """
    Shows CPU usage.
    The first output will be inacurate.

    Linux only

    .. rubric:: Available formatters

    * `{usage}`      — usage average of all cores
    * `{usage_cpu*}` — usage of one specific core. replace "*" by core number starting at 0
    * `{usage_all}`  — usage of all cores separate. usess natsort when available(relevant for more than 10 cores)

    """

    format = "{usage:02}%"
    format_all = "{core}:{usage:02}%"
    exclude_average = False
    interval = 1
    color = None
    settings = (
        ("format", "format string."),
        ("format_all", ("format string used for {usage_all} per core. "
                        "Available formaters are {core} and {usage}. ")),
        ("exclude_average", ("If True usage average of all cores will "
                             "not be in format_all.")),
        ("color", "HTML color code #RRGGBB")
    )

    def init(self):
        self.prev_total = defaultdict(int)
        self.prev_busy = defaultdict(int)
        self.formatter = Formatter()

    def get_cpu_timings(self):
        """
        reads and parses /proc/stat
        returns dictionary with all available cores including global average
        """
        timings = {}
        with open('/proc/stat', 'r') as file_obj:
            for line in file_obj:
                if 'cpu' in line:
                    line = line.strip().split()
                    timings[line[0]] = [int(x) for x in line[1:]]

        return timings

    def calculate_usage(self, cpu, total, busy):
        """
        calculates usage
        """
        diff_total = total - self.prev_total[cpu]
        diff_busy = busy - self.prev_busy[cpu]

        self.prev_total[cpu] = total
        self.prev_busy[cpu] = busy

        if diff_total == 0:
            return 0
        else:
            return int(diff_busy / diff_total * 100)

    def gen_format_all(self, usage):
        """
        generates string for format all
        """
        format_string = " "
        core_strings = []
        for core, usage in usage.items():
            if core == 'usage_cpu' and self.exclude_average:
                continue
            elif core == 'usage':
                continue

            core = core.replace('usage_', '')
            string = self.formatter.format(format_string=self.format_all,
                                           core=core,
                                           usage=usage)
            core_strings.append(string)

        core_strings = sorted(core_strings)

        return format_string.join(core_strings)

    def get_usage(self):
        """
        parses /proc/stat and calcualtes total and busy time
        (more specific USER_HZ see man 5 proc for further informations )
        """
        usage = {}

        for cpu, timings in self.get_cpu_timings().items():
            cpu_total = sum(timings)
            del timings[3:5]
            cpu_busy = sum(timings)
            cpu_usage = self.calculate_usage(cpu, cpu_total, cpu_busy)

            usage['usage_' + cpu] = cpu_usage

        return usage

    def run(self):
        usage = self.get_usage()
        usage['format_all'] = self.gen_format_all(usage)

        # for backward compatibility
        usage['usage'] = usage['usage_cpu']

        self.data = usage
        self.output = {
            "full_text": self.format.format_map(usage),
            "color": self.color
        }
<commit_msg>Fix wrong dictionary key<commit_after>from collections import defaultdict
from string import Formatter

from i3pystatus import IntervalModule

try:
    from natsort import natsorted as sorted
except ImportError:
    pass

class CpuUsage(IntervalModule):
    """
    Shows CPU usage.
    The first output will be inacurate.

    Linux only

    .. rubric:: Available formatters

    * `{usage}`      — usage average of all cores
    * `{usage_cpu*}` — usage of one specific core. replace "*" by core number starting at 0
    * `{usage_all}`  — usage of all cores separate. usess natsort when available(relevant for more than 10 cores)

    """

    format = "{usage:02}%"
    format_all = "{core}:{usage:02}%"
    exclude_average = False
    interval = 1
    color = None
    settings = (
        ("format", "format string."),
        ("format_all", ("format string used for {usage_all} per core. "
                        "Available formaters are {core} and {usage}. ")),
        ("exclude_average", ("If True usage average of all cores will "
                             "not be in format_all.")),
        ("color", "HTML color code #RRGGBB")
    )

    def init(self):
        self.prev_total = defaultdict(int)
        self.prev_busy = defaultdict(int)
        self.formatter = Formatter()

    def get_cpu_timings(self):
        """
        reads and parses /proc/stat
        returns dictionary with all available cores including global average
        """
        timings = {}
        with open('/proc/stat', 'r') as file_obj:
            for line in file_obj:
                if 'cpu' in line:
                    line = line.strip().split()
                    timings[line[0]] = [int(x) for x in line[1:]]

        return timings

    def calculate_usage(self, cpu, total, busy):
        """
        calculates usage
        """
        diff_total = total - self.prev_total[cpu]
        diff_busy = busy - self.prev_busy[cpu]

        self.prev_total[cpu] = total
        self.prev_busy[cpu] = busy

        if diff_total == 0:
            return 0
        else:
            return int(diff_busy / diff_total * 100)

    def gen_format_all(self, usage):
        """
        generates string for format all
        """
        format_string = " "
        core_strings = []
        for core, usage in usage.items():
            if core == 'usage_cpu' and self.exclude_average:
                continue
            elif core == 'usage':
                continue

            core = core.replace('usage_', '')
            string = self.formatter.format(format_string=self.format_all,
                                           core=core,
                                           usage=usage)
            core_strings.append(string)

        core_strings = sorted(core_strings)

        return format_string.join(core_strings)

    def get_usage(self):
        """
        parses /proc/stat and calcualtes total and busy time
        (more specific USER_HZ see man 5 proc for further informations )
        """
        usage = {}

        for cpu, timings in self.get_cpu_timings().items():
            cpu_total = sum(timings)
            del timings[3:5]
            cpu_busy = sum(timings)
            cpu_usage = self.calculate_usage(cpu, cpu_total, cpu_busy)

            usage['usage_' + cpu] = cpu_usage

        return usage

    def run(self):
        usage = self.get_usage()
        usage['usage_all'] = self.gen_format_all(usage)

        # for backward compatibility
        usage['usage'] = usage['usage_cpu']

        self.data = usage
        self.output = {
            "full_text": self.format.format_map(usage),
            "color": self.color
        }
<|endoftext|>


<commit_before># Lint as: python3
# Copyright 2020 Google LLC
#
# 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.
"""Tests for google3.third_party.gps_building_blocks.cloud.utils.cloud_composer."""

from typing import Dict
import unittest
from unittest import mock
from googleapiclient import errors
from googleapiclient import http

from gps_building_blocks.cloud.utils import cloud_auth
from gps_building_blocks.cloud.utils import cloud_composer
from gps_building_blocks.cloud.utils import utils

class CloudComposerUtilsTest(unittest.TestCase):

  def setUp(self):
    """Creates mock objects for googleapi client."""
    super(CloudComposerUtilsTest, self).setUp()
    self.addCleanup(mock.patch.stopall)
    self.project_id = 'project_id'
    self.location = cloud_composer._LOCATION
    self.environment_name = 'environment_name'
    self.zone = 'a'
    self.mock_build_service_client = mock.patch.object(
        cloud_auth, 'build_service_client', autospec=True).start()
    self.mock_wait_for_operation = mock.patch.object(
        utils, 'wait_for_operation', autospec=True).start()
    self.mock_execute_request = mock.patch.object(
        utils, 'execute_request', autospec=True).start()
    self.mock_client = mock.Mock()
    self.mock_build_service_client.return_value = self.mock_client

    self.service_account_key_file = '/tmp/service_account_key.json'
    self.composer = cloud_composer.CloudComposerUtils(
        project_id=self.project_id,
        location=self.location,
        service_account_key_file=self.service_account_key_file)
    self.operation_client = mock.Mock()
    self.operation = {}
    (self.mock_client.projects.return_value.locations.return_value.operations
     .return_value) = self.operation_client
    self.mock_execute_request.return_value = self.operation
    self.http_error = errors.HttpError(mock.MagicMock(status=400), b'')
    self.fully_qualified_name = (f'projects/{self.project_id}/locations/'
                                 f'{self.location}/environments/'
                                 f'{self.environment_name}')
    self.mock_environment_client = (
        self.mock_client.projects.return_value.locations.return_value
        .environments)
    self.mock_request = mock.Mock(http.HttpRequest)

  @mock.patch.object(cloud_auth, 'build_impersonated_client', autospec=True)
  def test_client_initializes_with_impersonated_service_account(
      self, mock_impersonated_client):
    service_account_name = 'my-svc-account@project-id.iam.gserviceaccount.com'
    version = 'v1beta1'

    cloud_composer.CloudComposerUtils(
        project_id=self.project_id,
        service_account_name=service_account_name,
        version=version)

    mock_impersonated_client.assert_called_once_with('composer',
                                                     service_account_name,
                                                     version)

  def test_get_fully_qualified_environment_name(self):
    expected_name = (f'projects/{self.project_id}/locations/'
                     f'{self.location}/environments/{self.environment_name}')

    fully_qualified_name = self.composer._get_fully_qualified_env_name(
        self.environment_name)

    self.assertEqual(fully_qualified_name, expected_name)

  def test_create_composer_environment(self):
    mock_create_env = self.mock_environment_client.return_value.create
    disk_size_gb = cloud_composer._DISC_SIZE
    machine_type = cloud_composer._MACHINE_TYPE
    python_version = cloud_composer._PYTHON_VERSION

    self.composer.create_environment(self.environment_name, self.zone,
                                     disk_size_gb, machine_type)

    composer_zone = '{}-{}'.format(self.location, self.zone)
    location = 'projects/{}/zones/{}'.format(self.project_id, composer_zone)
    parent = 'projects/{}/locations/{}'.format(self.project_id, self.location)
    machine_type = 'projects/{}/zones/{}/machineTypes/{}'.format(
        self.project_id, composer_zone, machine_type)
    request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'nodeConfig': {
                'location': location,
                'machineType': machine_type,
                'diskSizeGb': disk_size_gb
            },
            'softwareConfig': {
                'pythonVersion': python_version
            }
        }
    }
    mock_create_env.assert_called_once_with(parent=parent, body=request_body)
    self.mock_build_service_client.assert_called_once_with(
        'composer', self.service_account_key_file)
    self.assertEqual(self.mock_client, self.composer.client)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_create_composer_environment_raises_error(self):
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.create_environment(self.environment_name, self.zone)

  def test_create_environment_raises_error_if_disk_size_is_less_than_minimum(
      self):
    with self.assertRaises(cloud_composer.Error):
      self.composer.create_environment(self.environment_name, disk_size_gb=10)

  def test_create_environment_skips_if_already_exists(self):
    http_error_conflict = errors.HttpError(mock.MagicMock(status=409), b'')
    self.mock_client.projects.side_effect = http_error_conflict

    self.composer.create_environment(self.environment_name)

  def test_create_environment_with_specific_composer_image_version(self):
    image_version = 'composer-1.9.2-airflow-1.10.2'
    mock_create_env = self.mock_environment_client.return_value.create

    self.composer.create_environment(
        self.environment_name, image_version=image_version)

    _, mock_kwargs = mock_create_env.call_args
    self.assertEqual(
        mock_kwargs['body']['config']['softwareConfig']['imageVersion'],
        image_version)

  def test_install_python_packages(self):
    python_packages = {'lib1': '<=1.0.1', 'lib2': '==2.12.0', 'lib3': '>1.0.3'}
    mock_update_environment = self.mock_environment_client.return_value.patch
    mock_update_environment.return_value = self.mock_request

    self.composer.install_python_packages(self.environment_name,
                                          python_packages)

    request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'softwareConfig': {
                'pypiPackages': python_packages
            }
        }
    }
    mock_update_environment.assert_called_once_with(
        name=self.fully_qualified_name,
        body=request_body,
        updateMask='config.softwareConfig.pypiPackages')
    self.mock_execute_request.assert_called_once_with(self.mock_request)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_install_python_packages_raises_error_on_empty_packages(self):
    python_packages = []

    with self.assertRaises(cloud_composer.Error):
      self.composer.install_python_packages(self.environment_name,
                                            python_packages)

  def test_install_python_packages_raises_error_on_http_error(self):
    python_packages = {'lib1': '==1.0.1'}
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.install_python_packages(self.environment_name,
                                            python_packages)

  def test_set_environment_variables(self):
    environment_variables = {'name1': 'value1', 'name2': 'value2'}

    mock_update_environment = self.mock_environment_client.return_value.patch
    mock_update_environment.return_value = self.mock_request

    self.composer.set_environment_variables(self.environment_name,
                                            environment_variables)

    request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'softwareConfig': {
                'envVariables': environment_variables
            }
        }
    }
    mock_update_environment.assert_called_once_with(
        name=self.fully_qualified_name,
        body=request_body,
        updateMask='config.softwareConfig.envVariables')
    self.mock_execute_request.assert_called_once_with(self.mock_request)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_set_environment_variables_raises_error_on_http_error(self):
    environment_variables = {'name1': 'value1', 'name2': 'value2'}
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.set_environment_variables(self.environment_name,
                                              environment_variables)

  def test_override_airflow_configs(self):
    airflow_config_overrides = {
        'smtp-smtp_mail_from': 'no-reply@abc.com',
        'core-dags_are_paused_at_creation': 'True'
    }
    expected_request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'softwareConfig': {
                'airflowConfigOverrides': airflow_config_overrides
            }
        }
    }
    mock_update_environment = self.mock_environment_client.return_value.patch
    mock_update_environment.return_value = self.mock_request

    self.composer.override_airflow_configs(self.environment_name,
                                           airflow_config_overrides)

    mock_update_environment.assert_called_once_with(
        name=self.fully_qualified_name,
        body=expected_request_body,
        updateMask='config.softwareConfig.airflowConfigOverrides')
    self.mock_execute_request.assert_called_once_with(self.mock_request)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_override_airflow_configs_raises_error_on_http_error(self):
    airflow_config_overrides = {}
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.override_airflow_configs(self.environment_name,
                                             airflow_config_overrides)

  def test_get_environment(self):
    mock_environment = mock.Mock(Dict[str, str])
    self.mock_environment_client.return_value.get.return_value = (
        self.mock_request)
    self.mock_execute_request.return_value = mock_environment

    composer_environment = self.composer.get_environment(self.environment_name)

    self.assertEqual(composer_environment, mock_environment)
    self.mock_execute_request.assert_called_once_with(self.mock_request)

  def test_get_environment_raises_error_on_http_error(self):
    self.mock_environment_client.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.get_environment(self.environment_name)

  def test_get_dags_folder(self):
    dags_folder = 'gs://bucket_name/dags'
    composer_env_details = {'config': {'dagGcsPrefix': dags_folder}}
    self.mock_environment_client.return_value.get.return_value = (
        self.mock_request)
    self.mock_execute_request.return_value = composer_env_details

    actual_dags_folder = self.composer.get_dags_folder(self.environment_name)

    self.assertEqual(actual_dags_folder, dags_folder)

  def test_delete_environment(self):
    mock_delete_env = self.mock_environment_client.return_value.delete

    self.composer.delete_environment(self.environment_name)

    request_body = {
        'name': self.fully_qualified_name
    }
    mock_delete_env.assert_called_once_with(body=request_body)
    self.mock_build_service_client.assert_called_once_with(
        'composer', self.service_account_key_file)
    self.assertEqual(self.mock_client, self.composer.client)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_delete_environment_raises_error_on_http_error(self):
    self.mock_environment_client.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.delete_environment(self.environment_name)

  def test_delete_environment_sliently_skips_on_http_error(self):
    http_error_conflict = errors.HttpError(mock.MagicMock(status=404), b'')
    self.mock_client.projects.side_effect = http_error_conflict

    self.composer.delete_environment(self.environment_name)

if __name__ == '__main__':
  unittest.main()
<commit_msg>Internal change<commit_after># Lint as: python3
# Copyright 2020 Google LLC
#
# 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.
"""Tests for google3.third_party.gps_building_blocks.cloud.utils.cloud_composer."""

from typing import Dict
import unittest
from unittest import mock
from googleapiclient import errors
from googleapiclient import http

from gps_building_blocks.cloud.utils import cloud_auth
from gps_building_blocks.cloud.utils import cloud_composer
from gps_building_blocks.cloud.utils import utils

class CloudComposerUtilsTest(unittest.TestCase):

  def setUp(self):
    """Creates mock objects for googleapi client."""
    super(CloudComposerUtilsTest, self).setUp()
    self.addCleanup(mock.patch.stopall)
    self.project_id = 'project_id'
    self.location = cloud_composer._LOCATION
    self.environment_name = 'environment_name'
    self.zone = 'a'
    self.mock_build_service_client = mock.patch.object(
        cloud_auth, 'build_service_client', autospec=True).start()
    self.mock_wait_for_operation = mock.patch.object(
        utils, 'wait_for_operation', autospec=True).start()
    self.mock_execute_request = mock.patch.object(
        utils, 'execute_request', autospec=True).start()
    self.mock_client = mock.Mock()
    self.mock_build_service_client.return_value = self.mock_client

    self.service_account_key_file = '/tmp/service_account_key.json'
    self.composer = cloud_composer.CloudComposerUtils(
        project_id=self.project_id,
        location=self.location,
        service_account_key_file=self.service_account_key_file)
    self.operation_client = mock.Mock()
    self.operation = {}
    (self.mock_client.projects.return_value.locations.return_value.operations
     .return_value) = self.operation_client
    self.mock_execute_request.return_value = self.operation
    self.http_error = errors.HttpError(mock.MagicMock(status=400), b'')
    self.fully_qualified_name = (f'projects/{self.project_id}/locations/'
                                 f'{self.location}/environments/'
                                 f'{self.environment_name}')
    self.mock_environment_client = (
        self.mock_client.projects.return_value.locations.return_value
        .environments)
    self.mock_request = mock.Mock(http.HttpRequest)

  @mock.patch.object(cloud_auth, 'build_impersonated_client', autospec=True)
  def test_client_initializes_with_impersonated_service_account(
      self, mock_impersonated_client):
    service_account_name = 'my-svc-account@project-id.iam.gserviceaccount.com'
    version = 'v1beta1'

    cloud_composer.CloudComposerUtils(
        project_id=self.project_id,
        service_account_name=service_account_name,
        version=version)

    mock_impersonated_client.assert_called_once_with('composer',
                                                     service_account_name,
                                                     version)

  def test_get_fully_qualified_environment_name(self):
    expected_name = (f'projects/{self.project_id}/locations/'
                     f'{self.location}/environments/{self.environment_name}')

    fully_qualified_name = self.composer._get_fully_qualified_env_name(
        self.environment_name)

    self.assertEqual(fully_qualified_name, expected_name)

  def test_create_composer_environment(self):
    mock_create_env = self.mock_environment_client.return_value.create
    disk_size_gb = cloud_composer._DISC_SIZE
    machine_type = cloud_composer._MACHINE_TYPE
    python_version = cloud_composer._PYTHON_VERSION

    self.composer.create_environment(self.environment_name, self.zone,
                                     disk_size_gb, machine_type)

    composer_zone = '{}-{}'.format(self.location, self.zone)
    location = 'projects/{}/zones/{}'.format(self.project_id, composer_zone)
    parent = 'projects/{}/locations/{}'.format(self.project_id, self.location)
    machine_type = 'projects/{}/zones/{}/machineTypes/{}'.format(
        self.project_id, composer_zone, machine_type)
    request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'nodeConfig': {
                'location': location,
                'machineType': machine_type,
                'diskSizeGb': disk_size_gb
            },
            'softwareConfig': {
                'pythonVersion': python_version
            }
        }
    }
    mock_create_env.assert_called_once_with(parent=parent, body=request_body)
    self.mock_build_service_client.assert_called_once_with(
        'composer', self.service_account_key_file)
    self.assertEqual(self.mock_client, self.composer.client)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_create_composer_environment_raises_error(self):
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.create_environment(self.environment_name, self.zone)

  def test_create_environment_raises_error_if_disk_size_is_less_than_minimum(
      self):
    with self.assertRaises(cloud_composer.Error):
      self.composer.create_environment(self.environment_name, disk_size_gb=10)

  def test_create_environment_skips_if_already_exists(self):
    http_error_conflict = errors.HttpError(mock.MagicMock(status=409), b'')
    self.mock_client.projects.side_effect = http_error_conflict

    self.composer.create_environment(self.environment_name)

  def test_create_environment_with_specific_composer_image_version(self):
    image_version = 'composer-1.9.2-airflow-1.10.2'
    mock_create_env = self.mock_environment_client.return_value.create

    self.composer.create_environment(
        self.environment_name, image_version=image_version)

    _, mock_kwargs = mock_create_env.call_args
    self.assertEqual(
        mock_kwargs['body']['config']['softwareConfig']['imageVersion'],
        image_version)

  def test_install_python_packages(self):
    python_packages = {'lib1': '<=1.0.1', 'lib2': '==2.12.0', 'lib3': '>1.0.3'}
    mock_update_environment = self.mock_environment_client.return_value.patch
    mock_update_environment.return_value = self.mock_request

    self.composer.install_python_packages(self.environment_name,
                                          python_packages)

    request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'softwareConfig': {
                'pypiPackages': python_packages
            }
        }
    }
    mock_update_environment.assert_called_once_with(
        name=self.fully_qualified_name,
        body=request_body,
        updateMask='config.softwareConfig.pypiPackages')
    self.mock_execute_request.assert_called_once_with(self.mock_request)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_install_python_packages_raises_error_on_empty_packages(self):
    python_packages = []

    with self.assertRaises(cloud_composer.Error):
      self.composer.install_python_packages(self.environment_name,
                                            python_packages)

  def test_install_python_packages_raises_error_on_http_error(self):
    python_packages = {'lib1': '==1.0.1'}
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.install_python_packages(self.environment_name,
                                            python_packages)

  def test_set_environment_variables(self):
    environment_variables = {'name1': 'value1', 'name2': 'value2'}

    mock_update_environment = self.mock_environment_client.return_value.patch
    mock_update_environment.return_value = self.mock_request

    self.composer.set_environment_variables(self.environment_name,
                                            environment_variables)

    request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'softwareConfig': {
                'envVariables': environment_variables
            }
        }
    }
    mock_update_environment.assert_called_once_with(
        name=self.fully_qualified_name,
        body=request_body,
        updateMask='config.softwareConfig.envVariables')
    self.mock_execute_request.assert_called_once_with(self.mock_request)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_set_environment_variables_raises_error_on_http_error(self):
    environment_variables = {'name1': 'value1', 'name2': 'value2'}
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.set_environment_variables(self.environment_name,
                                              environment_variables)

  def test_override_airflow_configs(self):
    airflow_config_overrides = {
        'smtp-smtp_mail_from': 'no-reply@abc.com',
        'core-dags_are_paused_at_creation': 'True'
    }
    expected_request_body = {
        'name': self.fully_qualified_name,
        'config': {
            'softwareConfig': {
                'airflowConfigOverrides': airflow_config_overrides
            }
        }
    }
    mock_update_environment = self.mock_environment_client.return_value.patch
    mock_update_environment.return_value = self.mock_request

    self.composer.override_airflow_configs(self.environment_name,
                                           airflow_config_overrides)

    mock_update_environment.assert_called_once_with(
        name=self.fully_qualified_name,
        body=expected_request_body,
        updateMask='config.softwareConfig.airflowConfigOverrides')
    self.mock_execute_request.assert_called_once_with(self.mock_request)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_override_airflow_configs_raises_error_on_http_error(self):
    airflow_config_overrides = {}
    self.mock_client.projects.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.override_airflow_configs(self.environment_name,
                                             airflow_config_overrides)

  def test_get_environment(self):
    mock_environment = mock.Mock(Dict[str, str])
    self.mock_environment_client.return_value.get.return_value = (
        self.mock_request)
    self.mock_execute_request.return_value = mock_environment

    composer_environment = self.composer.get_environment(self.environment_name)

    self.assertEqual(composer_environment, mock_environment)
    self.mock_execute_request.assert_called_once_with(self.mock_request)

  def test_get_environment_raises_error_on_http_error(self):
    self.mock_environment_client.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.get_environment(self.environment_name)

  def test_get_dags_folder(self):
    dags_folder = 'gs://bucket_name/dags'
    composer_env_details = {'config': {'dagGcsPrefix': dags_folder}}
    self.mock_environment_client.return_value.get.return_value = (
        self.mock_request)
    self.mock_execute_request.return_value = composer_env_details

    actual_dags_folder = self.composer.get_dags_folder(self.environment_name)

    self.assertEqual(actual_dags_folder, dags_folder)

  def test_delete_environment(self):
    mock_delete_env = self.mock_environment_client.return_value.delete

    self.composer.delete_environment(self.environment_name)

    fully_qualified_name = self.fully_qualified_name

    mock_delete_env.assert_called_once_with(name=fully_qualified_name)
    self.mock_build_service_client.assert_called_once_with(
        'composer', self.service_account_key_file)
    self.assertEqual(self.mock_client, self.composer.client)
    self.mock_wait_for_operation.assert_called_once_with(
        self.operation_client, self.operation)

  def test_delete_environment_raises_error_on_http_error(self):
    self.mock_environment_client.side_effect = self.http_error

    with self.assertRaises(cloud_composer.Error):
      self.composer.delete_environment(self.environment_name)

  def test_delete_environment_sliently_skips_on_http_error(self):
    http_error_conflict = errors.HttpError(mock.MagicMock(status=404), b'')
    self.mock_client.projects.side_effect = http_error_conflict

    self.composer.delete_environment(self.environment_name)

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


<commit_before>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2018 CNRS

# 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.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import itertools
import numpy as np
from tqdm import tqdm
import torch
from torch.autograd import Variable
import torch.nn.functional as F
from pyannote.audio.generators.speaker import SpeechSegmentGenerator
from pyannote.audio.callback import LoggingCallbackPytorch
from torch.optim import Adam
from pyannote.audio.embedding.utils import to_condensed
from scipy.spatial.distance import pdist, squareform
from pyannote.metrics.binary_classification import det_curve

class TripletLoss(object):
    """

    delta = d(anchor, positive) - d(anchor, negative)

    * with 'positive' clamping:
        loss = max(0, delta + margin x D)
    * with 'sigmoid' clamping:
        loss = sigmoid(10 * delta)

    where d(., .) varies in range [0, D] (e.g. D=2 for euclidean distance).

    Parameters
    ----------
    duration : float, optional
        Defautls to 3.2 seconds.
    metric : {'euclidean', 'cosine', 'angular'}, optional
        Defaults to 'cosine'.
    margin: float, optional
        Margin factor. Defaults to 0.2.
    clamp : {'positive', 'sigmoid', 'softmargin'}, optional
        Defaults to 'positive'.
    sampling : {'all', 'hard', 'negative'}, optional
        Triplet sampling strategy.
    per_label : int, optional
        Number of sequences per speaker in each batch. Defaults to 3.
    per_fold : int, optional
        If provided, sample triplets from groups of `per_fold` speakers at a
        time. Defaults to sample triplets from the whole speaker set.
    trim : float, optional
        Do not use speech segments that are that close to the beginning/end of
        the annotated region. Useful when features are short-term normalized,
        as this can result in weird feature distribution at the boundaries.
        Defaults to 0s (i.e. do no trim).
    parallel : int, optional
        Number of prefetching background generators. Defaults to 1.
        Each generator will prefetch enough batches to cover a whole epoch.
        Set `parallel` to 0 to not use background generators.
    """

    def __init__(self, duration=3.2,
                 metric='cosine', margin=0.2, clamp='positive',
                 sampling='all', per_label=3, per_fold=None,
                 trim=0, parallel=1):

        super(TripletLoss, self).__init__()

        self.metric = metric
        self.margin = margin

        if self.metric == 'cosine':
            self.margin_ = self.margin * 2.
        elif self.metric == 'angular':
            self.margin_ = self.margin * np.pi
        elif self.metric == 'euclidean':
            self.margin_ = self.margin * 2.
        else:
            msg = "'metric' must be one of {'euclidean', 'cosine', 'angular'}."
            raise ValueError(msg)

        if clamp not in {'positive', 'sigmoid', 'softmargin'}:
            msg = "'clamp' must be one of {'positive', 'sigmoid', 'softmargin'}."
            raise ValueError(msg)
        self.clamp = clamp

        if sampling not in {'all', 'hard', 'negative'}:
            msg = "'sampling' must be one of {'all', 'hard', 'negative'}."
            raise ValueError(msg)
        self.sampling = sampling

        self.per_fold = per_fold
        self.per_label = per_label
        self.duration = duration
        self.trim = trim
        self.parallel = parallel

    def pdist(self, fX):
        """Compute pdist à-la scipy.spatial.distance.pdist

        Parameters
        ----------
        fX : (n, d) torch.autograd.Variable
            Embeddings.

        Returns
        -------
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix
        """

        n_sequences, _ = fX.size()
        distances = []

        for i in range(n_sequences - 1):

            if self.metric in ('cosine', 'angular'):
                d = 1. - F.cosine_similarity(
                    fX[i, :].expand(n_sequences - 1 - i, -1),
                    fX[i+1:, :], dim=1, eps=1e-8)

                if self.metric == 'angular':
                    d = torch.acos(torch.clamp(1. - d, -1 + 1e-6, 1 - 1e-6))

            elif self.metric == 'euclidean':
                d = F.pairwise_distance(
                    fX[i, :].expand(n_sequences - 1 - i, -1),
                    fX[i+1:, :], p=2, eps=1e-06).view(-1)

            distances.append(d)

        return torch.cat(distances)

    def batch_hard(self, y, distances):
        """Build triplet with both hardest positive and hardest negative

        Parameters
        ----------
        y : list
            Sequence labels.
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix

        Returns
        -------
        anchors, positives, negatives : list of int
            Triplets indices.
        """

        anchors, positives, negatives = [], [], []

        if distances.is_cuda:
            distances = squareform(distances.data.cpu().numpy())
        else:
            distances = squareform(distances.data.numpy())
        y = np.array(y)

        for anchor, y_anchor in enumerate(y):

            d = distances[anchor]

            # hardest positive
            pos = np.where(y == y_anchor)[0]
            pos = [p for p in pos if p != anchor]
            positive = int(pos[np.argmax(d[pos])])

            # hardest negative
            neg = np.where(y != y_anchor)[0]
            negative = int(neg[np.argmin(d[neg])])

            anchors.append(anchor)
            positives.append(positive)
            negatives.append(negative)

        return anchors, positives, negatives

    def batch_negative(self, y, distances):
        """Build triplet with hardest negative

        Parameters
        ----------
        y : list
            Sequence labels.
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix

        Returns
        -------
        anchors, positives, negatives : list of int
            Triplets indices.
        """

        anchors, positives, negatives = [], [], []

        if distances.is_cuda:
            distances = squareform(distances.data.cpu().numpy())
        else:
            distances = squareform(distances.data.numpy())
        y = np.array(y)

        for anchor, y_anchor in enumerate(y):

            # hardest negative
            d = distances[anchor]
            neg = np.where(y != y_anchor)[0]
            negative = int(neg[np.argmin(d[neg])])

            for positive in np.where(y == y_anchor)[0]:
                if positive == anchor:
                    continue

                anchors.append(anchor)
                positives.append(positive)
                negatives.append(negative)

        return anchors, positives, negatives

    def batch_all(self, y, distances):
        """Build all possible triplet

        Parameters
        ----------
        y : list
            Sequence labels.
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix

        Returns
        -------
        anchors, positives, negatives : list of int
            Triplets indices.
        """

        anchors, positives, negatives = [], [], []

        for anchor, y_anchor in enumerate(y):
            for positive, y_positive in enumerate(y):

                # if same embedding or different labels, skip
                if (anchor == positive) or (y_anchor != y_positive):
                    continue

                for negative, y_negative in enumerate(y):

                    if y_negative == y_anchor:
                        continue

                    anchors.append(anchor)
                    positives.append(positive)
                    negatives.append(negative)

        return anchors, positives, negatives

    def triplet_loss(self, distances, anchors, positives, negatives,
                     return_delta=False):
        """Compute triplet loss

        Parameters
        ----------
        distances : torch.autograd.Variable
            Condensed matrix of pairwise distances.
        anchors, positives, negatives : list of int
            Triplets indices.
        return_delta : bool, optional
            Return delta before clamping.

        Returns
        -------
        loss : torch.autograd.Variable
            Triplet loss.
        """

        # estimate total number of embeddings from pdist shape
        n = int(.5 * (1 + np.sqrt(1 + 8 * len(distances))))
        n = [n] * len(anchors)

        # convert indices from squared matrix
        # to condensed matrix referential
        pos = list(map(to_condensed, n, anchors, positives))
        neg = list(map(to_condensed, n, anchors, negatives))

        # compute raw triplet loss (no margin, no clamping)
        # the lower, the better
        delta = distances[pos] - distances[neg]

        # clamp triplet loss
        if self.clamp == 'positive':
            loss = torch.clamp(delta + self.margin_, min=0)

        elif self.clamp == 'softmargin':
            loss = torch.log1p(torch.exp(delta))

        elif self.clamp == 'sigmoid':
            loss = F.sigmoid(10 * delta)

        # return triplet losses
        if return_delta:
            return loss, delta.view((-1, 1)), pos, neg
        else:
            return loss

    def fit(self, model, feature_extraction, protocol, log_dir, subset='train',
            epochs=1000, restart=0, gpu=False):

        import tensorboardX
        writer = tensorboardX.SummaryWriter(log_dir=log_dir)

        logging_callback = LoggingCallbackPytorch(log_dir=log_dir,
                                                  restart=restart > 0)

        batch_generator = SpeechSegmentGenerator(
            feature_extraction,
            per_label=self.per_label, per_fold=self.per_fold,
            duration=self.duration, trim=self.trim,
            parallel=self.parallel)
        batches = batch_generator(protocol, subset=subset)
        batch = next(batches)

        batches_per_epoch = batch_generator.batches_per_epoch

        if restart > 0:
            weights_pt = logging_callback.WEIGHTS_PT.format(
                log_dir=log_dir, epoch=restart)
            model.load_state_dict(torch.load(weights_pt))

        if gpu:
            model = model.cuda()

        model.internal = False

        optimizer = Adam(model.parameters())
        if restart > 0:
            optimizer_pt = logging_callback.OPTIMIZER_PT.format(
                log_dir=log_dir, epoch=restart)
            optimizer.load_state_dict(torch.load(optimizer_pt))
            if gpu:
                for state in optimizer.state.values():
                    for k, v in state.items():
                        if torch.is_tensor(v):
                            state[k] = v.cuda()

        epoch = restart if restart > 0 else -1
        while True:
            epoch += 1
            if epoch > epochs:
                break

            tloss_avg = 0.

            if epoch % 10 == 0:
                positive, negative = [], []
                norms = []

            desc = 'Epoch #{0}'.format(epoch)
            for i in tqdm(range(batches_per_epoch), desc=desc):

                model.zero_grad()

                batch = next(batches)

                X = Variable(torch.from_numpy(
                    np.array(np.rollaxis(batch['X'], 0, 2),
                             dtype=np.float32)))

                if gpu:
                    X = X.cuda()

                fX = model(X)

                if epoch % 10 == 0:
                    if gpu:
                        fX_ = fX.data.cpu().numpy()
                    else:
                        fX_ = fX.data.numpy()
                    norms.append(np.linalg.norm(fX_, axis=1))

                # pre-compute pairwise distances
                distances = self.pdist(fX)

                if epoch % 10 == 0:
                    if gpu:
                        distances_ = distances.data.cpu().numpy()
                    else:
                        distances_ = distances.data.numpy()
                    is_positive = pdist(batch['y'].reshape((-1, 1)), metric='chebyshev') < 1
                    positive.append(distances_[np.where(is_positive)])
                    negative.append(distances_[np.where(~is_positive)])

                # sample triplets
                triplets = getattr(self, 'batch_{0}'.format(self.sampling))
                anchors, positives, negatives = triplets(batch['y'], distances)

                # compute triplet loss
                losses = self.triplet_loss(
                    distances, anchors, positives, negatives)

                loss = torch.mean(losses)

                # log loss
                if gpu:
                    loss_ = float(loss.data.cpu().numpy())
                else:
                    loss_ = float(loss.data.numpy())
                tloss_avg += loss_

                loss.backward()
                optimizer.step()

            tloss_avg /= batches_per_epoch
            writer.add_scalar('tloss', tloss_avg, global_step=epoch)

            if epoch % 10 == 0:

                positive = np.hstack(positive)
                negative = np.hstack(negative)
                writer.add_histogram(
                    'embedding/pairwise_distance/positive', positive,
                    global_step=epoch, bins='auto')
                writer.add_histogram(
                    'embedding/pairwise_distance/negative', negative,
                    global_step=epoch, bins='auto')

                _, _, _, eer = det_curve(
                    np.hstack([np.ones(len(positive)), np.zeros(len(negative))]),
                    np.hstack([positive, negative]), distances=True)
                writer.add_scalar('eer', eer, global_step=epoch)

                norms = np.hstack(norms)
                writer.add_histogram(
                    'embedding/norm', norms,
                    global_step=epoch, bins='auto')

            logging_callback.model = model
            logging_callback.optimizer = optimizer
            logging_callback.on_epoch_end(epoch)
<commit_msg>BREAKING: add support for margin in "sigmoid" triplet loss<commit_after>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2018 CNRS

# 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.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import itertools
import numpy as np
from tqdm import tqdm
import torch
from torch.autograd import Variable
import torch.nn.functional as F
from pyannote.audio.generators.speaker import SpeechSegmentGenerator
from pyannote.audio.callback import LoggingCallbackPytorch
from torch.optim import Adam
from pyannote.audio.embedding.utils import to_condensed
from scipy.spatial.distance import pdist, squareform
from pyannote.metrics.binary_classification import det_curve

class TripletLoss(object):
    """

    delta = d(anchor, positive) - d(anchor, negative)

    * with 'positive' clamping:
        loss = max(0, delta + margin x D)
    * with 'sigmoid' clamping:
        loss = sigmoid(10 * delta)

    where d(., .) varies in range [0, D] (e.g. D=2 for euclidean distance).

    Parameters
    ----------
    duration : float, optional
        Defautls to 3.2 seconds.
    metric : {'euclidean', 'cosine', 'angular'}, optional
        Defaults to 'cosine'.
    margin: float, optional
        Margin factor. Defaults to 0.2.
    clamp : {'positive', 'sigmoid', 'softmargin'}, optional
        Defaults to 'positive'.
    sampling : {'all', 'hard', 'negative'}, optional
        Triplet sampling strategy.
    per_label : int, optional
        Number of sequences per speaker in each batch. Defaults to 3.
    per_fold : int, optional
        If provided, sample triplets from groups of `per_fold` speakers at a
        time. Defaults to sample triplets from the whole speaker set.
    trim : float, optional
        Do not use speech segments that are that close to the beginning/end of
        the annotated region. Useful when features are short-term normalized,
        as this can result in weird feature distribution at the boundaries.
        Defaults to 0s (i.e. do no trim).
    parallel : int, optional
        Number of prefetching background generators. Defaults to 1.
        Each generator will prefetch enough batches to cover a whole epoch.
        Set `parallel` to 0 to not use background generators.
    """

    def __init__(self, duration=3.2,
                 metric='cosine', margin=0.2, clamp='positive',
                 sampling='all', per_label=3, per_fold=None,
                 trim=0, parallel=1):

        super(TripletLoss, self).__init__()

        self.metric = metric
        self.margin = margin

        if self.metric == 'cosine':
            self.margin_ = self.margin * 2.
        elif self.metric == 'angular':
            self.margin_ = self.margin * np.pi
        elif self.metric == 'euclidean':
            self.margin_ = self.margin * 2.
        else:
            msg = "'metric' must be one of {'euclidean', 'cosine', 'angular'}."
            raise ValueError(msg)

        if clamp not in {'positive', 'sigmoid', 'softmargin'}:
            msg = "'clamp' must be one of {'positive', 'sigmoid', 'softmargin'}."
            raise ValueError(msg)
        self.clamp = clamp

        if sampling not in {'all', 'hard', 'negative'}:
            msg = "'sampling' must be one of {'all', 'hard', 'negative'}."
            raise ValueError(msg)
        self.sampling = sampling

        self.per_fold = per_fold
        self.per_label = per_label
        self.duration = duration
        self.trim = trim
        self.parallel = parallel

    def pdist(self, fX):
        """Compute pdist à-la scipy.spatial.distance.pdist

        Parameters
        ----------
        fX : (n, d) torch.autograd.Variable
            Embeddings.

        Returns
        -------
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix
        """

        n_sequences, _ = fX.size()
        distances = []

        for i in range(n_sequences - 1):

            if self.metric in ('cosine', 'angular'):
                d = 1. - F.cosine_similarity(
                    fX[i, :].expand(n_sequences - 1 - i, -1),
                    fX[i+1:, :], dim=1, eps=1e-8)

                if self.metric == 'angular':
                    d = torch.acos(torch.clamp(1. - d, -1 + 1e-6, 1 - 1e-6))

            elif self.metric == 'euclidean':
                d = F.pairwise_distance(
                    fX[i, :].expand(n_sequences - 1 - i, -1),
                    fX[i+1:, :], p=2, eps=1e-06).view(-1)

            distances.append(d)

        return torch.cat(distances)

    def batch_hard(self, y, distances):
        """Build triplet with both hardest positive and hardest negative

        Parameters
        ----------
        y : list
            Sequence labels.
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix

        Returns
        -------
        anchors, positives, negatives : list of int
            Triplets indices.
        """

        anchors, positives, negatives = [], [], []

        if distances.is_cuda:
            distances = squareform(distances.data.cpu().numpy())
        else:
            distances = squareform(distances.data.numpy())
        y = np.array(y)

        for anchor, y_anchor in enumerate(y):

            d = distances[anchor]

            # hardest positive
            pos = np.where(y == y_anchor)[0]
            pos = [p for p in pos if p != anchor]
            positive = int(pos[np.argmax(d[pos])])

            # hardest negative
            neg = np.where(y != y_anchor)[0]
            negative = int(neg[np.argmin(d[neg])])

            anchors.append(anchor)
            positives.append(positive)
            negatives.append(negative)

        return anchors, positives, negatives

    def batch_negative(self, y, distances):
        """Build triplet with hardest negative

        Parameters
        ----------
        y : list
            Sequence labels.
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix

        Returns
        -------
        anchors, positives, negatives : list of int
            Triplets indices.
        """

        anchors, positives, negatives = [], [], []

        if distances.is_cuda:
            distances = squareform(distances.data.cpu().numpy())
        else:
            distances = squareform(distances.data.numpy())
        y = np.array(y)

        for anchor, y_anchor in enumerate(y):

            # hardest negative
            d = distances[anchor]
            neg = np.where(y != y_anchor)[0]
            negative = int(neg[np.argmin(d[neg])])

            for positive in np.where(y == y_anchor)[0]:
                if positive == anchor:
                    continue

                anchors.append(anchor)
                positives.append(positive)
                negatives.append(negative)

        return anchors, positives, negatives

    def batch_all(self, y, distances):
        """Build all possible triplet

        Parameters
        ----------
        y : list
            Sequence labels.
        distances : (n * (n-1) / 2,) torch.autograd.Variable
            Condensed pairwise distance matrix

        Returns
        -------
        anchors, positives, negatives : list of int
            Triplets indices.
        """

        anchors, positives, negatives = [], [], []

        for anchor, y_anchor in enumerate(y):
            for positive, y_positive in enumerate(y):

                # if same embedding or different labels, skip
                if (anchor == positive) or (y_anchor != y_positive):
                    continue

                for negative, y_negative in enumerate(y):

                    if y_negative == y_anchor:
                        continue

                    anchors.append(anchor)
                    positives.append(positive)
                    negatives.append(negative)

        return anchors, positives, negatives

    def triplet_loss(self, distances, anchors, positives, negatives,
                     return_delta=False):
        """Compute triplet loss

        Parameters
        ----------
        distances : torch.autograd.Variable
            Condensed matrix of pairwise distances.
        anchors, positives, negatives : list of int
            Triplets indices.
        return_delta : bool, optional
            Return delta before clamping.

        Returns
        -------
        loss : torch.autograd.Variable
            Triplet loss.
        """

        # estimate total number of embeddings from pdist shape
        n = int(.5 * (1 + np.sqrt(1 + 8 * len(distances))))
        n = [n] * len(anchors)

        # convert indices from squared matrix
        # to condensed matrix referential
        pos = list(map(to_condensed, n, anchors, positives))
        neg = list(map(to_condensed, n, anchors, negatives))

        # compute raw triplet loss (no margin, no clamping)
        # the lower, the better
        delta = distances[pos] - distances[neg]

        # clamp triplet loss
        if self.clamp == 'positive':
            loss = torch.clamp(delta + self.margin_, min=0)

        elif self.clamp == 'softmargin':
            loss = torch.log1p(torch.exp(delta))

        elif self.clamp == 'sigmoid':
            # TODO. tune this "10" hyperparameter
            loss = F.sigmoid(10 * (delta + self.margin_))

        # return triplet losses
        if return_delta:
            return loss, delta.view((-1, 1)), pos, neg
        else:
            return loss

    def fit(self, model, feature_extraction, protocol, log_dir, subset='train',
            epochs=1000, restart=0, gpu=False):

        import tensorboardX
        writer = tensorboardX.SummaryWriter(log_dir=log_dir)

        logging_callback = LoggingCallbackPytorch(log_dir=log_dir,
                                                  restart=restart > 0)

        batch_generator = SpeechSegmentGenerator(
            feature_extraction,
            per_label=self.per_label, per_fold=self.per_fold,
            duration=self.duration, trim=self.trim,
            parallel=self.parallel)
        batches = batch_generator(protocol, subset=subset)
        batch = next(batches)

        batches_per_epoch = batch_generator.batches_per_epoch

        if restart > 0:
            weights_pt = logging_callback.WEIGHTS_PT.format(
                log_dir=log_dir, epoch=restart)
            model.load_state_dict(torch.load(weights_pt))

        if gpu:
            model = model.cuda()

        model.internal = False

        optimizer = Adam(model.parameters())
        if restart > 0:
            optimizer_pt = logging_callback.OPTIMIZER_PT.format(
                log_dir=log_dir, epoch=restart)
            optimizer.load_state_dict(torch.load(optimizer_pt))
            if gpu:
                for state in optimizer.state.values():
                    for k, v in state.items():
                        if torch.is_tensor(v):
                            state[k] = v.cuda()

        epoch = restart if restart > 0 else -1
        while True:
            epoch += 1
            if epoch > epochs:
                break

            tloss_avg = 0.

            if epoch % 10 == 0:
                positive, negative = [], []
                norms = []

            desc = 'Epoch #{0}'.format(epoch)
            for i in tqdm(range(batches_per_epoch), desc=desc):

                model.zero_grad()

                batch = next(batches)

                X = Variable(torch.from_numpy(
                    np.array(np.rollaxis(batch['X'], 0, 2),
                             dtype=np.float32)))

                if gpu:
                    X = X.cuda()

                fX = model(X)

                if epoch % 10 == 0:
                    if gpu:
                        fX_ = fX.data.cpu().numpy()
                    else:
                        fX_ = fX.data.numpy()
                    norms.append(np.linalg.norm(fX_, axis=1))

                # pre-compute pairwise distances
                distances = self.pdist(fX)

                if epoch % 10 == 0:
                    if gpu:
                        distances_ = distances.data.cpu().numpy()
                    else:
                        distances_ = distances.data.numpy()
                    is_positive = pdist(batch['y'].reshape((-1, 1)), metric='chebyshev') < 1
                    positive.append(distances_[np.where(is_positive)])
                    negative.append(distances_[np.where(~is_positive)])

                # sample triplets
                triplets = getattr(self, 'batch_{0}'.format(self.sampling))
                anchors, positives, negatives = triplets(batch['y'], distances)

                # compute triplet loss
                losses = self.triplet_loss(
                    distances, anchors, positives, negatives)

                loss = torch.mean(losses)

                # log loss
                if gpu:
                    loss_ = float(loss.data.cpu().numpy())
                else:
                    loss_ = float(loss.data.numpy())
                tloss_avg += loss_

                loss.backward()
                optimizer.step()

            tloss_avg /= batches_per_epoch
            writer.add_scalar('tloss', tloss_avg, global_step=epoch)

            if epoch % 10 == 0:

                positive = np.hstack(positive)
                negative = np.hstack(negative)
                writer.add_histogram(
                    'embedding/pairwise_distance/positive', positive,
                    global_step=epoch, bins='auto')
                writer.add_histogram(
                    'embedding/pairwise_distance/negative', negative,
                    global_step=epoch, bins='auto')

                _, _, _, eer = det_curve(
                    np.hstack([np.ones(len(positive)), np.zeros(len(negative))]),
                    np.hstack([positive, negative]), distances=True)
                writer.add_scalar('eer', eer, global_step=epoch)

                norms = np.hstack(norms)
                writer.add_histogram(
                    'embedding/norm', norms,
                    global_step=epoch, bins='auto')

            logging_callback.model = model
            logging_callback.optimizer = optimizer
            logging_callback.on_epoch_end(epoch)
<|endoftext|>


<commit_before>from exceptions import ConfigError
from config_processing import loadConfigInformation
from copy import copy
from itertools import chain, product
from os.path import exists, abspath, join, split, relpath
import os
import sys
from glob import glob
from treedict import TreeDict
from cStringIO import StringIO
import re
import ctypes
import shutil
import cleaning

"""
Specifies the setup stuff
"""

# This list is to keep them from being unloaded and garbage collected.
loaded_ctype_dlls = []

# A setup regular expression for the

output_line_okay = re.compile(
    r"(^\s*running build_ext\s*$)|(^\s*skipping.+up-to-date\)\s*$)|(^\s*$)")

# Copied from python 2.7's subprocess.py file
from subprocess import Popen, PIPE, STDOUT

# Exception classes used by this module.
class CalledProcessError(Exception):
    """This exception is raised when a process run by check_call() or
    check_output() returns a non-zero exit status.
    The exit status will be stored in the returncode attribute;
    check_output() will also store the output in the output attribute.
    """
    def __init__(self, returncode, cmd, output=None):
        self.returncode = returncode
        self.cmd = cmd
        self.output = output
    def __str__(self):
        return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)

def check_output(*popenargs, **kwargs):
    r"""Run command with arguments and return its output as a byte string.

    If the exit code was non-zero it raises a CalledProcessError.  The
    CalledProcessError object will have the return code in the returncode
    attribute and output in the output attribute.

    The arguments are the same as for the Popen constructor.  Example:

    >>> check_output(["ls", "-l", "/dev/null"])
    'crw-rw-rw- 1 root root 1, 3 Oct 18  2007 /dev/null\n'

    The stdout argument is not allowed as it is used internally.
    To capture standard error in the result, use stderr=STDOUT.

    >>> check_output(["/bin/sh", "-c",
    ...               "ls -l non_existent_file ; exit 0"],
    ...              stderr=STDOUT)
    'ls: non_existent_file: No such file or directory\n'
    """
    if 'stdout' in kwargs:
        raise ValueError('stdout argument not allowed, it will be overridden.')
    process = Popen(stdout=PIPE, *popenargs, **kwargs)
    output, unused_err = process.communicate()
    retcode = process.poll()
    if retcode:
        cmd = kwargs.get("args")
        if cmd is None:
            cmd = popenargs[0]
        raise CalledProcessError(retcode, cmd, output=output)
    return output

def readyCMakeProjects(base_dir, opttree, config):
    """
    Compiles CMake Projects.
    """

    # may switch so it runs in a specified build directory.

    if opttree.verbose_mode:
        if opttree.no_compile:
            print "\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>"
            print "Loading cmake project library files.\n"
        else:
            print "\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>"
            print "Compiling cmake projects.\n"

    for k, b in config.cmake.iteritems(recursive=False, branch_mode = "only"):
        d = abspath(join(base_dir, b.directory))

        if not exists(d):
            raise ConfigError("CMake subproject '%s' directory does not exist (%s)"
                              % (k, d))

        def run(cmd):

            run_command = \
                "cd '%s' && CMAKE_INCLUDE_PATH=$INCLUDE_PATH CMAKE_LIBRARY_PATH=$LD_LIBRARY_PATH %s" % (d, cmd)

            try:
                check_output([run_command], shell=True, stderr=STDOUT)
            except CalledProcessError, ce:
                raise ConfigError("Error while compiling cmake project '%s' in '%s':\n%s\n%s"
                                  %(k, d, "Error code %d while running '%s':" % (ce.returncode, cmd), ce.output))

        if opttree.verbose_mode:
            print "CMake: '%s' in '%s' " % (k,d)

        first_attempt = True

        if not opttree.no_compile:
            while True:

                try:
                    if not exists(join(d, "Makefile")):
                        run("cmake ./")

                    run("make --jobs -f Makefile")

                except ConfigError, ce:

                    if first_attempt:
                        print ("WARNING: Error while compiling cmake project '%s';"
                               " removing cache files and retrying.") % k
                        cleaning.clean_cmake_project(opttree, config, b)
                        first_attempt = False
                        continue

                    else:
                        raise ce

                break

        if not exists(b.library_file):
            if opttree.no_compile:
                raise ConfigError("Expected shared library apparently not present, recompiling needed?: \n"
                                  + "  Subproject: %s " % k
                                  + "  Expected library file: %s" % (relpath(b.library_file)))
            else:
                raise ConfigError("Expected shared library apparently not produced by compiliation: \n"
                                  + "  Subproject: %s " % k
                                  + "  Expected library file: %s" % (relpath(b.library_file)))

        # Dynamically load them here; this means they will be preloaded and the program will work
        loaded_ctype_dlls.append(ctypes.cdll.LoadLibrary(b.library_file))

    if opttree.verbose_mode:
        print "Done compiling and loading cmake library projects."
        print "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n"

def runBuildExt(base_dir, opttree, config):
    """
    Sets up and runs the python build_ext setup configuration.
    """

    if opttree.no_compile:
        return

    ct = config.cython

    extra_include_dirs = ct.extra_include_dirs
    extra_library_dirs = ct.extra_library_dirs
    libraries          = ct.libraries
    library_map        = ct.library_map
    compiler_args      = ct.compiler_args
    link_args          = ct.link_args

    quiet = not opttree.verbose_mode

    from distutils.core import setup as dist_setup
    from distutils.extension import Extension

    if ct.numpy_needed:
	import numpy
	extra_include_dirs += [numpy.get_include()]

    ######################################################
    # First have to see if we're authorized to use cython files, or if we
    # should instead compile the included files

    # Get all the cython files in the sub directories and in this directory
    cython_files = config.cython_files

    # Set the compiler arguments -- Add in the environment path stuff
    ld_library_path = os.getenv("LD_LIBRARY_PATH")

    if ld_library_path is not None:
	extra_library_dirs += [p.strip() for p in ld_library_path.split(":") if len(p.strip()) > 0]

    include_path = os.getenv("INCLUDE_PATH")
    if include_path is not None:
	extra_include_dirs += [p.strip() for p in include_path.split(":") if len(p.strip()) > 0]

    # The rest is also shared with the setup.py file, in addition to
    # this one, so

    def strip_empty(l):
        return [e.strip() for e in l if len(e.strip()) != 0]

    def get_include_dirs(m):
        return strip_empty(extra_include_dirs)

    def get_libnode_library(t):
        return t.library_name

    def get_libnode_directory(t):
        return t.directory

    def get_library_dirs(m):
        l = strip_empty(extra_library_dirs)

        if m in library_map:
            for lib in library_map[m]:
                if type(lib) is TreeDict:
                    l.append(get_libnode_directory(lib))

        l = [abspath(ld) for ld in l]
        return l

    def get_libraries(m):
        def process_lib(lib):
            if type(lib) is TreeDict:
                return get_libnode_library(lib)
            else:
                return lib

        liblist = libraries + (library_map[m] if m in library_map else [])

        return strip_empty(process_lib(lib) for lib in liblist)

    def get_extra_compile_args(m):
        return strip_empty(compiler_args + (['-g'] if config.debug_mode else ["-DNDEBUG"]))

    def get_extra_link_args(m):
        return strip_empty(link_args + (['-g'] if config.debug_mode else ["-DNDEBUG"]))

    ############################################################
    # Cython extension lists
    ext_modules = []

    for f in cython_files:
        # extract the module names
        rel_f = relpath(f, base_dir)
        assert rel_f.endswith('.pyx')
        modname = rel_f[:-4].replace('/', '.')

        ext_modules.append(Extension(
            modname,
            [f],
            include_dirs = get_include_dirs(modname),
            library_dirs = get_library_dirs(modname),
            libraries = get_libraries(modname),
            extra_compile_args = get_extra_compile_args(modname),
            extra_link_args = get_extra_link_args(modname),
            language = "c++" if ct.use_cpp else "c",
            ))

    ############################################################
    # Now get all these ready to go

    from Cython.Distutils import build_ext

    cmdclass = {'build_ext' : build_ext}

    old_argv = copy(sys.argv)
    sys.argv = (old_argv[0], "build_ext", "--inplace")

    if not quiet:
        print ">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> "
        print "Compiling cython extension modules.\n"

    old_stdout = sys.stdout
    old_stderr = sys.stderr

    try:
        if quiet:
            output = sys.stderr = sys.stdout = StringIO()

        dist_setup(
            cmdclass = cmdclass,
            ext_modules = ext_modules)

    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr

        if quiet:
            output_string = output.getvalue()

            # Check for output
            if not all(output_line_okay.match(le) is not None for le in output_string.split('\n')):
                if quiet:
                    print "++++++++++++++++++++"
                    print "Compiling cython extension modules.\n"

                print output_string

    if not quiet:
        print "\nCython extension modules successfully compiled."
        print "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n"

    sys.argv = old_argv

def setup(base_dir, opttree, config):
    "The main setup function; calls the rest."
    readyCMakeProjects(base_dir, opttree, config)
    runBuildExt(base_dir, opttree, config)

<commit_msg>Added cmake project recompilation when libraries fail to load properly.<commit_after>from exceptions import ConfigError
from config_processing import loadConfigInformation
from copy import copy
from itertools import chain, product
from os.path import exists, abspath, join, split, relpath
import os
import sys
from glob import glob
from treedict import TreeDict
from cStringIO import StringIO
import re
import ctypes
import shutil
import cleaning

"""
Specifies the setup stuff
"""

# This list is to keep them from being unloaded and garbage collected.
loaded_ctype_dlls = []

# A setup regular expression for the

output_line_okay = re.compile(
    r"(^\s*running build_ext\s*$)|(^\s*skipping.+up-to-date\)\s*$)|(^\s*$)")

# Copied from python 2.7's subprocess.py file
from subprocess import Popen, PIPE, STDOUT

# Exception classes used by this module.
class CalledProcessError(Exception):
    """This exception is raised when a process run by check_call() or
    check_output() returns a non-zero exit status.
    The exit status will be stored in the returncode attribute;
    check_output() will also store the output in the output attribute.
    """
    def __init__(self, returncode, cmd, output=None):
        self.returncode = returncode
        self.cmd = cmd
        self.output = output
    def __str__(self):
        return "Command '%s' returned non-zero exit status %d" % (self.cmd, self.returncode)

def check_output(*popenargs, **kwargs):
    r"""Run command with arguments and return its output as a byte string.

    If the exit code was non-zero it raises a CalledProcessError.  The
    CalledProcessError object will have the return code in the returncode
    attribute and output in the output attribute.

    The arguments are the same as for the Popen constructor.  Example:

    >>> check_output(["ls", "-l", "/dev/null"])
    'crw-rw-rw- 1 root root 1, 3 Oct 18  2007 /dev/null\n'

    The stdout argument is not allowed as it is used internally.
    To capture standard error in the result, use stderr=STDOUT.

    >>> check_output(["/bin/sh", "-c",
    ...               "ls -l non_existent_file ; exit 0"],
    ...              stderr=STDOUT)
    'ls: non_existent_file: No such file or directory\n'
    """
    if 'stdout' in kwargs:
        raise ValueError('stdout argument not allowed, it will be overridden.')
    process = Popen(stdout=PIPE, *popenargs, **kwargs)
    output, unused_err = process.communicate()
    retcode = process.poll()
    if retcode:
        cmd = kwargs.get("args")
        if cmd is None:
            cmd = popenargs[0]
        raise CalledProcessError(retcode, cmd, output=output)
    return output

def readyCMakeProjects(base_dir, opttree, config):
    """
    Compiles CMake Projects.
    """

    # may switch so it runs in a specified build directory.

    if opttree.verbose_mode:
        if opttree.no_compile:
            print "\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>"
            print "Loading cmake project library files.\n"
        else:
            print "\n>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>"
            print "Compiling cmake projects.\n"

    for k, b in config.cmake.iteritems(recursive=False, branch_mode = "only"):
        d = abspath(join(base_dir, b.directory))

        if not exists(d):
            raise ConfigError("CMake subproject '%s' directory does not exist (%s)"
                              % (k, d))

        def run(cmd):

            run_command = \
                "cd '%s' && CMAKE_INCLUDE_PATH=$INCLUDE_PATH CMAKE_LIBRARY_PATH=$LD_LIBRARY_PATH %s" % (d, cmd)

            try:
                check_output([run_command], shell=True, stderr=STDOUT)
            except CalledProcessError, ce:
                raise ConfigError("Error while compiling cmake project '%s' in '%s':\n%s\n%s"
                                  %(k, d, "Error code %d while running '%s':" % (ce.returncode, cmd), ce.output))

        if opttree.verbose_mode:
            print "CMake: '%s' in '%s' " % (k,d)

        retry_allowed = True
        load_retry_allowed = True

        while True:

            if not opttree.no_compile:
                while True:

                    try:
                        if not exists(join(d, "Makefile")):
                            run("cmake ./")

                        run("make --jobs -f Makefile")

                    except ConfigError, ce:

                        if retry_allowed:
                            print ("WARNING: Error while compiling cmake project '%s';"
                                   " removing cache files and retrying.") % k
                            cleaning.clean_cmake_project(opttree, config, b)
                            retry_allowed = False
                            continue

                        else:
                            raise

                    break

            if not exists(b.library_file):
                if opttree.no_compile:
                    raise ConfigError("Expected shared library apparently not present, recompiling needed?: \n"
                                      + "  Subproject: %s " % k
                                      + "  Expected library file: %s" % (relpath(b.library_file)))
                else:
                    raise ConfigError("Expected shared library apparently not produced by compiliation: \n"
                                      + "  Subproject: %s " % k
                                      + "  Expected library file: %s" % (relpath(b.library_file)))

            # Dynamically load them here; this means they will be preloaded and the program will work
            try:
                loaded_dll = ctypes.cdll.LoadLibrary(b.library_file)
            except OSError, ose:

                if load_retry_allowed:
                    print "Error loading library: ", str(ose)
                    print "Cleaning, attempting again."
                    cleaning.clean_cmake_project(opttree, config, b)
                    load_retry_allowed = False
                    continue
                else:
                    raise

            break

        loaded_ctype_dlls.append(loaded_ctype_dlls)

    if opttree.verbose_mode:
        print "Done compiling and loading cmake library projects."
        print "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n"

def runBuildExt(base_dir, opttree, config):
    """
    Sets up and runs the python build_ext setup configuration.
    """

    if opttree.no_compile:
        return

    ct = config.cython

    extra_include_dirs = ct.extra_include_dirs
    extra_library_dirs = ct.extra_library_dirs
    libraries          = ct.libraries
    library_map        = ct.library_map
    compiler_args      = ct.compiler_args
    link_args          = ct.link_args

    quiet = not opttree.verbose_mode

    from distutils.core import setup as dist_setup
    from distutils.extension import Extension

    if ct.numpy_needed:
	import numpy
	extra_include_dirs += [numpy.get_include()]

    ######################################################
    # First have to see if we're authorized to use cython files, or if we
    # should instead compile the included files

    # Get all the cython files in the sub directories and in this directory
    cython_files = config.cython_files

    # Set the compiler arguments -- Add in the environment path stuff
    ld_library_path = os.getenv("LD_LIBRARY_PATH")

    if ld_library_path is not None:
	extra_library_dirs += [p.strip() for p in ld_library_path.split(":") if len(p.strip()) > 0]

    include_path = os.getenv("INCLUDE_PATH")
    if include_path is not None:
	extra_include_dirs += [p.strip() for p in include_path.split(":") if len(p.strip()) > 0]

    # The rest is also shared with the setup.py file, in addition to
    # this one, so

    def strip_empty(l):
        return [e.strip() for e in l if len(e.strip()) != 0]

    def get_include_dirs(m):
        return strip_empty(extra_include_dirs)

    def get_libnode_library(t):
        return t.library_name

    def get_libnode_directory(t):
        return t.directory

    def get_library_dirs(m):
        l = strip_empty(extra_library_dirs)

        if m in library_map:
            for lib in library_map[m]:
                if type(lib) is TreeDict:
                    l.append(get_libnode_directory(lib))

        l = [abspath(ld) for ld in l]
        return l

    def get_libraries(m):
        def process_lib(lib):
            if type(lib) is TreeDict:
                return get_libnode_library(lib)
            else:
                return lib

        liblist = libraries + (library_map[m] if m in library_map else [])

        return strip_empty(process_lib(lib) for lib in liblist)

    def get_extra_compile_args(m):
        return strip_empty(compiler_args + (['-g'] if config.debug_mode else ["-DNDEBUG"]))

    def get_extra_link_args(m):
        return strip_empty(link_args + (['-g'] if config.debug_mode else ["-DNDEBUG"]))

    ############################################################
    # Cython extension lists
    ext_modules = []

    for f in cython_files:
        # extract the module names
        rel_f = relpath(f, base_dir)
        assert rel_f.endswith('.pyx')
        modname = rel_f[:-4].replace('/', '.')

        ext_modules.append(Extension(
            modname,
            [f],
            include_dirs = get_include_dirs(modname),
            library_dirs = get_library_dirs(modname),
            libraries = get_libraries(modname),
            extra_compile_args = get_extra_compile_args(modname),
            extra_link_args = get_extra_link_args(modname),
            language = "c++" if ct.use_cpp else "c",
            ))

    ############################################################
    # Now get all these ready to go

    from Cython.Distutils import build_ext

    cmdclass = {'build_ext' : build_ext}

    old_argv = copy(sys.argv)
    sys.argv = (old_argv[0], "build_ext", "--inplace")

    if not quiet:
        print ">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> "
        print "Compiling cython extension modules.\n"

    old_stdout = sys.stdout
    old_stderr = sys.stderr

    try:
        if quiet:
            output = sys.stderr = sys.stdout = StringIO()

        dist_setup(
            cmdclass = cmdclass,
            ext_modules = ext_modules)

    finally:
        sys.stdout = old_stdout
        sys.stderr = old_stderr

        if quiet:
            output_string = output.getvalue()

            # Check for output
            if not all(output_line_okay.match(le) is not None for le in output_string.split('\n')):
                if quiet:
                    print "++++++++++++++++++++"
                    print "Compiling cython extension modules.\n"

                print output_string

    if not quiet:
        print "\nCython extension modules successfully compiled."
        print "<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<\n"

    sys.argv = old_argv

def setup(base_dir, opttree, config):
    "The main setup function; calls the rest."
    readyCMakeProjects(base_dir, opttree, config)
    runBuildExt(base_dir, opttree, config)

<|endoftext|>


<commit_before><commit_msg>fixes for MC plotmon 2D features<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from optparse import OptionParser
import logging
import os
import sys
import zeit.care.crawl
import zeit.connector.connector
from zeit.connector.resource import Resource

logger = logging.getLogger(__name__)

def commentthread_worker(resource, connector):
    pass

def main():
    usage = "usage: %prog [options] arg"
    parser = OptionParser(usage)
    parser.add_option("-c", "--collection", dest="collection",
                      help="entry collection for starting the conversion")
    parser.add_option("-w", "--webdav", dest="webdav",
                      help="webdav server uri")
    parser.add_option("-l", "--log", dest="logfile",
                      help="logfile for errors")

    (options, args) = parser.parse_args()

    if not options.collection:
        parser.error("missing entry point")

    if not options.webdav:
        parser.error("missing webdav uri")

    if options.logfile:
        add_file_logging(logger, options.logfile)

    connector = zeit.connector.connector.Connector(roots=dict(
        default=options.webdav))
    crawler = zeit.care.crawl.Crawler(connector, commentthread_worker)
    crawler.run(options.collection)

<commit_msg>Collect comments in commentthread worker.<commit_after># -*- coding: utf-8 -*-
from optparse import OptionParser
import logging
import os
import sys
import zeit.care.crawl
import zeit.connector.connector
from zeit.connector.resource import Resource

logger = logging.getLogger(__name__)

def commentthread_worker(resource, connector):
    properties = resource.properties
    comments = properties.get(
        ('comments', 'http://namespaces.zeit.de/CMS/document'))

def main():
    usage = "usage: %prog [options] arg"
    parser = OptionParser(usage)
    parser.add_option("-c", "--collection", dest="collection",
                      help="entry collection for starting the conversion")
    parser.add_option("-w", "--webdav", dest="webdav",
                      help="webdav server uri")
    parser.add_option("-l", "--log", dest="logfile",
                      help="logfile for errors")

    (options, args) = parser.parse_args()

    if not options.collection:
        parser.error("missing entry point")

    if not options.webdav:
        parser.error("missing webdav uri")

    if options.logfile:
        add_file_logging(logger, options.logfile)

    connector = zeit.connector.connector.Connector(roots=dict(
        default=options.webdav))
    crawler = zeit.care.crawl.Crawler(connector, commentthread_worker)
    crawler.run(options.collection)

<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
============================================
Copy a tree chain while overwriting branches
============================================

This is an example showing how to copy a tree chain while overwriting one or
more of its branches with new values.
"""
print __doc__
import rootpy
rootpy.log.basic_config_colorized()
from rootpy.tree import Tree, TreeModel, TreeChain, FloatCol, IntCol
from rootpy.io import root_open
from random import gauss

"""
This first section of code only creates an example tree chain.
"""

class Event(TreeModel):
    """Event model definition"""
    x = FloatCol()
    y = FloatCol()
    z = FloatCol()
    i = IntCol()

# first create several example trees in separate files
fnames = ["test_{0:d}.root".format(i) for i in xrange(10)]

for fname in fnames:
    with root_open(fname, "recreate") as f:

        tree = Tree("test", model=Event)

        # fill the tree
        for i in xrange(10000):
            tree.x = gauss(.5, 1.)
            tree.y = gauss(.3, 2.)
            tree.z = gauss(13., 42.)
            tree.i = i
            tree.fill()
        tree.write()

"""
This section below takes the example trees and copies it while overwriting a
branch with new values.
"""

# first define the chain of trees
chain = TreeChain(name="test", files=fnames)

# Now we want to copy the tree above into a new file while overwriting a branch
# First create a new file to save the new tree in:
f_copy = root_open("test_copy.root", "recreate")

# You may not know the entire model of the original tree but only the branches
# you intend to overwrite, so I am not specifying the model=Event below as an
# example of how to deal with this in general:
tree_copy = Tree("test_copy")

# If the original tree was not handed to you through rootpy don't forget to:
# >>> from rootpy import asrootpy
# >>> tree = asrootpy(tree)

# Here we specify the buffer for the new tree to use. We use the same buffer as
# the original tree. This creates all the same branches in the new tree but
# their addresses point to the same memory used by the original tree.
tree_copy.set_buffer(
        chain._buffer,
        create_branches=True)

# Now loop over the original tree and fill the new tree
for entry in chain:
    # Overwrite a branch value. This changes the value that will be written to
    # the new tree but leaves the value unchanged in the original tree on disk.
    entry.x = 3.141
    # "entry" is actually the buffer, which is shared between both trees.
    tree_copy.Fill()

# tree_copy is now a copy of tree where the "x" branch has been overwritten
# with new values
tree_copy.Write()
f_copy.Close()
<commit_msg>examples: speedup chain_overwrite.py by using fewer files and entries in each<commit_after>#!/usr/bin/env python
"""
============================================
Copy a tree chain while overwriting branches
============================================

This is an example showing how to copy a tree chain while overwriting one or
more of its branches with new values.
"""
print __doc__
import rootpy
rootpy.log.basic_config_colorized()
from rootpy.tree import Tree, TreeModel, TreeChain, FloatCol, IntCol
from rootpy.io import root_open
from random import gauss

"""
This first section of code only creates an example tree chain.
"""

class Event(TreeModel):
    """Event model definition"""
    x = FloatCol()
    y = FloatCol()
    z = FloatCol()
    i = IntCol()

# first create several example trees in separate files
fnames = ["test_{0:d}.root".format(i) for i in xrange(5)]

for fname in fnames:
    with root_open(fname, "recreate") as f:

        tree = Tree("test", model=Event)

        # fill the tree
        for i in xrange(100):
            tree.x = gauss(.5, 1.)
            tree.y = gauss(.3, 2.)
            tree.z = gauss(13., 42.)
            tree.i = i
            tree.fill()
        tree.write()

"""
This section below takes the example trees and copies it while overwriting a
branch with new values.
"""

# first define the chain of trees
chain = TreeChain(name="test", files=fnames)

# Now we want to copy the tree above into a new file while overwriting a branch
# First create a new file to save the new tree in:
f_copy = root_open("test_copy.root", "recreate")

# You may not know the entire model of the original tree but only the branches
# you intend to overwrite, so I am not specifying the model=Event below as an
# example of how to deal with this in general:
tree_copy = Tree("test_copy")

# If the original tree was not handed to you through rootpy don't forget to:
# >>> from rootpy import asrootpy
# >>> tree = asrootpy(tree)

# Here we specify the buffer for the new tree to use. We use the same buffer as
# the original tree. This creates all the same branches in the new tree but
# their addresses point to the same memory used by the original tree.
tree_copy.set_buffer(
        chain._buffer,
        create_branches=True)

# Now loop over the original tree and fill the new tree
for entry in chain:
    # Overwrite a branch value. This changes the value that will be written to
    # the new tree but leaves the value unchanged in the original tree on disk.
    entry.x = 3.141
    # "entry" is actually the buffer, which is shared between both trees.
    tree_copy.Fill()

# tree_copy is now a copy of tree where the "x" branch has been overwritten
# with new values
tree_copy.Write()
f_copy.Close()
<|endoftext|>