<commit_before># -*- coding: utf-8 -*-
"""
Created on Mon Jul 11 10:00:25 2016

@author: Charlie Shobe

for West Valley Project: Landlab component to calculate drainage density

"""
from landlab import Component
import numpy as np

class DrainageDensity(Component):

    """Calculate drainage density over a DEM.

    calc_drainage_density function returns drainage density for the model
    domain.

    Landlab component that implements the distance to channel algorithm of
    Tucker et al., 2001.

    Written by C. Shobe on 7/11/2016

    Construction::

        DrainageDensity(grid, channel_network_name='string')

    Parameters
    ----------
    grid : ModelGrid
    channel_network_name : string naming a boolean grid field with 'True'
        in channels and 'False' elsewhere.

    Examples
    ---------
    >>> import numpy as np
    >>> from landlab import RasterModelGrid
    >>> from landlab.components.flow_routing import FlowRouter
    >>> from landlab.components import FastscapeEroder

    >>> mg = RasterModelGrid((10, 10), 1.0)
    >>> _ = mg.add_zeros('node', 'topographic__elevation')

    >>> np.random.seed(50)
    >>> noise = np.random.rand(100)
    >>> mg.at_node['topographic__elevation'] += noise
    >>> mg.at_node['topographic__elevation'] # doctest: +NORMALIZE_WHITESPACE
    array([ 0.49460165,  0.2280831 ,  0.25547392,  0.39632991,  0.3773151 ,
        0.99657423,  0.4081972 ,  0.77189399,  0.76053669,  0.31000935,
        0.3465412 ,  0.35176482,  0.14546686,  0.97266468,  0.90917844,
        0.5599571 ,  0.31359075,  0.88820004,  0.67457307,  0.39108745,
        0.50718412,  0.5241035 ,  0.92800093,  0.57137307,  0.66833757,
        0.05225869,  0.3270573 ,  0.05640164,  0.17982769,  0.92593317,
        0.93801522,  0.71409271,  0.73268761,  0.46174768,  0.93132927,
        0.40642024,  0.68320577,  0.64991587,  0.59876518,  0.22203939,
        0.68235717,  0.8780563 ,  0.79671726,  0.43200225,  0.91787822,
        0.78183368,  0.72575028,  0.12485469,  0.91630845,  0.38771099,
        0.29492955,  0.61673141,  0.46784623,  0.25533891,  0.83899589,
        0.1786192 ,  0.22711417,  0.65987645,  0.47911625,  0.07344734,
        0.13896007,  0.11230718,  0.47778497,  0.54029623,  0.95807105,
        0.58379231,  0.52666409,  0.92226269,  0.91925702,  0.25200886,
        0.68263261,  0.96427612,  0.22696165,  0.7160172 ,  0.79776011,
        0.9367512 ,  0.8537225 ,  0.42154581,  0.00543987,  0.03486533,
        0.01390537,  0.58890993,  0.3829931 ,  0.11481895,  0.86445401,
        0.82165703,  0.73749168,  0.84034417,  0.4015291 ,  0.74862   ,
        0.55962945,  0.61323757,  0.29810165,  0.60237917,  0.42567684,
        0.53854438,  0.48672986,  0.49989164,  0.91745948,  0.26287702])

    >>> fr = FlowRouter(mg)
    >>> fsc = FastscapeEroder(mg, K_sp=.01, m_sp=.5, n_sp=1)
    >>> for x in range(100):
    ...     fr.run_one_step()
    ...     fsc.run_one_step(dt = 10.0)
    ...     mg.at_node['topographic__elevation'][mg.core_nodes] += .01

    >>> channels = mg.at_node['drainage_area'] > 5
    >>> _ = mg.add_field('channel_network', channels, at='node')

    >>> dd = DrainageDensity(mg, channel_network_name='channel_network')
    >>> mean_drainage_density = dd.calc_drainage_density()
    >>> np.isclose(mean_drainage_density, 0.3831100571)
    True
    """

    _name = 'DrainageDensity'

    _input_var_names = (
        'flow__receiver_node',
        'flow__link_to_receiver_node',
        'channel_network',
    )

    _output_var_names = (
        'distance_to_channel',
    )

    _var_units = {
        'flow__receiver_node': '-',
        'flow__link_to_receiver_node': '-',
        'channel_network': '-',
        'distance_to_channel': 'm',
    }

    _var_mapping = {
        'flow__receiver_node': 'node',
        'flow__link_to_receiver_node': 'node',
        'channel_network': 'node',
        'distance_to_channel': 'node',
    }

    _var_doc = {
        'flow__receiver_node':
            'Node array of receivers (node that receives flow from current '
            'node)',
        'flow__link_to_receiver_node':
            'ID of link downstream of each node, which carries the discharge',
        'channel_network':
            'Logical map of at which grid nodes channels are present',
        'distance_to_channel':
            'Distance from each node to the nearest channel',
    }

    def __init__(self, grid, channel_network_name='channel_network', **kwds):
        """Initialize the DrainageDensity component.

        Parameters
        ----------
        grid : ModelGrid
            Landlab ModelGrid object
        channel_network_name : string, optional (defaults to 'channel_network')
            Name of Landlab grid field that holds 1's where channels exist
            and 0's elsewhere
        """

        # Store grid
        self._grid = grid

        # Create fields... Input fields should raise an error if not present,
        # rather than silently creating a nonsensical blank field.
        # Channel network
        self.channel_network_name = channel_network_name
        if channel_network_name in grid.at_node:
            self.channel_network = grid.at_node[channel_network_name].view(
                dtype=np.uint8)
            # for this component to work with Cython acceleration,
            # the channel_network must be int, not bool...
        else:
            raise FieldError(
                'DrainageDensity needs the field channel_network!')

        # Flow receivers
        if 'flow__receiver_node' in grid.at_node:
            self.flow_receivers = grid.at_node['flow__receiver_node']
        else:
            raise FieldError(
                'DrainageDensity needs the field flow__receiver_node!')

        # Links to receiver nodes
        if 'flow__link_to_receiver_node' in grid.at_node:
            # ^flow__link_to_receiver_node
            self.stack_links = grid.at_node['flow__link_to_receiver_node']
        else:
            raise FieldError(
                'DrainageDensity needs the field flow__link_to_receiver_node!')

        #   Distance to channel
        if 'distance_to_channel' in grid.at_node:
            self.distance_to_channel = grid.at_node['distance_to_channel']
        else:
            self.distance_to_channel = grid.add_zeros(
                'node', 'distance_to_channel', dtype=float)

    def calc_drainage_density(self, **kwds):
        # ^there is no 'run_one_step' methid b/c this is a tool, not a model.
        """Calculate distance to channel and drainage density, after
        Tucker et al., 2001.

        Returns
        -------
        landscape_drainage_density : float (1/m)
            Drainage density over the model domain.
        """
        from .cfuncs import _calc_dists_to_channel
        _calc_dists_to_channel(self.channel_network,
                               self.flow_receivers,
                               self.grid._length_of_link_with_diagonals,
                               self.stack_links,
                               self.distance_to_channel,
                               self.grid.number_of_nodes)
        landscape_drainage_density = 1. / (2.0 * np.mean(self.grid.at_node[
            'distance_to_channel'][self.grid.core_nodes]))
        # self.distance_to_channel))  # this is THE drainage density
        return landscape_drainage_density
<commit_msg>Moved field checks to top of __init__.<commit_after># -*- coding: utf-8 -*-
"""
Created on Mon Jul 11 10:00:25 2016

@author: Charlie Shobe

for West Valley Project: Landlab component to calculate drainage density

"""
from landlab import Component
import numpy as np

class DrainageDensity(Component):

    """Calculate drainage density over a DEM.

    calc_drainage_density function returns drainage density for the model
    domain.

    Landlab component that implements the distance to channel algorithm of
    Tucker et al., 2001.

    Written by C. Shobe on 7/11/2016

    Construction::

        DrainageDensity(grid, channel_network_name='string')

    Parameters
    ----------
    grid : ModelGrid
    channel_network_name : string naming a boolean grid field with 'True'
        in channels and 'False' elsewhere.

    Examples
    ---------
    >>> import numpy as np
    >>> from landlab import RasterModelGrid
    >>> from landlab.components.flow_routing import FlowRouter
    >>> from landlab.components import FastscapeEroder

    >>> mg = RasterModelGrid((10, 10), 1.0)
    >>> _ = mg.add_zeros('node', 'topographic__elevation')

    >>> np.random.seed(50)
    >>> noise = np.random.rand(100)
    >>> mg.at_node['topographic__elevation'] += noise
    >>> mg.at_node['topographic__elevation'] # doctest: +NORMALIZE_WHITESPACE
    array([ 0.49460165,  0.2280831 ,  0.25547392,  0.39632991,  0.3773151 ,
        0.99657423,  0.4081972 ,  0.77189399,  0.76053669,  0.31000935,
        0.3465412 ,  0.35176482,  0.14546686,  0.97266468,  0.90917844,
        0.5599571 ,  0.31359075,  0.88820004,  0.67457307,  0.39108745,
        0.50718412,  0.5241035 ,  0.92800093,  0.57137307,  0.66833757,
        0.05225869,  0.3270573 ,  0.05640164,  0.17982769,  0.92593317,
        0.93801522,  0.71409271,  0.73268761,  0.46174768,  0.93132927,
        0.40642024,  0.68320577,  0.64991587,  0.59876518,  0.22203939,
        0.68235717,  0.8780563 ,  0.79671726,  0.43200225,  0.91787822,
        0.78183368,  0.72575028,  0.12485469,  0.91630845,  0.38771099,
        0.29492955,  0.61673141,  0.46784623,  0.25533891,  0.83899589,
        0.1786192 ,  0.22711417,  0.65987645,  0.47911625,  0.07344734,
        0.13896007,  0.11230718,  0.47778497,  0.54029623,  0.95807105,
        0.58379231,  0.52666409,  0.92226269,  0.91925702,  0.25200886,
        0.68263261,  0.96427612,  0.22696165,  0.7160172 ,  0.79776011,
        0.9367512 ,  0.8537225 ,  0.42154581,  0.00543987,  0.03486533,
        0.01390537,  0.58890993,  0.3829931 ,  0.11481895,  0.86445401,
        0.82165703,  0.73749168,  0.84034417,  0.4015291 ,  0.74862   ,
        0.55962945,  0.61323757,  0.29810165,  0.60237917,  0.42567684,
        0.53854438,  0.48672986,  0.49989164,  0.91745948,  0.26287702])

    >>> fr = FlowRouter(mg)
    >>> fsc = FastscapeEroder(mg, K_sp=.01, m_sp=.5, n_sp=1)
    >>> for x in range(100):
    ...     fr.run_one_step()
    ...     fsc.run_one_step(dt = 10.0)
    ...     mg.at_node['topographic__elevation'][mg.core_nodes] += .01

    >>> channels = mg.at_node['drainage_area'] > 5
    >>> _ = mg.add_field('channel_network', channels, at='node')

    >>> dd = DrainageDensity(mg, channel_network_name='channel_network')
    >>> mean_drainage_density = dd.calc_drainage_density()
    >>> np.isclose(mean_drainage_density, 0.3831100571)
    True
    """

    _name = 'DrainageDensity'

    _input_var_names = (
        'flow__receiver_node',
        'flow__link_to_receiver_node',
        'channel_network',
    )

    _output_var_names = (
        'distance_to_channel',
    )

    _var_units = {
        'flow__receiver_node': '-',
        'flow__link_to_receiver_node': '-',
        'channel_network': '-',
        'distance_to_channel': 'm',
    }

    _var_mapping = {
        'flow__receiver_node': 'node',
        'flow__link_to_receiver_node': 'node',
        'channel_network': 'node',
        'distance_to_channel': 'node',
    }

    _var_doc = {
        'flow__receiver_node':
            'Node array of receivers (node that receives flow from current '
            'node)',
        'flow__link_to_receiver_node':
            'ID of link downstream of each node, which carries the discharge',
        'channel_network':
            'Logical map of at which grid nodes channels are present',
        'distance_to_channel':
            'Distance from each node to the nearest channel',
    }

    def __init__(self, grid, channel_network_name='channel_network', **kwds):
        """Initialize the DrainageDensity component.

        Parameters
        ----------
        grid : ModelGrid
            Landlab ModelGrid object
        channel_network_name : string, optional (defaults to 'channel_network')
            Name of Landlab grid field that holds 1's where channels exist
            and 0's elsewhere
        """
        required = ('flow__receiver_node', 'flow__link_to_receiver_node',
                    channel_network_name)
        for name in required:
            if name not in grid.at_node:
                raise FieldError(
                    '{name]: missing required field'.format(name=name))

        # Store grid
        self._grid = grid

        # Create fields... Input fields should raise an error if not present,
        # rather than silently creating a nonsensical blank field.
        # Channel network
        self.channel_network_name = channel_network_name

        # for this component to work with Cython acceleration,
        # the channel_network must be uint8, not bool...
        self.channel_network = grid.at_node[channel_network_name].view(
            dtype=np.uint8)

        # Flow receivers
        self.flow_receivers = grid.at_node['flow__receiver_node']

        # Links to receiver nodes
        self.stack_links = grid.at_node['flow__link_to_receiver_node']

        # Distance to channel
        try:
            self.distance_to_channel = grid.at_node['distance_to_channel']
        except KeyError:
            self.distance_to_channel = grid.add_zeros(
                'distance_to_channel', at='node', dtype=float)

    def calc_drainage_density(self):
        """Calculate drainage density.

        Returns
        -------
        float
            The drainage density.
        """

        # ^there is no 'run_one_step' methid b/c this is a tool, not a model.
        """Calculate distance to channel and drainage density, after
        Tucker et al., 2001.

        Returns
        -------
        landscape_drainage_density : float (1/m)
            Drainage density over the model domain.
        """
        from .cfuncs import _calc_dists_to_channel
        _calc_dists_to_channel(self.channel_network,
                               self.flow_receivers,
                               self.grid._length_of_link_with_diagonals,
                               self.stack_links,
                               self.distance_to_channel,
                               self.grid.number_of_nodes)
        landscape_drainage_density = 1. / (2.0 * np.mean(self.grid.at_node[
            'distance_to_channel'][self.grid.core_nodes]))
        # self.distance_to_channel))  # this is THE drainage density
        return landscape_drainage_density
<|endoftext|>


<commit_before>#!/usr/bin/python

DOCUMENTATION = '''
module: route53_zone
short_description: add or delete Route53 zones
description:
    - Creates and deletes Route53 private and public zones
options:
    zone:
        description:
            - The DNS zone record (eg: foo.com.)
        required: true
        default: null
    command:
        description:
            - whether or not the zone should exist or not
        required: false
        default: true
    vpc_id:
        description:
            - The VPC ID the zone should be a part of (if this is going to be a private zone)
        required: false
        default: null
    vpc_region:
        description:
            - The VPC Region the zone should be a part of (if this is going to be a private zone)
        required: false
        default: null
    comment:
        description:
            - Comment associated with the zone
        required: false
        default: ''
extends_documentation_fragment: aws
'''

import time

try:
    import boto
    import boto.ec2
    from boto import route53
    from boto.route53 import Route53Connection
    from boto.route53.zone import Zone
    HAS_BOTO = True
except ImportError:
    HAS_BOTO = False

def main():
    module = AnsibleModule(
        argument_spec=dict(
            zone=dict(required=True),
            command=dict(default='create', choices=['create', 'delete']),
            vpc_id=dict(default=None),
            vpc_region=dict(default=None),
            comment=dict(default=''),
        )
    )

    if not HAS_BOTO:
        module.fail_json(msg='boto required for this module')

    zone_in = module.params.get('zone').lower()
    command = module.params.get('command').lower()
    vpc_id = module.params.get('vpc_id')
    vpc_region = module.params.get('vpc_region')
    comment = module.params.get('comment')

    private_zone = vpc_id is not None and vpc_region is not None

    _, _, aws_connect_kwargs = get_aws_connection_info(module)

    # connect to the route53 endpoint
    try:
        conn = Route53Connection(**aws_connect_kwargs)
    except boto.exception.BotoServerError, e:
        module.fail_json(msg=e.error_message)

    results = conn.get_all_hosted_zones()
    zones = {}

    for r53zone in results['ListHostedZonesResponse']['HostedZones']:
        zone_id = r53zone['Id'].replace('/hostedzone/', '')
        zone_details = conn.get_hosted_zone(zone_id)['GetHostedZoneResponse']
        if vpc_id and 'VPCs' in zone_details:
            # this is to deal with this boto bug: https://github.com/boto/boto/pull/2882
            if isinstance(zone_details['VPCs'], dict):
                if zone_details['VPCs']['VPC']['VPCId'] == vpc_id:
                    zones[r53zone['Name']] = zone_id
            else: # Forward compatibility for when boto fixes that bug
                if vpc_id in [v['VPCId'] for v in zone_details['VPCs']]:
                    zones[r53zone['Name']] = zone_id
        else:
            zones[r53zone['Name']] = zone_id

    record = {
        'private_zone': private_zone,
        'vpc_id': vpc_id,
        'vpc_region': vpc_region,
        'comment': comment,
    }

    if command == 'create' and zone_in in zones:
        if private_zone:
            details = conn.get_hosted_zone(zones[zone_in])

            if 'VPCs' not in details['GetHostedZoneResponse']:
                module.fail_json(
                    msg="Can't change VPC from public to private"
                )

            vpc_details = details['GetHostedZoneResponse']['VPCs']['VPC']
            current_vpc_id = vpc_details['VPCId']
            current_vpc_region = vpc_details['VPCRegion']

            if current_vpc_id != vpc_id:
                module.fail_json(
                    msg="Can't change VPC ID once a zone has been created"
                )
            if current_vpc_region != vpc_region:
                module.fail_json(
                    msg="Can't change VPC Region once a zone has been created"
                )

        record['zone_id'] = zones[zone_in]
        record['name'] = zone_in
        module.exit_json(changed=False, set=record)

    elif command == 'create':
        result = conn.create_hosted_zone(zone_in, **record)
        hosted_zone = result['CreateHostedZoneResponse']['HostedZone']
        zone_id = hosted_zone['Id'].replace('/hostedzone/', '')
        record['zone_id'] = zone_id
        record['name'] = zone_in
        module.exit_json(changed=True, set=record)

    elif command == 'delete' and zone_in in zones:
        conn.delete_hosted_zone(zones[zone_in])
        module.exit_json(changed=True)

    elif command == 'delete':
        module.exit_json(changed=False)

from ansible.module_utils.basic import *
from ansible.module_utils.ec2 import *

main()
<commit_msg>Update documentation to include usual fields<commit_after>#!/usr/bin/python

DOCUMENTATION = '''
module: route53_zone
short_description: add or delete Route53 zones
description:
    - Creates and deletes Route53 private and public zones
version_added: "2.0"
options:
    zone:
        description:
            - The DNS zone record (eg: foo.com.)
        required: true
    command:
        description:
            - whether or not the zone should exist or not
        required: false
        default: true
    vpc_id:
        description:
            - The VPC ID the zone should be a part of (if this is going to be a private zone)
        required: false
        default: null
    vpc_region:
        description:
            - The VPC Region the zone should be a part of (if this is going to be a private zone)
        required: false
        default: null
    comment:
        description:
            - Comment associated with the zone
        required: false
        default: ''
extends_documentation_fragment: aws
author: "Christopher Troup (@minichate)"
'''

import time

try:
    import boto
    import boto.ec2
    from boto import route53
    from boto.route53 import Route53Connection
    from boto.route53.zone import Zone
    HAS_BOTO = True
except ImportError:
    HAS_BOTO = False

def main():
    module = AnsibleModule(
        argument_spec=dict(
            zone=dict(required=True),
            command=dict(default='create', choices=['create', 'delete']),
            vpc_id=dict(default=None),
            vpc_region=dict(default=None),
            comment=dict(default=''),
        )
    )

    if not HAS_BOTO:
        module.fail_json(msg='boto required for this module')

    zone_in = module.params.get('zone').lower()
    command = module.params.get('command').lower()
    vpc_id = module.params.get('vpc_id')
    vpc_region = module.params.get('vpc_region')
    comment = module.params.get('comment')

    private_zone = vpc_id is not None and vpc_region is not None

    _, _, aws_connect_kwargs = get_aws_connection_info(module)

    # connect to the route53 endpoint
    try:
        conn = Route53Connection(**aws_connect_kwargs)
    except boto.exception.BotoServerError, e:
        module.fail_json(msg=e.error_message)

    results = conn.get_all_hosted_zones()
    zones = {}

    for r53zone in results['ListHostedZonesResponse']['HostedZones']:
        zone_id = r53zone['Id'].replace('/hostedzone/', '')
        zone_details = conn.get_hosted_zone(zone_id)['GetHostedZoneResponse']
        if vpc_id and 'VPCs' in zone_details:
            # this is to deal with this boto bug: https://github.com/boto/boto/pull/2882
            if isinstance(zone_details['VPCs'], dict):
                if zone_details['VPCs']['VPC']['VPCId'] == vpc_id:
                    zones[r53zone['Name']] = zone_id
            else: # Forward compatibility for when boto fixes that bug
                if vpc_id in [v['VPCId'] for v in zone_details['VPCs']]:
                    zones[r53zone['Name']] = zone_id
        else:
            zones[r53zone['Name']] = zone_id

    record = {
        'private_zone': private_zone,
        'vpc_id': vpc_id,
        'vpc_region': vpc_region,
        'comment': comment,
    }

    if command == 'create' and zone_in in zones:
        if private_zone:
            details = conn.get_hosted_zone(zones[zone_in])

            if 'VPCs' not in details['GetHostedZoneResponse']:
                module.fail_json(
                    msg="Can't change VPC from public to private"
                )

            vpc_details = details['GetHostedZoneResponse']['VPCs']['VPC']
            current_vpc_id = vpc_details['VPCId']
            current_vpc_region = vpc_details['VPCRegion']

            if current_vpc_id != vpc_id:
                module.fail_json(
                    msg="Can't change VPC ID once a zone has been created"
                )
            if current_vpc_region != vpc_region:
                module.fail_json(
                    msg="Can't change VPC Region once a zone has been created"
                )

        record['zone_id'] = zones[zone_in]
        record['name'] = zone_in
        module.exit_json(changed=False, set=record)

    elif command == 'create':
        result = conn.create_hosted_zone(zone_in, **record)
        hosted_zone = result['CreateHostedZoneResponse']['HostedZone']
        zone_id = hosted_zone['Id'].replace('/hostedzone/', '')
        record['zone_id'] = zone_id
        record['name'] = zone_in
        module.exit_json(changed=True, set=record)

    elif command == 'delete' and zone_in in zones:
        conn.delete_hosted_zone(zones[zone_in])
        module.exit_json(changed=True)

    elif command == 'delete':
        module.exit_json(changed=False)

from ansible.module_utils.basic import *
from ansible.module_utils.ec2 import *

main()
<|endoftext|>


<commit_before># coding=utf-8
# 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.

from absl import logging

import numpy as np
import os
import pandas as pd
from sklearn.utils import shuffle
import torch

from typing import Optional

from cabby.geo import regions
from cabby.geo import util as gutil
from cabby.model import dataset_item
from cabby.model import util
from transformers import DistilBertTokenizerFast, T5Tokenizer

MODELS = [
  'Dual-Encoder-Bert',
  'Classification-Bert',
  'S2-Generation-T5',
  'S2-Generation-T5-Landmarks',
  'S2-Generation-T5-Path']

T5_TYPE = "t5-small"
BERT_TYPE = 'distilbert-base-uncased'

# DISTRIBUTION_SCALE_DISTANCEA is a factor (in meters) that gives the overall
# scale in meters for the distribution.
DISTRIBUTION_SCALE_DISTANCE = 1000
dprob = util.DistanceProbability(DISTRIBUTION_SCALE_DISTANCE)

tokenizerT5 = T5Tokenizer.from_pretrained(T5_TYPE)

class Dataset:
  def __init__(self, data_dir: str, s2level: int, region: Optional[str], model_type: str):
    self.data_dir = data_dir
    self.s2level = s2level
    self.region = region
    self.model_type = model_type

    self.unique_cellid = {}
    self.cellid_to_label = {}
    self.label_to_cellid = {}

    self.train = None
    self.valid = None
    self.test = None

    self.set_tokenizers()

  def set_tokenizers(self):
    assert self.model_type in MODELS
    if self.model_type in ['Dual-Encoder-Bert', 'Classification-Bert']:
      self.text_tokenizer = DistilBertTokenizerFast.from_pretrained(BERT_TYPE)
      self.s2_tokenizer = util.binary_representation
    elif 'S2-Generation-T5' in self.model_type:
      self.text_tokenizer = T5Tokenizer.from_pretrained(T5_TYPE)
      self.s2_tokenizer = self.tokenize_cell

  def tokenize_cell(self, list_cells):
    if isinstance(list_cells[0], list):
      labels = []
      for c_list in list_cells:
        list_lables = []
        for c in c_list:
          list_lables.append(str(self.cellid_to_label[c]))

        labels.append('; '.join(list_lables))

    else:
      labels = [str(self.cellid_to_label[c]) for c in list_cells]

    return tokenizerT5(
      labels, return_tensors="pt", padding=True, truncation=True).input_ids

  def process_route(self, route_str):
    route_str = route_str.replace('LINESTRING', "").replace('(', "").replace(')', "")
    ladmarks_str_list = route_str.split(',')
    return [
      gutil.point_from_str_coord_xy(landmark_str) for landmark_str in ladmarks_str_list]

  def process_landmarks(self, landmarks_str_one_line):
    ladmarks_str_list = landmarks_str_one_line.split(';')
    return [gutil.point_from_str_coord_yx(
      landmark_str.split(':')[-1]) for landmark_str in ladmarks_str_list]

  def create_dataset(self, infer_only: bool = False
  ) -> dataset_item.TextGeoDataset:
    '''Loads data and creates datasets and train, validate and test sets.
    Returns:
      The train, validate and test sets and the dictionary of labels to cellids.
    '''

    points = gutil.get_centers_from_s2cellids(self.unique_cellid)

    unique_cells_df = pd.DataFrame(
      {'point': points, 'cellid': self.unique_cellid})

    unique_cells_df['far'] = unique_cells_df.point.swifter.apply(
        lambda x: gutil.far_cellid(x, unique_cells_df))

    vec_cells = self.s2_tokenizer(unique_cells_df.cellid.tolist())
    tens_cells = torch.tensor(vec_cells)

    # Create RUN dataset.
    train_dataset = None
    val_dataset = None
    logging.info("Starting to create the splits")
    if infer_only == False:
      train_dataset = dataset_item.TextGeoSplit(
        self.text_tokenizer,
        self.s2_tokenizer,
        self.train, self.s2level, unique_cells_df,
        self.cellid_to_label, dprob)
      logging.info(
        f"Finished to create the train-set with {len(train_dataset)} samples")
      val_dataset = dataset_item.TextGeoSplit(
        self.text_tokenizer,
        self.s2_tokenizer,
        self.valid, self.s2level, unique_cells_df,
        self.cellid_to_label, dprob)
      logging.info(
        f"Finished to create the valid-set with {len(val_dataset)} samples")
    test_dataset = dataset_item.TextGeoSplit(
      self.text_tokenizer,
      self.s2_tokenizer,
      self.test, self.s2level, unique_cells_df,
      self.cellid_to_label, dprob)
    logging.info(
      f"Finished to create the test-set with {len(test_dataset)} samples")

    return dataset_item.TextGeoDataset.from_TextGeoSplit(
      train_dataset, val_dataset, test_dataset,
      np.array(self.unique_cellid),
      tens_cells, self.label_to_cellid)

class HumanDataset(Dataset):
  def __init__(
    self, data_dir: str,
    s2level: int,
    region: Optional[str],
    model_type: str = "Dual-Encoder-Bert"):

    Dataset.__init__(self, data_dir, s2level, region, model_type)
    self.train = self.load_data(data_dir, 'train', lines=True)
    self.valid = self.load_data(data_dir, 'dev', lines=True)
    self.test = self.load_data(data_dir, 'test', lines=True)

    # Get labels.
    active_region = regions.get_region(region)
    unique_cellid = gutil.cellids_from_polygon(active_region.polygon, s2level)
    label_to_cellid = {idx: cellid for idx, cellid in enumerate(unique_cellid)}
    cellid_to_label = {cellid: idx for idx, cellid in enumerate(unique_cellid)}

    self.unique_cellid = unique_cellid
    self.label_to_cellid = label_to_cellid
    self.cellid_to_label = cellid_to_label

  def load_data(self, data_dir: str, ds_set: str, lines: bool):

    ds_path = os.path.join(data_dir, ds_set + '.json')
    assert os.path.exists(ds_path), f"{ds_path} doesn't exsits"

    ds = pd.read_json(ds_path, lines=lines)
    ds['instructions'] = ds['content']
    ds['end_point'] = ds['rvs_goal_point'].apply(gutil.point_from_str_coord_xy)
    ds['start_point'] = ds['rvs_start_point'].apply(gutil.point_from_str_coord_xy)

    if 'landmarks' in ds:
      ds['landmarks'] = ds.landmarks.apply(self.process_landmarks)

    if 'route' in ds:
      ds['route'] = ds.route.apply(self.process_route)

    columns_keep = ds.columns.difference(
      ['instructions', 'end_point', 'start_point', 'landmarks', 'route'])
    ds.drop(columns_keep, 1, inplace=True)

    ds = shuffle(ds)
    ds.reset_index(inplace=True, drop=True)
    return ds

class RUNDataset(Dataset):
  def __init__(
    self,
    data_dir: str,
    s2level: int,
    region: Optional[str],
    model_type: str = "Dual-Encoder-Bert"):

    Dataset.__init__(self, data_dir, s2level, None, model_type)

    train_ds, valid_ds, test_ds, ds = self.load_data(data_dir, lines=False)

    # Get labels.
    map_1 = regions.get_region("RUN-map1")
    map_2 = regions.get_region("RUN-map2")
    map_3 = regions.get_region("RUN-map3")

    logging.info(map_1.polygon.wkt)
    logging.info(map_2.polygon.wkt)
    logging.info(map_3.polygon.wkt)

    unique_cellid_map_1 = gutil.cellids_from_polygon(map_1.polygon, s2level)
    unique_cellid_map_2 = gutil.cellids_from_polygon(map_2.polygon, s2level)
    unique_cellid_map_3 = gutil.cellids_from_polygon(map_3.polygon, s2level)

    unique_cellid = (
      unique_cellid_map_1 + unique_cellid_map_2 + unique_cellid_map_3)
    label_to_cellid = {idx: cellid for idx, cellid in enumerate(unique_cellid)}
    cellid_to_label = {cellid: idx for idx, cellid in enumerate(unique_cellid)}

    self.train = train_ds
    self.valid = valid_ds
    self.test = test_ds
    self.ds = ds
    self.unique_cellid = unique_cellid
    self.label_to_cellid = label_to_cellid
    self.cellid_to_label = cellid_to_label

  def load_data(self, data_dir: str, lines: bool):
    ds = pd.read_json(os.path.join(data_dir, 'dataset.json'), lines=lines)
    ds['instructions'] = ds.groupby(
      ['id'])['instruction'].transform(lambda x: ' '.join(x))

    ds = ds.drop_duplicates(subset='id', keep="last")

    columns_keep = ds.columns.difference(
      ['map', 'id', 'instructions', 'end_point', 'start_point'])
    ds.drop(columns_keep, 1, inplace=True)

    ds = shuffle(ds)
    ds.reset_index(inplace=True, drop=True)

    dataset_size = ds.shape[0]
    logging.info(f"Size of dataset: {ds.shape[0]}")
    train_size = round(dataset_size * 80 / 100)
    valid_size = round(dataset_size * 10 / 100)

    train_ds = ds.iloc[:train_size]
    valid_ds = ds.iloc[train_size:train_size + valid_size]
    test_ds = ds.iloc[train_size + valid_size:]
    return train_ds, valid_ds, test_ds, ds

class RVSDataset(Dataset):
  def __init__(
    self, data_dir: str, s2level: int, region: Optional[str], model_type: str = "RVS"):
    Dataset.__init__(self, data_dir, s2level, region, model_type)
    train_ds = self.load_data(data_dir, 'train', True)
    valid_ds = self.load_data(data_dir, 'dev', True)
    test_ds = self.load_data(data_dir, 'test', True)

    # Get labels.
    active_region = regions.get_region(region)
    unique_cellid = gutil.cellids_from_polygon(active_region.polygon, s2level)
    label_to_cellid = {idx: cellid for idx, cellid in enumerate(unique_cellid)}
    cellid_to_label = {cellid: idx for idx, cellid in enumerate(unique_cellid)}

    self.train = train_ds
    self.valid = valid_ds
    self.test = test_ds
    self.unique_cellid = unique_cellid
    self.label_to_cellid = label_to_cellid
    self.cellid_to_label = cellid_to_label

  def load_data(self, data_dir: str, split: str, lines: bool):
    path_ds= os.path.join(data_dir, f'ds_{split}.json')
    assert os.path.exists(path_ds), path_ds
    ds = pd.read_json(path_ds, lines=lines)

    logging.info(f"Size of dataset before removal of duplication: {ds.shape[0]}")
    ds['landmarks'] =  ds.geo_landmarks.apply(self.process_landmarks)
    ds['route'] = ds.route.apply(self.process_route)

    ds = pd.concat([ds.drop(['geo_landmarks'], axis=1), ds['geo_landmarks'].apply(pd.Series)], axis=1)

    ds['end_osmid'] = ds.end_point.apply(lambda x: x[1])
    ds['start_osmid'] = ds.start_point.apply(lambda x: x[1])
    ds['end_pivot'] = ds.end_point
    ds['end_point'] = ds.end_point.apply(lambda x: gutil.point_from_list_coord_yx(x[3]))
    ds['start_point'] = ds.start_point.apply(lambda x: gutil.point_from_list_coord_yx(x[3]))
    ds = ds.drop_duplicates(subset=['end_osmid', 'start_osmid'], keep='last')
    return ds

  def process_landmarks(self, landmarks_dict):
    ladmarks_list = list(landmarks_dict.values())
    return [gutil.point_from_list_coord_yx(
      landmark_l[-1]) for landmark_l in ladmarks_list if landmark_l[-1]]

  def process_route(self, route_list):
    return [
      gutil.point_from_list_coord_xy(landmark) for landmark in route_list]

<commit_msg>fixed for prev versions<commit_after># coding=utf-8
# 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.

from absl import logging

import numpy as np
import os
import pandas as pd
from sklearn.utils import shuffle
import torch

from typing import Optional

from cabby.geo import regions
from cabby.geo import util as gutil
from cabby.model import dataset_item
from cabby.model import util
from transformers import DistilBertTokenizerFast, T5Tokenizer

MODELS = [
  'Dual-Encoder-Bert',
  'Classification-Bert',
  'S2-Generation-T5',
  'S2-Generation-T5-Landmarks',
  'S2-Generation-T5-Path']

T5_TYPE = "t5-small"
BERT_TYPE = 'distilbert-base-uncased'

# DISTRIBUTION_SCALE_DISTANCEA is a factor (in meters) that gives the overall
# scale in meters for the distribution.
DISTRIBUTION_SCALE_DISTANCE = 1000
dprob = util.DistanceProbability(DISTRIBUTION_SCALE_DISTANCE)

tokenizerT5 = T5Tokenizer.from_pretrained(T5_TYPE)

class Dataset:
  def __init__(self, data_dir: str, s2level: int, region: Optional[str], model_type: str):
    self.data_dir = data_dir
    self.s2level = s2level
    self.region = region
    self.model_type = model_type

    self.unique_cellid = {}
    self.cellid_to_label = {}
    self.label_to_cellid = {}

    self.train = None
    self.valid = None
    self.test = None

    self.set_tokenizers()

  def set_tokenizers(self):
    assert self.model_type in MODELS
    if self.model_type in ['Dual-Encoder-Bert', 'Classification-Bert']:
      self.text_tokenizer = DistilBertTokenizerFast.from_pretrained(BERT_TYPE)
      self.s2_tokenizer = util.binary_representation
    elif 'S2-Generation-T5' in self.model_type:
      self.text_tokenizer = T5Tokenizer.from_pretrained(T5_TYPE)
      self.s2_tokenizer = self.tokenize_cell

  def tokenize_cell(self, list_cells):
    if isinstance(list_cells[0], list):
      labels = []
      for c_list in list_cells:
        list_lables = []
        for c in c_list:
          list_lables.append(str(self.cellid_to_label[c]))

        labels.append('; '.join(list_lables))

    else:
      labels = [str(self.cellid_to_label[c]) for c in list_cells]

    return tokenizerT5(
      labels, return_tensors="pt", padding=True, truncation=True).input_ids

  def process_route(self, route_str):
    route_str = route_str.replace('LINESTRING', "").replace('(', "").replace(')', "")
    ladmarks_str_list = route_str.split(',')
    return [
      gutil.point_from_str_coord_xy(landmark_str) for landmark_str in ladmarks_str_list]

  def process_landmarks(self, landmarks_str_one_line):
    ladmarks_str_list = landmarks_str_one_line.split(';')
    return [gutil.point_from_str_coord_yx(
      landmark_str.split(':')[-1]) for landmark_str in ladmarks_str_list]

  def create_dataset(self, infer_only: bool = False
  ) -> dataset_item.TextGeoDataset:
    '''Loads data and creates datasets and train, validate and test sets.
    Returns:
      The train, validate and test sets and the dictionary of labels to cellids.
    '''

    points = gutil.get_centers_from_s2cellids(self.unique_cellid)

    unique_cells_df = pd.DataFrame(
      {'point': points, 'cellid': self.unique_cellid})

    unique_cells_df['far'] = unique_cells_df.point.swifter.apply(
        lambda x: gutil.far_cellid(x, unique_cells_df))

    vec_cells = self.s2_tokenizer(unique_cells_df.cellid.tolist())
    tens_cells = torch.tensor(vec_cells)

    # Create RUN dataset.
    train_dataset = None
    val_dataset = None
    logging.info("Starting to create the splits")
    if infer_only == False:
      train_dataset = dataset_item.TextGeoSplit(
        self.text_tokenizer,
        self.s2_tokenizer,
        self.train, self.s2level, unique_cells_df,
        self.cellid_to_label, dprob)
      logging.info(
        f"Finished to create the train-set with {len(train_dataset)} samples")
      val_dataset = dataset_item.TextGeoSplit(
        self.text_tokenizer,
        self.s2_tokenizer,
        self.valid, self.s2level, unique_cells_df,
        self.cellid_to_label, dprob)
      logging.info(
        f"Finished to create the valid-set with {len(val_dataset)} samples")
    test_dataset = dataset_item.TextGeoSplit(
      self.text_tokenizer,
      self.s2_tokenizer,
      self.test, self.s2level, unique_cells_df,
      self.cellid_to_label, dprob)
    logging.info(
      f"Finished to create the test-set with {len(test_dataset)} samples")

    return dataset_item.TextGeoDataset.from_TextGeoSplit(
      train_dataset, val_dataset, test_dataset,
      np.array(self.unique_cellid),
      tens_cells, self.label_to_cellid)

class HumanDataset(Dataset):
  def __init__(
    self, data_dir: str,
    s2level: int,
    region: Optional[str],
    model_type: str = "Dual-Encoder-Bert"):

    Dataset.__init__(self, data_dir, s2level, region, model_type)
    self.train = self.load_data(data_dir, 'train', lines=True)
    self.valid = self.load_data(data_dir, 'dev', lines=True)
    self.test = self.load_data(data_dir, 'test', lines=True)

    # Get labels.
    active_region = regions.get_region(region)
    unique_cellid = gutil.cellids_from_polygon(active_region.polygon, s2level)
    label_to_cellid = {idx: cellid for idx, cellid in enumerate(unique_cellid)}
    cellid_to_label = {cellid: idx for idx, cellid in enumerate(unique_cellid)}

    self.unique_cellid = unique_cellid
    self.label_to_cellid = label_to_cellid
    self.cellid_to_label = cellid_to_label

  def load_data(self, data_dir: str, ds_set: str, lines: bool):

    ds_path = os.path.join(data_dir, ds_set + '.json')
    assert os.path.exists(ds_path), f"{ds_path} doesn't exsits"

    ds = pd.read_json(ds_path, lines=lines)
    ds['instructions'] = ds['content']
    ds['end_point'] = ds['rvs_goal_point'].apply(gutil.point_from_str_coord_xy)
    ds['start_point'] = ds['rvs_start_point'].apply(gutil.point_from_str_coord_xy)

    if 'landmarks' in ds:
      ds['landmarks'] = ds.landmarks.apply(self.process_landmarks)

    if 'route' in ds:
      ds['route'] = ds.route.apply(self.process_route)

    columns_keep = ds.columns.difference(
      ['instructions', 'end_point', 'start_point', 'landmarks', 'route'])
    ds.drop(columns_keep, 1, inplace=True)

    ds = shuffle(ds)
    ds.reset_index(inplace=True, drop=True)
    return ds

class RUNDataset(Dataset):
  def __init__(
    self,
    data_dir: str,
    s2level: int,
    region: Optional[str],
    model_type: str = "Dual-Encoder-Bert"):

    Dataset.__init__(self, data_dir, s2level, None, model_type)

    train_ds, valid_ds, test_ds, ds = self.load_data(data_dir, lines=False)

    # Get labels.
    map_1 = regions.get_region("RUN-map1")
    map_2 = regions.get_region("RUN-map2")
    map_3 = regions.get_region("RUN-map3")

    logging.info(map_1.polygon.wkt)
    logging.info(map_2.polygon.wkt)
    logging.info(map_3.polygon.wkt)

    unique_cellid_map_1 = gutil.cellids_from_polygon(map_1.polygon, s2level)
    unique_cellid_map_2 = gutil.cellids_from_polygon(map_2.polygon, s2level)
    unique_cellid_map_3 = gutil.cellids_from_polygon(map_3.polygon, s2level)

    unique_cellid = (
      unique_cellid_map_1 + unique_cellid_map_2 + unique_cellid_map_3)
    label_to_cellid = {idx: cellid for idx, cellid in enumerate(unique_cellid)}
    cellid_to_label = {cellid: idx for idx, cellid in enumerate(unique_cellid)}

    self.train = train_ds
    self.valid = valid_ds
    self.test = test_ds
    self.ds = ds
    self.unique_cellid = unique_cellid
    self.label_to_cellid = label_to_cellid
    self.cellid_to_label = cellid_to_label

  def load_data(self, data_dir: str, lines: bool):
    ds = pd.read_json(os.path.join(data_dir, 'dataset.json'), lines=lines)
    ds['instructions'] = ds.groupby(
      ['id'])['instruction'].transform(lambda x: ' '.join(x))

    ds = ds.drop_duplicates(subset='id', keep="last")

    columns_keep = ds.columns.difference(
      ['map', 'id', 'instructions', 'end_point', 'start_point'])
    ds.drop(columns_keep, 1, inplace=True)

    ds = shuffle(ds)
    ds.reset_index(inplace=True, drop=True)

    dataset_size = ds.shape[0]
    logging.info(f"Size of dataset: {ds.shape[0]}")
    train_size = round(dataset_size * 80 / 100)
    valid_size = round(dataset_size * 10 / 100)

    train_ds = ds.iloc[:train_size]
    valid_ds = ds.iloc[train_size:train_size + valid_size]
    test_ds = ds.iloc[train_size + valid_size:]
    return train_ds, valid_ds, test_ds, ds

class RVSDataset(Dataset):
  def __init__(
    self, data_dir: str, s2level: int, region: Optional[str], model_type: str = "RVS"):
    Dataset.__init__(self, data_dir, s2level, region, model_type)
    train_ds = self.load_data(data_dir, 'train', True)
    valid_ds = self.load_data(data_dir, 'dev', True)
    test_ds = self.load_data(data_dir, 'test', True)

    # Get labels.
    active_region = regions.get_region(region)
    unique_cellid = gutil.cellids_from_polygon(active_region.polygon, s2level)
    label_to_cellid = {idx: cellid for idx, cellid in enumerate(unique_cellid)}
    cellid_to_label = {cellid: idx for idx, cellid in enumerate(unique_cellid)}

    self.train = train_ds
    self.valid = valid_ds
    self.test = test_ds
    self.unique_cellid = unique_cellid
    self.label_to_cellid = label_to_cellid
    self.cellid_to_label = cellid_to_label

  def load_data(self, data_dir: str, split: str, lines: bool):
    path_ds= os.path.join(data_dir, f'ds_{split}.json')
    assert os.path.exists(path_ds), path_ds
    ds = pd.read_json(path_ds, lines=lines)

    logging.info(f"Size of dataset before removal of duplication: {ds.shape[0]}")

    if 'landmarks' in ds:
      ds['landmarks'] =  ds.geo_landmarks.apply(self.process_landmarks)

    if 'route' in ds:
      ds['route'] = ds.route.apply(self.process_route)

    ds = pd.concat([ds.drop(['geo_landmarks'], axis=1), ds['geo_landmarks'].apply(pd.Series)], axis=1)

    ds['end_osmid'] = ds.end_point.apply(lambda x: x[1])
    ds['start_osmid'] = ds.start_point.apply(lambda x: x[1])
    ds['end_pivot'] = ds.end_point
    ds['end_point'] = ds.end_point.apply(lambda x: gutil.point_from_list_coord_yx(x[3]))
    ds['start_point'] = ds.start_point.apply(lambda x: gutil.point_from_list_coord_yx(x[3]))
    ds = ds.drop_duplicates(subset=['end_osmid', 'start_osmid'], keep='last')
    return ds

  def process_landmarks(self, landmarks_dict):
    ladmarks_list = list(landmarks_dict.values())
    return [gutil.point_from_list_coord_yx(
      landmark_l[-1]) for landmark_l in ladmarks_list if landmark_l[-1]]

  def process_route(self, route_list):
    return [
      gutil.point_from_list_coord_xy(landmark) for landmark in route_list]

<|endoftext|>


<commit_before>"""
bootstrap.validators
~~~~~~~~~~~~~~~~~~~~

Validators for use with Click_.

.. _Click: http://click.pocoo.org/

:Copyright: 2006-2018 Jochen Kupperschmidt
:License: Modified BSD, see LICENSE for details.
"""

import click

from byceps.services.brand import service as brand_service
from byceps.services.party import service as party_service
from byceps.services.user import service as user_service

def validate_brand(ctx, param, brand_id):
    brand = brand_service.find_brand(brand_id)

    if not brand:
        raise click.BadParameter('Unknown brand ID "{}".'.format(brand_id))

    return brand

def validate_party(ctx, param, party_id):
    party = party_service.find_party(party_id)

    if not party:
        raise click.BadParameter('Unknown party ID "{}".'.format(party_id))

    return party

def validate_user_id(ctx, param, user_id):
    user = user_service.find_user(user_id)

    if not user:
        raise click.BadParameter('Unknown user ID "{}".'.format(user_id))

    return user

def validate_user_screen_name(ctx, param, screen_name):
    user = user_service.find_user_by_screen_name(screen_name)

    if not user:
        raise click.BadParameter('Unknown user screen name "{}".'
                                 .format(screen_name))

    return user
<commit_msg>Add (non-Click) type hints to script parameter validators<commit_after>"""
bootstrap.validators
~~~~~~~~~~~~~~~~~~~~

Validators for use with Click_.

.. _Click: http://click.pocoo.org/

:Copyright: 2006-2018 Jochen Kupperschmidt
:License: Modified BSD, see LICENSE for details.
"""

import click

from byceps.services.brand.models.brand import Brand
from byceps.services.brand import service as brand_service
from byceps.services.party.models.party import Party
from byceps.services.party import service as party_service
from byceps.services.user.models.user import User
from byceps.services.user import service as user_service
from byceps.typing import BrandID, PartyID, UserID

def validate_brand(ctx, param, brand_id: BrandID) -> Brand:
    brand = brand_service.find_brand(brand_id)

    if not brand:
        raise click.BadParameter('Unknown brand ID "{}".'.format(brand_id))

    return brand

def validate_party(ctx, param, party_id: PartyID) -> Party:
    party = party_service.find_party(party_id)

    if not party:
        raise click.BadParameter('Unknown party ID "{}".'.format(party_id))

    return party

def validate_user_id(ctx, param, user_id: UserID) -> User:
    user = user_service.find_user(user_id)

    if not user:
        raise click.BadParameter('Unknown user ID "{}".'.format(user_id))

    return user

def validate_user_screen_name(ctx, param, screen_name: str) -> User:
    user = user_service.find_user_by_screen_name(screen_name)

    if not user:
        raise click.BadParameter('Unknown user screen name "{}".'
                                 .format(screen_name))

    return user
<|endoftext|>


<commit_before>#!/usr/bin/python
#
# 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.
"""Functions required to interact with Docker to build and run images, shells
and notebooks in a Docker environment.

"""

from __future__ import absolute_import, division, print_function

import subprocess
import sys
from typing import Any, Dict, Iterable, List, Optional

import tqdm
from absl import logging
from blessings import Terminal
from tqdm.utils import _screen_shape_wrapper

import caliban.config as c
import caliban.config.experiment as ce
import caliban.docker.build as b
import caliban.util.fs as ufs
import caliban.util.tqdm as ut
from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform
from caliban.history.util import (create_experiments, generate_container_spec,
                                  get_mem_engine, get_sql_engine, session_scope)

t = Terminal()

def _run_cmd(job_mode: c.JobMode,
             run_args: Optional[List[str]] = None) -> List[str]:
  """Returns the sequence of commands for the subprocess run functions required
  to execute `docker run`. in CPU or GPU mode, depending on the value of
  job_mode.

  Keyword args:
  - run_args: list of args to pass to docker run.

  """
  if run_args is None:
    run_args = []

  runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else []
  return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args

def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec:
  """Prints logging as a side effect for the supplied sequence of job specs
  generated from an experiment definition; returns the input job spec.

  """
  args = ce.experiment_to_args(job_spec.experiment.kwargs,
                               job_spec.experiment.args)
  logging.info("")
  logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args))))
  return job_spec

def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]:
  """Generates an iterable of job specs that should be passed to `docker run` to
  execute the experiments defined by the supplied iterable.

  """
  for i, s in enumerate(job_specs, 1):
    yield log_job_spec_instance(s, i)

def execute_dry_run(job_specs: Iterable[JobSpec]) -> None:
  """Expands the supplied sequence of experiments into sequences of args and logs
  the jobs that WOULD have been executed, had the dry run flag not been
  applied.

  """
  list(logged_job_specs(job_specs))

  logging.info('')
  logging.info(
      t.yellow("To build your image and execute these jobs, \
run your command again without {}.".format(c.DRY_RUN_FLAG)))
  logging.info('')
  return None

def local_callback(idx: int, job: Job) -> None:
  """Provides logging feedback for jobs run locally. If the return code is 0,
  logs success; else, logs the failure as an error and logs the script args
  that provided the failure.

  """
  if job.status == JobStatus.SUCCEEDED:
    logging.info(t.green(f'Job {idx} succeeded!'))
  else:
    logging.error(
        t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.'))
    args = ce.experiment_to_args(job.spec.experiment.kwargs,
                                 job.spec.experiment.args)
    logging.error(t.red(f'Failing args for job {idx}: {args}'))

def window_size_env_cmds():
  """Returns a sequence of `docker run` arguments that will internally configure
  the terminal columns and lines, so that progress bars and other terminal
  interactions will work properly.

  These aren't required for interactive Docker commands like those triggered by
  `caliban shell`.

  """
  ret = []
  cols, lines = _screen_shape_wrapper()(0)
  if cols:
    ret += ["-e", f"COLUMNS={cols}"]
  if lines:
    ret += ["-e", f"LINES={lines}"]
  return ret

# ----------------------------------------------------------------------------
def _create_job_spec_dict(
    experiment: Experiment,
    job_mode: c.JobMode,
    image_id: str,
    run_args: Optional[List[str]] = None,
) -> Dict[str, Any]:
  '''creates a job spec dictionary for a local job'''

  # Without the unbuffered environment variable, stderr and stdout won't be
  # emitted in the proper order from inside the container.
  terminal_cmds = ["-e" "PYTHONUNBUFFERED=1"] + window_size_env_cmds()

  base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id]
  command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args)
  return {'command': command, 'container': image_id}

# ----------------------------------------------------------------------------
def execute_jobs(
    job_specs: Iterable[JobSpec],
    dry_run: bool = False,
):
  '''executes a sequence of jobs based on job specs

  Arg:
  job_specs: specifications for jobs to be executed
  dry_run: if True, only print what would be done
  '''

  with ut.tqdm_logging() as orig_stream:
    pbar = tqdm.tqdm(logged_job_specs(job_specs),
                     file=orig_stream,
                     total=len(job_specs),
                     ascii=True,
                     unit="experiment",
                     desc="Executing")
    for idx, job_spec in enumerate(pbar, 1):
      command = job_spec.spec['command']
      logging.info(f'Running command: {" ".join(command)}')
      if not dry_run:
        _, ret_code = ufs.capture_stdout(command, "", ut.TqdmFile(sys.stderr))
      else:
        ret_code = 0
      j = Job(spec=job_spec,
              container=job_spec.spec['container'],
              details={'ret_code': ret_code},
              status=JobStatus.SUCCEEDED if ret_code == 0 else JobStatus.FAILED)
      local_callback(idx=idx, job=j)

  if dry_run:
    logging.info(
        t.yellow(f'\nTo build your image and execute these jobs, '
                 f'run your command again without {c.DRY_RUN_FLAG}\n'))

  return None

def run_experiments(job_mode: c.JobMode,
                    run_args: Optional[List[str]] = None,
                    script_args: Optional[List[str]] = None,
                    image_id: Optional[str] = None,
                    dry_run: bool = False,
                    experiment_config: Optional[ce.ExpConf] = None,
                    xgroup: Optional[str] = None,
                    **build_image_kwargs) -> None:
  """Builds an image using the supplied **build_image_kwargs and calls `docker
  run` on the resulting image using sensible defaults.

  Keyword args:

  - job_mode: c.JobMode.

  - run_args: extra arguments to supply to `docker run` after our defaults.
  - script_args: extra arguments to supply to the entrypoint. (You can
  - override the default container entrypoint by supplying a new one inside
    run_args.)
  - image_id: ID of the image to run. Supplying this will skip an image build.
  - experiment_config: dict of string to list, boolean, string or int. Any
    lists will trigger a cartesian product out with the rest of the config. A
    job will be executed for every combination of parameters in the experiment
    config.
  - dry_run: if True, no actual jobs will be executed and docker won't
    actually build; logging side effects will show the user what will happen
    without dry_run=True.

  any extra kwargs supplied are passed through to build_image.
  """
  if run_args is None:
    run_args = []

  if script_args is None:
    script_args = []

  if experiment_config is None:
    experiment_config = {}

  docker_args = {k: v for k, v in build_image_kwargs.items()}
  docker_args['job_mode'] = job_mode

  engine = get_mem_engine() if dry_run else get_sql_engine()

  with session_scope(engine) as session:
    container_spec = generate_container_spec(session, docker_args, image_id)

    if image_id is None:
      if dry_run:
        logging.info("Dry run - skipping actual 'docker build'.")
        image_id = 'dry_run_tag'
      else:
        image_id = b.build_image(**docker_args)

    experiments = create_experiments(
        session=session,
        container_spec=container_spec,
        script_args=script_args,
        experiment_config=experiment_config,
        xgroup=xgroup,
    )

    job_specs = [
        JobSpec.get_or_create(
            experiment=x,
            spec=_create_job_spec_dict(
                experiment=x,
                job_mode=job_mode,
                run_args=run_args,
                image_id=image_id,
            ),
            platform=Platform.LOCAL,
        ) for x in experiments
    ]

    try:
      execute_jobs(job_specs=job_specs, dry_run=dry_run)
    except Exception as e:
      logging.error(f'exception: {e}')
      session.commit()  # commit here, otherwise will be rolled back

def run(job_mode: c.JobMode,
        run_args: Optional[List[str]] = None,
        script_args: Optional[List[str]] = None,
        image_id: Optional[str] = None,
        **build_image_kwargs) -> None:
  """Builds an image using the supplied **build_image_kwargs and calls `docker
  run` on the resulting image using sensible defaults.
  Keyword args:
  - job_mode: c.JobMode.
  - run_args: extra arguments to supply to `docker run` after our defaults.
  - script_args: extra arguments to supply to the entrypoint. (You can
  - override the default container entrypoint by supplying a new one inside
    run_args.)
  - image_id: ID of the image to run. Supplying this will skip an image build.
  any extra kwargs supplied are passed through to build_image.
  """
  if run_args is None:
    run_args = []

  if script_args is None:
    script_args = []

  if image_id is None:
    image_id = b.build_image(job_mode, **build_image_kwargs)

  base_cmd = _run_cmd(job_mode, run_args)

  command = base_cmd + [image_id] + script_args

  logging.info("Running command: {}".format(' '.join(command)))
  subprocess.call(command)
  return None
<commit_msg>Add missing comma in list (#47)<commit_after>#!/usr/bin/python
#
# 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.
"""Functions required to interact with Docker to build and run images, shells
and notebooks in a Docker environment.

"""

from __future__ import absolute_import, division, print_function

import subprocess
import sys
from typing import Any, Dict, Iterable, List, Optional

import tqdm
from absl import logging
from blessings import Terminal
from tqdm.utils import _screen_shape_wrapper

import caliban.config as c
import caliban.config.experiment as ce
import caliban.docker.build as b
import caliban.util.fs as ufs
import caliban.util.tqdm as ut
from caliban.history.types import Experiment, Job, JobSpec, JobStatus, Platform
from caliban.history.util import (create_experiments, generate_container_spec,
                                  get_mem_engine, get_sql_engine, session_scope)

t = Terminal()

def _run_cmd(job_mode: c.JobMode,
             run_args: Optional[List[str]] = None) -> List[str]:
  """Returns the sequence of commands for the subprocess run functions required
  to execute `docker run`. in CPU or GPU mode, depending on the value of
  job_mode.

  Keyword args:
  - run_args: list of args to pass to docker run.

  """
  if run_args is None:
    run_args = []

  runtime = ["--runtime", "nvidia"] if c.gpu(job_mode) else []
  return ["docker", "run"] + runtime + ["--ipc", "host"] + run_args

def log_job_spec_instance(job_spec: JobSpec, i: int) -> JobSpec:
  """Prints logging as a side effect for the supplied sequence of job specs
  generated from an experiment definition; returns the input job spec.

  """
  args = ce.experiment_to_args(job_spec.experiment.kwargs,
                               job_spec.experiment.args)
  logging.info("")
  logging.info("Job {} - Experiment args: {}".format(i, t.yellow(str(args))))
  return job_spec

def logged_job_specs(job_specs: Iterable[JobSpec]) -> Iterable[JobSpec]:
  """Generates an iterable of job specs that should be passed to `docker run` to
  execute the experiments defined by the supplied iterable.

  """
  for i, s in enumerate(job_specs, 1):
    yield log_job_spec_instance(s, i)

def execute_dry_run(job_specs: Iterable[JobSpec]) -> None:
  """Expands the supplied sequence of experiments into sequences of args and logs
  the jobs that WOULD have been executed, had the dry run flag not been
  applied.

  """
  list(logged_job_specs(job_specs))

  logging.info('')
  logging.info(
      t.yellow("To build your image and execute these jobs, \
run your command again without {}.".format(c.DRY_RUN_FLAG)))
  logging.info('')
  return None

def local_callback(idx: int, job: Job) -> None:
  """Provides logging feedback for jobs run locally. If the return code is 0,
  logs success; else, logs the failure as an error and logs the script args
  that provided the failure.

  """
  if job.status == JobStatus.SUCCEEDED:
    logging.info(t.green(f'Job {idx} succeeded!'))
  else:
    logging.error(
        t.red(f'Job {idx} failed with return code {job.details["ret_code"]}.'))
    args = ce.experiment_to_args(job.spec.experiment.kwargs,
                                 job.spec.experiment.args)
    logging.error(t.red(f'Failing args for job {idx}: {args}'))

def window_size_env_cmds():
  """Returns a sequence of `docker run` arguments that will internally configure
  the terminal columns and lines, so that progress bars and other terminal
  interactions will work properly.

  These aren't required for interactive Docker commands like those triggered by
  `caliban shell`.

  """
  ret = []
  cols, lines = _screen_shape_wrapper()(0)
  if cols:
    ret += ["-e", f"COLUMNS={cols}"]
  if lines:
    ret += ["-e", f"LINES={lines}"]
  return ret

# ----------------------------------------------------------------------------
def _create_job_spec_dict(
    experiment: Experiment,
    job_mode: c.JobMode,
    image_id: str,
    run_args: Optional[List[str]] = None,
) -> Dict[str, Any]:
  '''creates a job spec dictionary for a local job'''

  # Without the unbuffered environment variable, stderr and stdout won't be
  # emitted in the proper order from inside the container.
  terminal_cmds = ["-e", "PYTHONUNBUFFERED=1"] + window_size_env_cmds()

  base_cmd = _run_cmd(job_mode, run_args) + terminal_cmds + [image_id]
  command = base_cmd + ce.experiment_to_args(experiment.kwargs, experiment.args)
  return {'command': command, 'container': image_id}

# ----------------------------------------------------------------------------
def execute_jobs(
    job_specs: Iterable[JobSpec],
    dry_run: bool = False,
):
  '''executes a sequence of jobs based on job specs

  Arg:
  job_specs: specifications for jobs to be executed
  dry_run: if True, only print what would be done
  '''

  with ut.tqdm_logging() as orig_stream:
    pbar = tqdm.tqdm(logged_job_specs(job_specs),
                     file=orig_stream,
                     total=len(job_specs),
                     ascii=True,
                     unit="experiment",
                     desc="Executing")
    for idx, job_spec in enumerate(pbar, 1):
      command = job_spec.spec['command']
      logging.info(f'Running command: {" ".join(command)}')
      if not dry_run:
        _, ret_code = ufs.capture_stdout(command, "", ut.TqdmFile(sys.stderr))
      else:
        ret_code = 0
      j = Job(spec=job_spec,
              container=job_spec.spec['container'],
              details={'ret_code': ret_code},
              status=JobStatus.SUCCEEDED if ret_code == 0 else JobStatus.FAILED)
      local_callback(idx=idx, job=j)

  if dry_run:
    logging.info(
        t.yellow(f'\nTo build your image and execute these jobs, '
                 f'run your command again without {c.DRY_RUN_FLAG}\n'))

  return None

def run_experiments(job_mode: c.JobMode,
                    run_args: Optional[List[str]] = None,
                    script_args: Optional[List[str]] = None,
                    image_id: Optional[str] = None,
                    dry_run: bool = False,
                    experiment_config: Optional[ce.ExpConf] = None,
                    xgroup: Optional[str] = None,
                    **build_image_kwargs) -> None:
  """Builds an image using the supplied **build_image_kwargs and calls `docker
  run` on the resulting image using sensible defaults.

  Keyword args:

  - job_mode: c.JobMode.

  - run_args: extra arguments to supply to `docker run` after our defaults.
  - script_args: extra arguments to supply to the entrypoint. (You can
  - override the default container entrypoint by supplying a new one inside
    run_args.)
  - image_id: ID of the image to run. Supplying this will skip an image build.
  - experiment_config: dict of string to list, boolean, string or int. Any
    lists will trigger a cartesian product out with the rest of the config. A
    job will be executed for every combination of parameters in the experiment
    config.
  - dry_run: if True, no actual jobs will be executed and docker won't
    actually build; logging side effects will show the user what will happen
    without dry_run=True.

  any extra kwargs supplied are passed through to build_image.
  """
  if run_args is None:
    run_args = []

  if script_args is None:
    script_args = []

  if experiment_config is None:
    experiment_config = {}

  docker_args = {k: v for k, v in build_image_kwargs.items()}
  docker_args['job_mode'] = job_mode

  engine = get_mem_engine() if dry_run else get_sql_engine()

  with session_scope(engine) as session:
    container_spec = generate_container_spec(session, docker_args, image_id)

    if image_id is None:
      if dry_run:
        logging.info("Dry run - skipping actual 'docker build'.")
        image_id = 'dry_run_tag'
      else:
        image_id = b.build_image(**docker_args)

    experiments = create_experiments(
        session=session,
        container_spec=container_spec,
        script_args=script_args,
        experiment_config=experiment_config,
        xgroup=xgroup,
    )

    job_specs = [
        JobSpec.get_or_create(
            experiment=x,
            spec=_create_job_spec_dict(
                experiment=x,
                job_mode=job_mode,
                run_args=run_args,
                image_id=image_id,
            ),
            platform=Platform.LOCAL,
        ) for x in experiments
    ]

    try:
      execute_jobs(job_specs=job_specs, dry_run=dry_run)
    except Exception as e:
      logging.error(f'exception: {e}')
      session.commit()  # commit here, otherwise will be rolled back

def run(job_mode: c.JobMode,
        run_args: Optional[List[str]] = None,
        script_args: Optional[List[str]] = None,
        image_id: Optional[str] = None,
        **build_image_kwargs) -> None:
  """Builds an image using the supplied **build_image_kwargs and calls `docker
  run` on the resulting image using sensible defaults.
  Keyword args:
  - job_mode: c.JobMode.
  - run_args: extra arguments to supply to `docker run` after our defaults.
  - script_args: extra arguments to supply to the entrypoint. (You can
  - override the default container entrypoint by supplying a new one inside
    run_args.)
  - image_id: ID of the image to run. Supplying this will skip an image build.
  any extra kwargs supplied are passed through to build_image.
  """
  if run_args is None:
    run_args = []

  if script_args is None:
    script_args = []

  if image_id is None:
    image_id = b.build_image(job_mode, **build_image_kwargs)

  base_cmd = _run_cmd(job_mode, run_args)

  command = base_cmd + [image_id] + script_args

  logging.info("Running command: {}".format(' '.join(command)))
  subprocess.call(command)
  return None
<|endoftext|>


<commit_before>from calvin.utilities.nodecontrol import dispatch_node
from calvin.utilities.attribute_resolver import format_index_string
from calvin.utilities import calvinconfig
from calvin.utilities import calvinlogger
from calvin.requests.request_handler import RT
import os
import time
import multiprocessing
import copy
import numbers

_log = calvinlogger.get_logger(__name__)
_conf = calvinconfig.get()

def retry(retries, function, criterion, error_msg):
    """
        Executes 'result = function()' until 'criterion(result)' evaluates to a true value.
        Raises 'Exception(error_msg)' if criterion is not fulfilled after 'retries' attempts

    """
    delay = 0.1
    retry = 0
    result = None
    while retry < retries:
        try:
            result = function()
            try:
                if criterion(result):
                    if retry > 0:
                        _log.info("Criterion %s(%s) satisfied after %d retries" %
                                    (str(criterion), str(function), retry,))
                    return result
            except Exception as e:
                _log.error("Erroneous criteria '%r" % (e, ))
                raise e
        except Exception as e:
            _log.info("Encountered exception when retrying '%s'" % (e,))
        delay = min(2, delay * 1.5); retry += 1
        time.sleep(delay)
        try:
            r = result if isinstance(result, (numbers.Number, type(None))) else len(result)
        except:
            r = None
        _log.info("Criterion still not satisfied after %d retries, result (length) %s" % (retry, r))
    _log.info("Criterion %s(%s) never satisfied, last full result %s" % (str(criterion), str(function), str(result)))
    raise Exception(error_msg)

def wait_for_tokens(request_handler, rt, actor_id, size=5, retries=10):
    """An alias for 'actual_tokens'"""
    return actual_tokens(request_handler, rt, actor_id, size, retries)

def actual_tokens(request_handler, rt, actor_id, size=5, retries=10):
    """
    Uses 'request_handler' to fetch the report from actor 'actor_id' on runtime 'rt'.
    """
    from functools import partial
    func = partial(request_handler.report, rt, actor_id)
    criterion = lambda tokens: len(tokens) >= size
    return retry(retries, func, criterion, "Not enough tokens, expected %d" % size)

def multi_report(request_handler, rt, actor_ids):
    """
    Helper uses 'request_handler' to fetch the report from actors in 'actor_ids' list on runtime(s) 'rt'.
    """
    result = []
    if isinstance(rt, (list, tuple, set)):
        args = zip(rt, actor_ids)
    else:
        args = zip([rt]*len(actor_ids), actor_ids)
    for runtime, actor_id in args:
        result.append(request_handler.report(runtime, actor_id))
    return result

def actual_tokens_multiple(request_handler, rt, actor_ids, size=5, retries=10):
    """
    Uses 'request_handler' to fetch the report from actors in 'actor_ids' list on runtime(s) 'rt'.
    """
    from functools import partial
    func = partial(multi_report, request_handler, rt, actor_ids)
    criterion = lambda tokens: sum([len(t) for t in tokens]) >= size
    return retry(retries, func, criterion, "Not enough tokens, expected %d" % size)

def destroy_app(deployer, retries=10):
    """
    Tries to destroy the app connected with deployer.
    """
    return delete_app(deployer.request_handler, deployer.runtime, deployer.app_id, retries=retries)

def deploy_app(request_handler, deployer, runtimes, retries=10):
    """
    Deploys app associated w/ deployer and then tries to verify its
    presence in registry (for all runtimes).
    """
    deployer.deploy()

    def check_application():
        for rt in runtimes:
            try:
                if request_handler.get_application(rt, deployer.app_id) is None:
                    return False
            except:
                return False
        _log.info("Application found on all peers, continuing")
        return True

    return retry(retries, check_application, lambda r: r, "Application not found on all peers")

def delete_app(request_handler, runtime, app_id, check_actor_ids=None, retries=10):
    """
    Deletes an app and then tries to verify it is actually gone.
    """
    from functools import partial

    def verify_app_gone(request_handler, runtime, app_id):
        try:
            request_handler.get_application(runtime, app_id)
            return False
        except:
            return True

    def verify_actors_gone(request_handler, runtime, actor_ids):
        responses = []
        for actor_id in actor_ids:
            responses.append(request_handler.async_get_actor(runtime, actor_id))
        gone = True
        for r in responses:
            try:
                response = request_handler.async_response(r)
                if response is not None:
                    gone = False
            except:
                pass
        return gone

    try:
        request_handler.delete_application(runtime, app_id)
    except Exception as e:
        msg = str(e.message)
        if msg.startswith("500"):
            _log.error("Delete App got 500")
        elif msg.startswith("404"):
            _log.error("Delete App got 404")
        else:
            _log.error("Delete App got unknown error %s" % str(msg))

    retry(retries, partial(verify_app_gone, request_handler, runtime, app_id), lambda r: r, "Application not deleted")
    if check_actor_ids:
        retry(retries, partial(verify_actors_gone, request_handler, runtime, check_actor_ids),
              lambda r: r, "Application actors not deleted")

def deploy_script(request_handler, name, script, runtime, retries=10):
    """
    Deploys script and then tries to verify its
    presence in registry on the runtime.
    """

    response = request_handler.deploy_application(runtime, name, script)
    app_id = response['application_id']

    def check_application():
        try:
            if request_handler.get_application(runtime, app_id) is None:
                return False
        except:
            return False
        _log.info("Application found, continuing")
        return True

    retry(retries, check_application, lambda r: r, "Application not found")
    return response

def flatten_zip(lz):
    return [] if not lz else [ lz[0][0], lz[0][1] ] + flatten_zip(lz[1:])

# Helper for 'std.CountTimer' actor
def expected_counter(n):
    return [i for i in range(1, n + 1)]

# Helper for 'std.Sum'
def expected_sum(n):
    def cumsum(l):
        s = 0
        for n in l:
            s = s + n
            yield s

    return list(cumsum(range(1, n + 1)))

def expected_tokens(request_handler, rt, actor_id, t_type):

    tokens = request_handler.report(rt, actor_id)

    if t_type == 'seq':
        return expected_counter(tokens)

    if t_type == 'sum':
        return expected_sum(tokens)

    return None

def setup_distributed(control_uri, purpose, request_handler):
    from functools import partial

    remote_node_count = 3
    test_peers = None
    runtimes = []

    runtime = RT(control_uri)
    index = {"node_name": {"organization": "com.ericsson", "purpose": purpose}}
    index_string = format_index_string(index)

    get_index = partial(request_handler.get_index, runtime, index_string)

    def criteria(peers):
        return peers and peers.get("result", None) and len(peers["result"]) >= remote_node_count

    test_peers = retry(10, get_index, criteria, "Not all nodes found")
    test_peers = test_peers["result"]

    for peer_id in test_peers:
        peer = request_handler.get_node(runtime, peer_id)
        if not peer:
            _log.warning("Runtime '%r' peer '%r' does not exist" % (runtime, peer_id, ))
            continue
        rt = RT(peer["control_uri"])
        rt.id = peer_id
        rt.uri = peer["uri"]
        runtimes.append(rt)

    return runtimes

def setup_local(ip_addr, request_handler, nbr, proxy_storage):
    def check_storage(rt, n, index):
        index_string = format_index_string(index)
        retries = 0
        while retries < 120:
            try:
                retries += 1
                peers = request_handler.get_index(rt, index_string, timeout=60)
            except:
                _log.info("Timed out when finding peers retrying")
                retries += 39  # A timeout counts more we don't want to wait 60*100 seconds
                continue
            if len(peers['result']) >= n:
                _log.info("Found %d peers (%r)", len(peers['result']), peers['result'])
                return
            _log.info("Only %d peers found (%r)", len(peers['result']), peers['result'])
            time.sleep(1)
        # No more retrying
        raise Exception("Storage check failed, could not find peers.")

    hosts = [
        ("calvinip://%s:%d" % (ip_addr, d), "http://%s:%d" % (ip_addr, d+1)) for d in range(5200, 5200 + 2 * nbr, 2)
    ]

    runtimes = []

    host = hosts[0]
    attr = {u'indexed_public': {u'node_name': {u'organization': u'com.ericsson', u'purpose': u'distributed-test'}}}
    attr_first = copy.deepcopy(attr)
    attr_first['indexed_public']['node_name']['group'] = u'first'
    attr_first['indexed_public']['node_name']['name'] = u'runtime1'
    attr_rest = copy.deepcopy(attr)
    attr_rest['indexed_public']['node_name']['group'] = u'rest'

    _log.info("starting runtime %s %s" % host)

    if proxy_storage:
        import calvin.runtime.north.storage
        calvin.runtime.north.storage._conf.set('global', 'storage_type', 'local')
    rt, _ = dispatch_node([host[0]], host[1], attributes=attr_first)
    check_storage(rt, len(runtimes)+1, attr['indexed_public'])
    runtimes += [rt]
    if proxy_storage:
        import calvin.runtime.north.storage
        calvin.runtime.north.storage._conf.set('global', 'storage_type', 'proxy')
        calvin.runtime.north.storage._conf.set('global', 'storage_proxy', host[0])
    _log.info("started runtime %s %s" % host)

    count = 2
    for host in hosts[1:]:
        if nbr > 3:
            # Improve likelihood of success if runtimes started with a time interval
            time.sleep(10.0)
        _log.info("starting runtime %s %s" % host)
        attr_rt = copy.deepcopy(attr_rest)
        attr_rt['indexed_public']['node_name']['name'] = u'runtime' + str(count)
        count += 1
        rt, _ = dispatch_node([host[0]], host[1], attributes=attr_rt)
        check_storage(rt, len(runtimes)+1, attr['indexed_public'])
        _log.info("started runtime %s %s" % host)
        runtimes += [rt]

    for host in hosts:
        check_storage(RT(host[1]), nbr, attr['indexed_public'])

    for host in hosts:
        request_handler.peer_setup(RT(host[1]), [h[0] for h in hosts if h != host])

    return runtimes

def setup_bluetooth(bt_master_controluri, request_handler):
    runtime = RT(bt_master_controluri)
    runtimes = []
    bt_master_id = request_handler.get_node_id(bt_master_controluri)
    data = request_handler.get_node(runtime, bt_master_id)
    if data:
        runtime.id = bt_master_id
        runtime.uri = data["uri"]
        test_peers = request_handler.get_nodes(runtime)
        test_peer2_id = test_peers[0]
        test_peer2 = request_handler.get_node(runtime, test_peer2_id)
        if test_peer2:
            rt2 = RT(test_peer2["control_uri"])
            rt2.id = test_peer2_id
            rt2.uri = test_peer2["uri"]
            runtimes.append(rt2)
        test_peer3_id = test_peers[1]
        if test_peer3_id:
            test_peer3 = request_handler.get_node(runtime, test_peer3_id)
            if test_peer3:
                rt3 = request_handler.RT(test_peer3["control_uri"])
                rt3.id = test_peer3_id
                rt3.uri = test_peer3["uri"]
                runtimes.append(rt3)
    return [runtime] + runtimes

def setup_test_type(request_handler, nbr=3, proxy_storage=False):
    control_uri = None
    ip_addr = None
    purpose = None
    bt_master_controluri = None
    test_type = None

    try:
        control_uri = os.environ["CALVIN_TEST_CTRL_URI"]
        purpose = os.environ["CALVIN_TEST_UUID"]
        test_type = "distributed"
    except KeyError:
        pass

    if not test_type:
        # Bluetooth tests assumes one master runtime with two connected peers
        # CALVIN_TEST_BT_MASTERCONTROLURI is the control uri of the master runtime
        try:
            bt_master_controluri = os.environ["CALVIN_TEST_BT_MASTERCONTROLURI"]
            _log.debug("Running Bluetooth tests")
            test_type = "bluetooth"
        except KeyError:
            pass

    if not test_type:
        try:
            ip_addr = os.environ["CALVIN_TEST_LOCALHOST"]
        except Exception:
            import socket
            # If this fails add hostname to the /etc/hosts file for 127.0.0.1
            ip_addr = socket.gethostbyname(socket.gethostname())
        test_type = "local"

    if test_type == "distributed":
        runtimes = setup_distributed(control_uri, purpose, request_handler)
    elif test_type == "bluetooth":
        runtimes = setup_bluetooth(bt_master_controluri, request_handler)
    else:
        proxy_storage = bool(int(os.environ.get("CALVIN_TESTING_PROXY_STORAGE", proxy_storage)))
        runtimes = setup_local(ip_addr, request_handler, nbr, proxy_storage)

    return test_type, runtimes

def teardown_test_type(test_type, runtimes, request_handler):
    from functools import partial
    def wait_for_it(peer):
        while True:
            try:
                request_handler.get_node_id(peer)
            except Exception:
                return True
        return False

    if test_type == "local":
        for peer in runtimes:
            request_handler.quit(peer)
            retry(10, partial(request_handler.get_node_id, peer), lambda _: True, "Failed to stop peer %r" % (peer,))
            # wait_for_it(peer)
        for p in multiprocessing.active_children():
            p.terminate()
            time.sleep(1)

<commit_msg>Improved check on if storage is up and running.<commit_after>from calvin.utilities.nodecontrol import dispatch_node
from calvin.utilities.attribute_resolver import format_index_string
from calvin.utilities import calvinconfig
from calvin.utilities import calvinlogger
from calvin.requests.request_handler import RT
import os
import time
import multiprocessing
import copy
import numbers

_log = calvinlogger.get_logger(__name__)
_conf = calvinconfig.get()

def retry(retries, function, criterion, error_msg):
    """
        Executes 'result = function()' until 'criterion(result)' evaluates to a true value.
        Raises 'Exception(error_msg)' if criterion is not fulfilled after 'retries' attempts

    """
    delay = 0.1
    retry = 0
    result = None
    while retry < retries:
        try:
            result = function()
            try:
                if criterion(result):
                    if retry > 0:
                        _log.info("Criterion %s(%s) satisfied after %d retries" %
                                    (str(criterion), str(function), retry,))
                    return result
            except Exception as e:
                _log.error("Erroneous criteria '%r" % (e, ))
                raise e
        except Exception as e:
            _log.info("Encountered exception when retrying '%s'" % (e,))
        delay = min(2, delay * 1.5); retry += 1
        time.sleep(delay)
        try:
            r = result if isinstance(result, (numbers.Number, type(None))) else len(result)
        except:
            r = None
        _log.info("Criterion still not satisfied after %d retries, result (length) %s" % (retry, r))
    _log.info("Criterion %s(%s) never satisfied, last full result %s" % (str(criterion), str(function), str(result)))
    raise Exception(error_msg)

def wait_for_tokens(request_handler, rt, actor_id, size=5, retries=10):
    """An alias for 'actual_tokens'"""
    return actual_tokens(request_handler, rt, actor_id, size, retries)

def actual_tokens(request_handler, rt, actor_id, size=5, retries=10):
    """
    Uses 'request_handler' to fetch the report from actor 'actor_id' on runtime 'rt'.
    """
    from functools import partial
    func = partial(request_handler.report, rt, actor_id)
    criterion = lambda tokens: len(tokens) >= size
    return retry(retries, func, criterion, "Not enough tokens, expected %d" % size)

def multi_report(request_handler, rt, actor_ids):
    """
    Helper uses 'request_handler' to fetch the report from actors in 'actor_ids' list on runtime(s) 'rt'.
    """
    result = []
    if isinstance(rt, (list, tuple, set)):
        args = zip(rt, actor_ids)
    else:
        args = zip([rt]*len(actor_ids), actor_ids)
    for runtime, actor_id in args:
        result.append(request_handler.report(runtime, actor_id))
    return result

def actual_tokens_multiple(request_handler, rt, actor_ids, size=5, retries=10):
    """
    Uses 'request_handler' to fetch the report from actors in 'actor_ids' list on runtime(s) 'rt'.
    """
    from functools import partial
    func = partial(multi_report, request_handler, rt, actor_ids)
    criterion = lambda tokens: sum([len(t) for t in tokens]) >= size
    return retry(retries, func, criterion, "Not enough tokens, expected %d" % size)

def destroy_app(deployer, retries=10):
    """
    Tries to destroy the app connected with deployer.
    """
    return delete_app(deployer.request_handler, deployer.runtime, deployer.app_id, retries=retries)

def deploy_app(request_handler, deployer, runtimes, retries=10):
    """
    Deploys app associated w/ deployer and then tries to verify its
    presence in registry (for all runtimes).
    """
    deployer.deploy()

    def check_application():
        for rt in runtimes:
            try:
                if request_handler.get_application(rt, deployer.app_id) is None:
                    return False
            except:
                return False
        _log.info("Application found on all peers, continuing")
        return True

    return retry(retries, check_application, lambda r: r, "Application not found on all peers")

def delete_app(request_handler, runtime, app_id, check_actor_ids=None, retries=10):
    """
    Deletes an app and then tries to verify it is actually gone.
    """
    from functools import partial

    def verify_app_gone(request_handler, runtime, app_id):
        try:
            request_handler.get_application(runtime, app_id)
            return False
        except:
            return True

    def verify_actors_gone(request_handler, runtime, actor_ids):
        responses = []
        for actor_id in actor_ids:
            responses.append(request_handler.async_get_actor(runtime, actor_id))
        gone = True
        for r in responses:
            try:
                response = request_handler.async_response(r)
                if response is not None:
                    gone = False
            except:
                pass
        return gone

    try:
        request_handler.delete_application(runtime, app_id)
    except Exception as e:
        msg = str(e.message)
        if msg.startswith("500"):
            _log.error("Delete App got 500")
        elif msg.startswith("404"):
            _log.error("Delete App got 404")
        else:
            _log.error("Delete App got unknown error %s" % str(msg))

    retry(retries, partial(verify_app_gone, request_handler, runtime, app_id), lambda r: r, "Application not deleted")
    if check_actor_ids:
        retry(retries, partial(verify_actors_gone, request_handler, runtime, check_actor_ids),
              lambda r: r, "Application actors not deleted")

def deploy_script(request_handler, name, script, runtime, retries=10):
    """
    Deploys script and then tries to verify its
    presence in registry on the runtime.
    """

    response = request_handler.deploy_application(runtime, name, script)
    app_id = response['application_id']

    def check_application():
        try:
            if request_handler.get_application(runtime, app_id) is None:
                return False
        except:
            return False
        _log.info("Application found, continuing")
        return True

    retry(retries, check_application, lambda r: r, "Application not found")
    return response

def flatten_zip(lz):
    return [] if not lz else [ lz[0][0], lz[0][1] ] + flatten_zip(lz[1:])

# Helper for 'std.CountTimer' actor
def expected_counter(n):
    return [i for i in range(1, n + 1)]

# Helper for 'std.Sum'
def expected_sum(n):
    def cumsum(l):
        s = 0
        for n in l:
            s = s + n
            yield s

    return list(cumsum(range(1, n + 1)))

def expected_tokens(request_handler, rt, actor_id, t_type):

    tokens = request_handler.report(rt, actor_id)

    if t_type == 'seq':
        return expected_counter(tokens)

    if t_type == 'sum':
        return expected_sum(tokens)

    return None

def setup_distributed(control_uri, purpose, request_handler):
    from functools import partial

    remote_node_count = 3
    test_peers = None
    runtimes = []

    runtime = RT(control_uri)
    index = {"node_name": {"organization": "com.ericsson", "purpose": purpose}}
    index_string = format_index_string(index)

    get_index = partial(request_handler.get_index, runtime, index_string)

    def criteria(peers):
        return peers and peers.get("result", None) and len(peers["result"]) >= remote_node_count

    test_peers = retry(10, get_index, criteria, "Not all nodes found")
    test_peers = test_peers["result"]

    for peer_id in test_peers:
        peer = request_handler.get_node(runtime, peer_id)
        if not peer:
            _log.warning("Runtime '%r' peer '%r' does not exist" % (runtime, peer_id, ))
            continue
        rt = RT(peer["control_uri"])
        rt.id = peer_id
        rt.uri = peer["uri"]
        runtimes.append(rt)

    return runtimes

def setup_local(ip_addr, request_handler, nbr, proxy_storage):
    def check_storage(rt, n, index):
        index_string = format_index_string(index)
        retries = 0
        while retries < 120:
            try:
                retries += 1
                peers = request_handler.get_index(rt, index_string, timeout=60)
            except Exception as e:
                try:
                    notfound = e.message.startswith("404")
                except:
                    notfound = False
                if notfound:
                    peers={'result':[]}
                else:
                    _log.info("Timed out when finding peers retrying")
                    retries += 39  # A timeout counts more we don't want to wait 60*100 seconds
                    continue
            if len(peers['result']) >= n:
                _log.info("Found %d peers (%r)", len(peers['result']), peers['result'])
                return
            _log.info("Only %d peers found (%r)", len(peers['result']), peers['result'])
            time.sleep(1)
        # No more retrying
        raise Exception("Storage check failed, could not find peers.")

    hosts = [
        ("calvinip://%s:%d" % (ip_addr, d), "http://%s:%d" % (ip_addr, d+1)) for d in range(5200, 5200 + 2 * nbr, 2)
    ]

    runtimes = []

    host = hosts[0]
    attr = {u'indexed_public': {u'node_name': {u'organization': u'com.ericsson', u'purpose': u'distributed-test'}}}
    attr_first = copy.deepcopy(attr)
    attr_first['indexed_public']['node_name']['group'] = u'first'
    attr_first['indexed_public']['node_name']['name'] = u'runtime1'
    attr_rest = copy.deepcopy(attr)
    attr_rest['indexed_public']['node_name']['group'] = u'rest'

    _log.info("starting runtime %s %s" % host)

    if proxy_storage:
        import calvin.runtime.north.storage
        calvin.runtime.north.storage._conf.set('global', 'storage_type', 'local')
    rt, _ = dispatch_node([host[0]], host[1], attributes=attr_first)
    check_storage(rt, len(runtimes)+1, attr['indexed_public'])
    runtimes += [rt]
    if proxy_storage:
        import calvin.runtime.north.storage
        calvin.runtime.north.storage._conf.set('global', 'storage_type', 'proxy')
        calvin.runtime.north.storage._conf.set('global', 'storage_proxy', host[0])
    _log.info("started runtime %s %s" % host)

    count = 2
    for host in hosts[1:]:
        if nbr > 3:
            # Improve likelihood of success if runtimes started with a time interval
            time.sleep(10.0)
        _log.info("starting runtime %s %s" % host)
        attr_rt = copy.deepcopy(attr_rest)
        attr_rt['indexed_public']['node_name']['name'] = u'runtime' + str(count)
        count += 1
        rt, _ = dispatch_node([host[0]], host[1], attributes=attr_rt)
        check_storage(rt, len(runtimes)+1, attr['indexed_public'])
        _log.info("started runtime %s %s" % host)
        runtimes += [rt]

    for host in hosts:
        check_storage(RT(host[1]), nbr, attr['indexed_public'])

    for host in hosts:
        request_handler.peer_setup(RT(host[1]), [h[0] for h in hosts if h != host])

    return runtimes

def setup_bluetooth(bt_master_controluri, request_handler):
    runtime = RT(bt_master_controluri)
    runtimes = []
    bt_master_id = request_handler.get_node_id(bt_master_controluri)
    data = request_handler.get_node(runtime, bt_master_id)
    if data:
        runtime.id = bt_master_id
        runtime.uri = data["uri"]
        test_peers = request_handler.get_nodes(runtime)
        test_peer2_id = test_peers[0]
        test_peer2 = request_handler.get_node(runtime, test_peer2_id)
        if test_peer2:
            rt2 = RT(test_peer2["control_uri"])
            rt2.id = test_peer2_id
            rt2.uri = test_peer2["uri"]
            runtimes.append(rt2)
        test_peer3_id = test_peers[1]
        if test_peer3_id:
            test_peer3 = request_handler.get_node(runtime, test_peer3_id)
            if test_peer3:
                rt3 = request_handler.RT(test_peer3["control_uri"])
                rt3.id = test_peer3_id
                rt3.uri = test_peer3["uri"]
                runtimes.append(rt3)
    return [runtime] + runtimes

def setup_test_type(request_handler, nbr=3, proxy_storage=False):
    control_uri = None
    ip_addr = None
    purpose = None
    bt_master_controluri = None
    test_type = None

    try:
        control_uri = os.environ["CALVIN_TEST_CTRL_URI"]
        purpose = os.environ["CALVIN_TEST_UUID"]
        test_type = "distributed"
    except KeyError:
        pass

    if not test_type:
        # Bluetooth tests assumes one master runtime with two connected peers
        # CALVIN_TEST_BT_MASTERCONTROLURI is the control uri of the master runtime
        try:
            bt_master_controluri = os.environ["CALVIN_TEST_BT_MASTERCONTROLURI"]
            _log.debug("Running Bluetooth tests")
            test_type = "bluetooth"
        except KeyError:
            pass

    if not test_type:
        try:
            ip_addr = os.environ["CALVIN_TEST_LOCALHOST"]
        except Exception:
            import socket
            # If this fails add hostname to the /etc/hosts file for 127.0.0.1
            ip_addr = socket.gethostbyname(socket.gethostname())
        test_type = "local"

    if test_type == "distributed":
        runtimes = setup_distributed(control_uri, purpose, request_handler)
    elif test_type == "bluetooth":
        runtimes = setup_bluetooth(bt_master_controluri, request_handler)
    else:
        proxy_storage = bool(int(os.environ.get("CALVIN_TESTING_PROXY_STORAGE", proxy_storage)))
        runtimes = setup_local(ip_addr, request_handler, nbr, proxy_storage)

    return test_type, runtimes

def teardown_test_type(test_type, runtimes, request_handler):
    from functools import partial
    def wait_for_it(peer):
        while True:
            try:
                request_handler.get_node_id(peer)
            except Exception:
                return True
        return False

    if test_type == "local":
        for peer in runtimes:
            request_handler.quit(peer)
            retry(10, partial(request_handler.get_node_id, peer), lambda _: True, "Failed to stop peer %r" % (peer,))
            # wait_for_it(peer)
        for p in multiprocessing.active_children():
            p.terminate()
            time.sleep(1)

<|endoftext|>