<commit_before>
import os
import socket
import base64
import string
import random
from collections import defaultdict
from nova.exception import ImageNotFound
from nova.api.ec2.cloud import CloudController
from sfa.util.faults import SfaAPIError
from sfa.rspecs.rspec import RSpec
from sfa.rspecs.elements.hardware_type import HardwareType
from sfa.rspecs.elements.node import Node
from sfa.rspecs.elements.sliver import Sliver
from sfa.rspecs.elements.login import Login
from sfa.rspecs.elements.disk_image import DiskImage
from sfa.rspecs.elements.services import Services
from sfa.rspecs.elements.interface import Interface
from sfa.util.xrn import Xrn
from sfa.planetlab.plxrn import PlXrn
from sfa.openstack.osxrn import OSXrn, hrn_to_os_slicename
from sfa.rspecs.version_manager import VersionManager
from sfa.openstack.security_group import SecurityGroup
from sfa.server.threadmanager import ThreadManager
from sfa.util.sfalogging import logger

def pubkeys_to_user_data(pubkeys):
    user_data = "#!/bin/bash\n\n"
    for pubkey in pubkeys:
        pubkey = pubkey.replace('\n', '')
        user_data += "echo %s >> /root/.ssh/authorized_keys" % pubkey
        user_data += "\n"
        user_data += "echo >> /root/.ssh/authorized_keys"
        user_data += "\n"
    return user_data

def instance_to_sliver(instance, slice_xrn=None):
    sliver_id = None
    if slice_xrn:
        xrn = Xrn(slice_xrn, 'slice')
        sliver_id = xrn.get_sliver_id(instance.project_id, instance.hostname, instance.id)

    sliver = Sliver({'slice_id': sliver_id,
                     'name': instance.name,
                     'type': instance.name,
                     'cpus': str(instance.vcpus),
                     'memory': str(instance.ram),
                     'storage':  str(instance.disk)})
    return sliver

def image_to_rspec_disk_image(image):
    img = DiskImage()
    img['name'] = image['name']
    img['description'] = image['name']
    img['os'] = image['name']
    img['version'] = image['name']
    return img

class OSAggregate:

    def __init__(self, driver):
        self.driver = driver

    def get_rspec(self, slice_xrn=None, version=None, options={}):
        version_manager = VersionManager()
        version = version_manager.get_version(version)
        if not slice_xrn:
            rspec_version = version_manager._get_version(version.type, version.version, 'ad')
            nodes = self.get_aggregate_nodes()
        else:
            rspec_version = version_manager._get_version(version.type, version.version, 'manifest')
            nodes = self.get_slice_nodes(slice_xrn)
        rspec = RSpec(version=rspec_version, user_options=options)
        rspec.version.add_nodes(nodes)
        return rspec.toxml()

    def get_availability_zones(self):
        # essex release
        zones = self.driver.shell.nova_manager.dns_domains.domains()

        if not zones:
            zones = ['cloud']
        else:
            zones = [zone.name for zone in zones]
        return zones

    def get_slice_nodes(self, slice_xrn):
        # update nova connection
        tenant_name = OSXrn(xrn=slice_xrn, type='slice').get_tenant_name()
        self.driver.shell.nova_manager.connect(tenant=tenant_name)

        zones = self.get_availability_zones()
        name = hrn_to_os_slicename(slice_xrn)
        instances = self.driver.shell.nova_manager.servers.findall(name=name)
        node_dict = {}
        for instance in instances:
            # determine node urn
            node_xrn = instance.metadata.get('component_id')
            if not node_xrn:
                node_xrn = OSXrn('cloud', type='node')
            else:
                node_xrn = OSXrn(xrn=node_xrn, type='node')

            if not node_xrn.urn in node_dict:
                rspec_node = Node()
                rspec_node['component_id'] = node_xrn.urn
                rspec_node['component_name'] = node_xrn.name
                rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
                rspec_node['slivers'] = []
                node_dict[node_xrn.urn] = rspec_node
            else:
                rspec_node = node_dict[node_xrn.urn]

            if instance.metadata.get('client_id'):
                rspec_node['client_id'] = instance.metadata.get('client_id')

            flavor = self.driver.shell.nova_manager.flavors.find(id=instance.flavor['id'])
            sliver = instance_to_sliver(flavor)
            rspec_node['slivers'].append(sliver)
            image = self.driver.shell.image_manager.get_images(id=instance.image['id'])
            if isinstance(image, list) and len(image) > 0:
                image = image[0]
            disk_image = image_to_rspec_disk_image(image)
            sliver['disk_image'] = [disk_image]

            # build interfaces
            rspec_node['services'] = []
            rspec_node['interfaces'] = []
            addresses = instance.addresses
            # HACK: public ips are stored in the list of private, but
            # this seems wrong. Assume pub ip is the last in the list of
            # private ips until openstack bug is fixed.
            if addresses.get('private'):
                login = Login({'authentication': 'ssh-keys',
                               'hostname': addresses.get('private')[-1]['addr'],
                               'port':'22', 'username': 'root'})
                service = Services({'login': login})
                rspec_node['services'].append(service)

            for private_ip in addresses.get('private', []):
                if_xrn = PlXrn(auth=self.driver.hrn,
                               interface='node%s:eth0' % (instance.hostId))
                interface = Interface({'component_id': if_xrn.urn})
                interface['ips'] =  [{'address': private_ip['addr'],
                                     #'netmask': private_ip['network'],
                                     'type': private_ip['version']}]
                rspec_node['interfaces'].append(interface)

            # slivers always provide the ssh service
            for public_ip in addresses.get('public', []):
                login = Login({'authentication': 'ssh-keys',
                               'hostname': public_ip['addr'],
                               'port':'22', 'username': 'root'})
                service = Services({'login': login})
                rspec_node['services'].append(service)
            rspec_nodes.append(rspec_node)
        return node_dict.values()

    def get_aggregate_nodes(self):
        zones = self.get_availability_zones()
        # available sliver/instance/vm types
        instances = self.driver.shell.nova_manager.flavors.list()
        if isinstance(instances, dict):
            instances = instances.values()
        # available images
        images = self.driver.shell.image_manager.get_images_detailed()
        disk_images  = [image_to_rspec_disk_image(img) for img in images if img['container_format'] in ['ami', 'ovf']]
        rspec_nodes = []
        for zone in zones:
            rspec_node = Node()
            xrn = OSXrn(zone, type='node')
            rspec_node['component_id'] = xrn.urn
            rspec_node['component_name'] = xrn.name
            rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
            rspec_node['exclusive'] = 'false'
            rspec_node['hardware_types'] = [HardwareType({'name': 'plos-pc'}),
                                                HardwareType({'name': 'pc'})]
            slivers = []
            for instance in instances:
                sliver = instance_to_sliver(instance)
                sliver['disk_image'] = disk_images
                slivers.append(sliver)

            rspec_node['slivers'] = slivers
            rspec_nodes.append(rspec_node)

        return rspec_nodes

    def create_tenant(self, tenant_name):
        tenants = self.driver.shell.auth_manager.tenants.findall(name=tenant_name)
        if not tenants:
            self.driver.shell.auth_manager.tenants.create(tenant_name, tenant_name)
            tenant = self.driver.shell.auth_manager.tenants.find(name=tenant_name)
        else:
            tenant = tenants[0]
        return tenant

    def create_instance_key(self, slice_hrn, user):
        slice_name = Xrn(slice_hrn).leaf
        user_name = Xrn(user['urn']).leaf
        key_name = "%s_%s" % (slice_name, user_name)
        pubkey = user['keys'][0]
        key_found = False
        existing_keys = self.driver.shell.nova_manager.keypairs.findall(name=key_name)
        for existing_key in existing_keys:
            if existing_key.public_key != pubkey:
                self.driver.shell.nova_manager.keypairs.delete(existing_key)
            elif existing_key.public_key == pubkey:
                key_found = True

        if not key_found:
            self.driver.shell.nova_manager.keypairs.create(key_name, pubkey)
        return key_name

    def create_security_group(self, slicename, fw_rules=[]):
        # use default group by default
        group_name = 'default'
        if isinstance(fw_rules, list) and fw_rules:
            # Each sliver get's its own security group.
            # Keep security group names unique by appending some random
            # characters on end.
            random_name = "".join([random.choice(string.letters+string.digits)
                                           for i in xrange(6)])
            group_name = slicename + random_name
            security_group = SecurityGroup(self.driver)
            security_group.create_security_group(group_name)
            for rule in fw_rules:
                security_group.add_rule_to_group(group_name,
                                             protocol = rule.get('protocol'),
                                             cidr_ip = rule.get('cidr_ip'),
                                             port_range = rule.get('port_range'),
                                             icmp_type_code = rule.get('icmp_type_code'))
        return group_name

    def add_rule_to_security_group(self, group_name, **kwds):
        security_group = SecurityGroup(self.driver)
        security_group.add_rule_to_group(group_name=group_name,
                                         protocol=kwds.get('protocol'),
                                         cidr_ip =kwds.get('cidr_ip'),
                                         icmp_type_code = kwds.get('icmp_type_code'))

    def run_instances(self, instance_name, tenant_name, rspec, key_name, pubkeys):
        #logger.debug('Reserving an instance: image: %s, flavor: ' \
        #            '%s, key: %s, name: %s' % \
        #            (image_id, flavor_id, key_name, slicename))

        # make sure a tenant exists for this slice
        tenant = self.create_tenant(tenant_name)

        # add the sfa admin user to this tenant and update our nova client connection
        # to use these credentials for the rest of this session. This emsures that the instances
        # we create will be assigned to the correct tenant.
        sfa_admin_user = self.driver.shell.auth_manager.users.find(name=self.driver.shell.auth_manager.opts['OS_USERNAME'])
        user_role = self.driver.shell.auth_manager.roles.find(name='user')
        admin_role = self.driver.shell.auth_manager.roles.find(name='admin')
        self.driver.shell.auth_manager.roles.add_user_role(sfa_admin_user, admin_role, tenant)
        self.driver.shell.auth_manager.roles.add_user_role(sfa_admin_user, user_role, tenant)
        self.driver.shell.nova_manager.connect(tenant=tenant.name)

        authorized_keys = "\n".join(pubkeys)
        files = {'/root/.ssh/authorized_keys': authorized_keys}
        rspec = RSpec(rspec)
        requested_instances = defaultdict(list)
        # iterate over clouds/zones/nodes
        for node in rspec.version.get_nodes_with_slivers():
            instances = node.get('slivers', [])
            if not instances:
                continue
            for instance in instances:
                try:
                    metadata = {}
                    flavor_id = self.driver.shell.nova_manager.flavors.find(name=instance['name'])
                    image = instance.get('disk_image')
                    if image and isinstance(image, list):
                        image = image[0]
                    image_id = self.driver.shell.nova_manager.images.find(name=image['name'])
                    fw_rules = instance.get('fw_rules', [])
                    group_name = self.create_security_group(instance_name, fw_rules)
                    metadata['security_groups'] = group_name
                    if node.get('component_id'):
                        metadata['component_id'] = node['component_id']
                    if node.get('client_id'):
                        metadata['client_id'] = node['client_id']
                    self.driver.shell.nova_manager.servers.create(flavor=flavor_id,
                                                            image=image_id,
                                                            key_name = key_name,
                                                            security_groups = [group_name],
                                                            files=files,
                                                            meta=metadata,
                                                            name=instance_name)
                except Exception, err:
                    logger.log_exc(err)

    def delete_instances(self, instance_name, tenant_name):

        def _delete_security_group(instance):
            security_group = instance.metadata.get('security_groups', '')
            if security_group:
                manager = SecurityGroup(self.driver)
                timeout = 10.0 # wait a maximum of 10 seconds before forcing the security group delete
                start_time = time.time()
                instance_deleted = False
                while instance_deleted == False and (time.time() - start_time) < timeout:
                    inst = self.driver.shell.nova_manager.servers.findall(id=instance.id)
                    if not inst:
                        instance_deleted = True
                manager.delete_security_group(security_group)

        thread_manager = ThreadManager()
        self.driver.shell.nova_manager.connect(tenant=tenant_name)
        instances = self.driver.shell.nova_manager.servers.findall(name=instance_name)
        for instance in instances:
            # destroy instance
            self.driver.shell.nova_manager.servers.delete(instance)
            # deleate this instance's security groups
            thread_manager.run(_delete_security_group, instance)
        return 1

    def stop_instances(self, instance_name, tenant_name):
        self.driver.shell.nova_manager.connect(tenant=tenant_name)
        instances = self.driver.shell.nova_manager.servers.findall(name=instance_name)
        for instance in instances:
            self.driver.shell.nova_manager.servers.pause(instance)
        return 1

    def update_instances(self, project_name):
        pass
<commit_msg>sleep in between polls<commit_after>
import os
import socket
import base64
import string
import random
from collections import defaultdict
from nova.exception import ImageNotFound
from nova.api.ec2.cloud import CloudController
from sfa.util.faults import SfaAPIError
from sfa.rspecs.rspec import RSpec
from sfa.rspecs.elements.hardware_type import HardwareType
from sfa.rspecs.elements.node import Node
from sfa.rspecs.elements.sliver import Sliver
from sfa.rspecs.elements.login import Login
from sfa.rspecs.elements.disk_image import DiskImage
from sfa.rspecs.elements.services import Services
from sfa.rspecs.elements.interface import Interface
from sfa.util.xrn import Xrn
from sfa.planetlab.plxrn import PlXrn
from sfa.openstack.osxrn import OSXrn, hrn_to_os_slicename
from sfa.rspecs.version_manager import VersionManager
from sfa.openstack.security_group import SecurityGroup
from sfa.server.threadmanager import ThreadManager
from sfa.util.sfalogging import logger

def pubkeys_to_user_data(pubkeys):
    user_data = "#!/bin/bash\n\n"
    for pubkey in pubkeys:
        pubkey = pubkey.replace('\n', '')
        user_data += "echo %s >> /root/.ssh/authorized_keys" % pubkey
        user_data += "\n"
        user_data += "echo >> /root/.ssh/authorized_keys"
        user_data += "\n"
    return user_data

def instance_to_sliver(instance, slice_xrn=None):
    sliver_id = None
    if slice_xrn:
        xrn = Xrn(slice_xrn, 'slice')
        sliver_id = xrn.get_sliver_id(instance.project_id, instance.hostname, instance.id)

    sliver = Sliver({'slice_id': sliver_id,
                     'name': instance.name,
                     'type': instance.name,
                     'cpus': str(instance.vcpus),
                     'memory': str(instance.ram),
                     'storage':  str(instance.disk)})
    return sliver

def image_to_rspec_disk_image(image):
    img = DiskImage()
    img['name'] = image['name']
    img['description'] = image['name']
    img['os'] = image['name']
    img['version'] = image['name']
    return img

class OSAggregate:

    def __init__(self, driver):
        self.driver = driver

    def get_rspec(self, slice_xrn=None, version=None, options={}):
        version_manager = VersionManager()
        version = version_manager.get_version(version)
        if not slice_xrn:
            rspec_version = version_manager._get_version(version.type, version.version, 'ad')
            nodes = self.get_aggregate_nodes()
        else:
            rspec_version = version_manager._get_version(version.type, version.version, 'manifest')
            nodes = self.get_slice_nodes(slice_xrn)
        rspec = RSpec(version=rspec_version, user_options=options)
        rspec.version.add_nodes(nodes)
        return rspec.toxml()

    def get_availability_zones(self):
        # essex release
        zones = self.driver.shell.nova_manager.dns_domains.domains()

        if not zones:
            zones = ['cloud']
        else:
            zones = [zone.name for zone in zones]
        return zones

    def get_slice_nodes(self, slice_xrn):
        # update nova connection
        tenant_name = OSXrn(xrn=slice_xrn, type='slice').get_tenant_name()
        self.driver.shell.nova_manager.connect(tenant=tenant_name)

        zones = self.get_availability_zones()
        name = hrn_to_os_slicename(slice_xrn)
        instances = self.driver.shell.nova_manager.servers.findall(name=name)
        node_dict = {}
        for instance in instances:
            # determine node urn
            node_xrn = instance.metadata.get('component_id')
            if not node_xrn:
                node_xrn = OSXrn('cloud', type='node')
            else:
                node_xrn = OSXrn(xrn=node_xrn, type='node')

            if not node_xrn.urn in node_dict:
                rspec_node = Node()
                rspec_node['component_id'] = node_xrn.urn
                rspec_node['component_name'] = node_xrn.name
                rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
                rspec_node['slivers'] = []
                node_dict[node_xrn.urn] = rspec_node
            else:
                rspec_node = node_dict[node_xrn.urn]

            if instance.metadata.get('client_id'):
                rspec_node['client_id'] = instance.metadata.get('client_id')

            flavor = self.driver.shell.nova_manager.flavors.find(id=instance.flavor['id'])
            sliver = instance_to_sliver(flavor)
            rspec_node['slivers'].append(sliver)
            image = self.driver.shell.image_manager.get_images(id=instance.image['id'])
            if isinstance(image, list) and len(image) > 0:
                image = image[0]
            disk_image = image_to_rspec_disk_image(image)
            sliver['disk_image'] = [disk_image]

            # build interfaces
            rspec_node['services'] = []
            rspec_node['interfaces'] = []
            addresses = instance.addresses
            # HACK: public ips are stored in the list of private, but
            # this seems wrong. Assume pub ip is the last in the list of
            # private ips until openstack bug is fixed.
            if addresses.get('private'):
                login = Login({'authentication': 'ssh-keys',
                               'hostname': addresses.get('private')[-1]['addr'],
                               'port':'22', 'username': 'root'})
                service = Services({'login': login})
                rspec_node['services'].append(service)

            for private_ip in addresses.get('private', []):
                if_xrn = PlXrn(auth=self.driver.hrn,
                               interface='node%s:eth0' % (instance.hostId))
                interface = Interface({'component_id': if_xrn.urn})
                interface['ips'] =  [{'address': private_ip['addr'],
                                     #'netmask': private_ip['network'],
                                     'type': private_ip['version']}]
                rspec_node['interfaces'].append(interface)

            # slivers always provide the ssh service
            for public_ip in addresses.get('public', []):
                login = Login({'authentication': 'ssh-keys',
                               'hostname': public_ip['addr'],
                               'port':'22', 'username': 'root'})
                service = Services({'login': login})
                rspec_node['services'].append(service)
            rspec_nodes.append(rspec_node)
        return node_dict.values()

    def get_aggregate_nodes(self):
        zones = self.get_availability_zones()
        # available sliver/instance/vm types
        instances = self.driver.shell.nova_manager.flavors.list()
        if isinstance(instances, dict):
            instances = instances.values()
        # available images
        images = self.driver.shell.image_manager.get_images_detailed()
        disk_images  = [image_to_rspec_disk_image(img) for img in images if img['container_format'] in ['ami', 'ovf']]
        rspec_nodes = []
        for zone in zones:
            rspec_node = Node()
            xrn = OSXrn(zone, type='node')
            rspec_node['component_id'] = xrn.urn
            rspec_node['component_name'] = xrn.name
            rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
            rspec_node['exclusive'] = 'false'
            rspec_node['hardware_types'] = [HardwareType({'name': 'plos-pc'}),
                                                HardwareType({'name': 'pc'})]
            slivers = []
            for instance in instances:
                sliver = instance_to_sliver(instance)
                sliver['disk_image'] = disk_images
                slivers.append(sliver)

            rspec_node['slivers'] = slivers
            rspec_nodes.append(rspec_node)

        return rspec_nodes

    def create_tenant(self, tenant_name):
        tenants = self.driver.shell.auth_manager.tenants.findall(name=tenant_name)
        if not tenants:
            self.driver.shell.auth_manager.tenants.create(tenant_name, tenant_name)
            tenant = self.driver.shell.auth_manager.tenants.find(name=tenant_name)
        else:
            tenant = tenants[0]
        return tenant

    def create_instance_key(self, slice_hrn, user):
        slice_name = Xrn(slice_hrn).leaf
        user_name = Xrn(user['urn']).leaf
        key_name = "%s_%s" % (slice_name, user_name)
        pubkey = user['keys'][0]
        key_found = False
        existing_keys = self.driver.shell.nova_manager.keypairs.findall(name=key_name)
        for existing_key in existing_keys:
            if existing_key.public_key != pubkey:
                self.driver.shell.nova_manager.keypairs.delete(existing_key)
            elif existing_key.public_key == pubkey:
                key_found = True

        if not key_found:
            self.driver.shell.nova_manager.keypairs.create(key_name, pubkey)
        return key_name

    def create_security_group(self, slicename, fw_rules=[]):
        # use default group by default
        group_name = 'default'
        if isinstance(fw_rules, list) and fw_rules:
            # Each sliver get's its own security group.
            # Keep security group names unique by appending some random
            # characters on end.
            random_name = "".join([random.choice(string.letters+string.digits)
                                           for i in xrange(6)])
            group_name = slicename + random_name
            security_group = SecurityGroup(self.driver)
            security_group.create_security_group(group_name)
            for rule in fw_rules:
                security_group.add_rule_to_group(group_name,
                                             protocol = rule.get('protocol'),
                                             cidr_ip = rule.get('cidr_ip'),
                                             port_range = rule.get('port_range'),
                                             icmp_type_code = rule.get('icmp_type_code'))
        return group_name

    def add_rule_to_security_group(self, group_name, **kwds):
        security_group = SecurityGroup(self.driver)
        security_group.add_rule_to_group(group_name=group_name,
                                         protocol=kwds.get('protocol'),
                                         cidr_ip =kwds.get('cidr_ip'),
                                         icmp_type_code = kwds.get('icmp_type_code'))

    def run_instances(self, instance_name, tenant_name, rspec, key_name, pubkeys):
        #logger.debug('Reserving an instance: image: %s, flavor: ' \
        #            '%s, key: %s, name: %s' % \
        #            (image_id, flavor_id, key_name, slicename))

        # make sure a tenant exists for this slice
        tenant = self.create_tenant(tenant_name)

        # add the sfa admin user to this tenant and update our nova client connection
        # to use these credentials for the rest of this session. This emsures that the instances
        # we create will be assigned to the correct tenant.
        sfa_admin_user = self.driver.shell.auth_manager.users.find(name=self.driver.shell.auth_manager.opts['OS_USERNAME'])
        user_role = self.driver.shell.auth_manager.roles.find(name='user')
        admin_role = self.driver.shell.auth_manager.roles.find(name='admin')
        self.driver.shell.auth_manager.roles.add_user_role(sfa_admin_user, admin_role, tenant)
        self.driver.shell.auth_manager.roles.add_user_role(sfa_admin_user, user_role, tenant)
        self.driver.shell.nova_manager.connect(tenant=tenant.name)

        authorized_keys = "\n".join(pubkeys)
        files = {'/root/.ssh/authorized_keys': authorized_keys}
        rspec = RSpec(rspec)
        requested_instances = defaultdict(list)
        # iterate over clouds/zones/nodes
        for node in rspec.version.get_nodes_with_slivers():
            instances = node.get('slivers', [])
            if not instances:
                continue
            for instance in instances:
                try:
                    metadata = {}
                    flavor_id = self.driver.shell.nova_manager.flavors.find(name=instance['name'])
                    image = instance.get('disk_image')
                    if image and isinstance(image, list):
                        image = image[0]
                    image_id = self.driver.shell.nova_manager.images.find(name=image['name'])
                    fw_rules = instance.get('fw_rules', [])
                    group_name = self.create_security_group(instance_name, fw_rules)
                    metadata['security_groups'] = group_name
                    if node.get('component_id'):
                        metadata['component_id'] = node['component_id']
                    if node.get('client_id'):
                        metadata['client_id'] = node['client_id']
                    self.driver.shell.nova_manager.servers.create(flavor=flavor_id,
                                                            image=image_id,
                                                            key_name = key_name,
                                                            security_groups = [group_name],
                                                            files=files,
                                                            meta=metadata,
                                                            name=instance_name)
                except Exception, err:
                    logger.log_exc(err)

    def delete_instances(self, instance_name, tenant_name):

        def _delete_security_group(instance):
            security_group = instance.metadata.get('security_groups', '')
            if security_group:
                manager = SecurityGroup(self.driver)
                timeout = 10.0 # wait a maximum of 10 seconds before forcing the security group delete
                start_time = time.time()
                instance_deleted = False
                while instance_deleted == False and (time.time() - start_time) < timeout:
                    inst = self.driver.shell.nova_manager.servers.findall(id=instance.id)
                    if not inst:
                        instance_deleted = True
                    time.sleep(.5)
                manager.delete_security_group(security_group)

        thread_manager = ThreadManager()
        self.driver.shell.nova_manager.connect(tenant=tenant_name)
        instances = self.driver.shell.nova_manager.servers.findall(name=instance_name)
        for instance in instances:
            # destroy instance
            self.driver.shell.nova_manager.servers.delete(instance)
            # deleate this instance's security groups
            thread_manager.run(_delete_security_group, instance)
        return 1

    def stop_instances(self, instance_name, tenant_name):
        self.driver.shell.nova_manager.connect(tenant=tenant_name)
        instances = self.driver.shell.nova_manager.servers.findall(name=instance_name)
        for instance in instances:
            self.driver.shell.nova_manager.servers.pause(instance)
        return 1

    def update_instances(self, project_name):
        pass
<|endoftext|>


<commit_before>import base64

from nova.exception import ImageNotFound
from nova.api.ec2.cloud import CloudController
from sfa.util.faults import SfaAPIError
from sfa.rspecs.rspec import RSpec
from sfa.rspecs.elements.hardware_type import HardwareType
from sfa.rspecs.elements.node import Node
from sfa.rspecs.elements.sliver import Sliver
from sfa.rspecs.elements.login import Login
from sfa.rspecs.elements.disk_image import DiskImage
from sfa.rspecs.elements.services import Services
from sfa.util.xrn import Xrn
from sfa.util.osxrn import OSXrn
from sfa.rspecs.version_manager import VersionManager
from sfa.openstack.image import Image

def instance_to_sliver(instance, slice_xrn=None):
    # should include?
    # * instance.image_ref
    # * instance.kernel_id
    # * instance.ramdisk_id
    import nova.db.sqlalchemy.models
    name=None
    type=None
    sliver_id = None
    if isinstance(instance, dict):
        # this is an isntance type dict
        name = instance['name']
        type = instance['name']
    elif isinstance(instance, nova.db.sqlalchemy.models.Instance):
        # this is an object that describes a running instance
        name = instance.display_name
        type = instance.instance_type.name
    else:
        raise SfaAPIError("instnace must be an instance_type dict or" + \
                           " a nova.db.sqlalchemy.models.Instance object")
    if slice_xrn:
        xrn = Xrn(slice_xrn, 'slice')
        sliver_id = xrn.get_sliver_id(instance.project_id, instance.hostname, instance.id)

    sliver = Sliver({'slice_id': sliver_id,
                     'name': name,
                     'type': 'plos-' + type,
                     'tags': []})
    return sliver

class OSAggregate:

    def __init__(self, driver):
        self.driver = driver

    def get_rspec(self, slice_xrn=None, version=None, options={}):
        version_manager = VersionManager()
        version = version_manager.get_version(version)
        if not slice_xrn:
            rspec_version = version_manager._get_version(version.type, version.version, 'ad')
            nodes = self.get_aggregate_nodes()
        else:
            rspec_version = version_manager._get_version(version.type, version.version, 'manifest')
            nodes = self.get_slice_nodes(slice_xrn)
        rspec = RSpec(version=rspec_version, user_options=options)
        rspec.version.add_nodes(nodes)
        return rspec.toxml()

    def get_slice_nodes(self, slice_xrn):
        image_manager = Image(self.driver)
        name = OSXrn(xrn = slice_xrn).name
        instances = self.driver.shell.db.instance_get_all_by_project(name)
        rspec_nodes = []
        for instance in instances:
            rspec_node = Node()
            xrn = OSXrn(instance.hostname, 'node')
            rspec_node['component_id'] = xrn.urn
            rspec_node['component_name'] = xrn.name
            rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
            sliver = instance_to_sliver(instance)
            disk_image = image_manager.get_disk_image(instance.image_ref)
            sliver['disk_images'] = [Image.disk_image_to_rspec_object(disk_image)]
            rspec_node['slivers'] = [sliver]
            rspec_nodes.append(rspec_node)
        return rspec_nodes

    def get_aggregate_nodes(self):

        zones = self.driver.shell.db.zone_get_all()
        if not zones:
            zones = ['cloud']
        else:
            zones = [zone.name for zone in zones]

        # available sliver/instance/vm types
        instances = self.driver.shell.db.instance_type_get_all().values()
        # available images
        image_manager = Image(self.driver)
        disk_images = image_manager.get_available_disk_images()
        disk_image_objects = [Image.disk_image_to_rspec_object(image) \
                               for image in disk_images]
        rspec_nodes = []
        for zone in zones:
            rspec_node = Node()
            xrn = OSXrn(zone, 'node')
            rspec_node['component_id'] = xrn.urn
            rspec_node['component_name'] = xrn.name
            rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
            rspec_node['exclusive'] = 'false'
            rspec_node['hardware_types'] = [HardwareType({'name': 'plos-pc'}),
                                                HardwareType({'name': 'pc'})]
            slivers = []
            for instance in instances:
                sliver = instance_to_sliver(instance)
                sliver['disk_images'] = disk_image_objects
                slivers.append(sliver)

            rspec_node['slivers'] = slivers
            rspec_nodes.append(rspec_node)

        return rspec_nodes

    def create_project(self, slicename, users, options={}):
        """
        Create the slice if it doesn't alredy exist. Create user
        accounts that don't already exist
        """
        import nova.exception.ProjectNotFound
        try:
            slice = self.driver.shell.auth_manager.get_project(slicename)
        except nova.exception.ProjectNotFound:
            # assume that the first user is the project manager
            proj_manager = Xrn(users[0]['urn']).get_leaf()
            self.driver.shell.auth_manager.create_project(slicename, proj_manager)

        for user in users:
            username = Xrn(user['urn']).get_leaf()
            try:
                self.driver.shell.auth_manager.get_user(username)
            except nova.exception.UserNotFound:
                self.driver.shell.auth_manager.create_user(username)
            self.verify_user_keys(username, user['keys'], options)

    def verify_user_keys(self, username, keys, options={}):
        """
        Add requested keys.
        """
        append = options.get('append', True)
        existing_keys = self.driver.shell.db.key_pair_get_all_by_user(username)
        existing_pub_keys = [key.public_key for key in existing_keys]
        removed_pub_keys = set(existing_pub_keys).difference(keys)
        added_pub_keys = set(keys).difference(existing_pub_keys)
        pubkeys = []
        # add new keys
        for public_key in added_pub_keys:
            key = {}
            key['user_id'] = username
            key['name'] =  username
            key['public'] = public_key
            self.driver.shell.db.key_pair_create(key)

        # remove old keys
        if not append:
            for key in existing_keys:
                if key.public_key in removed_pub_keys:
                    self.driver.shell.db.key_pair_destroy(username, key.name)

    def create_security_group(self, group_name, fw_rules=[]):
        security_group = SecurityGroup(self.driver)
        security_group.create_security_group(group_name)
        for rule in fw_rules:
            security_group.add_rule_to_group(group_name,
                                             protocol = rule.get('protocol'),
                                             cidr_ip = rule.get('cidr_ip'),
                                             port_range = rule.get('port_range'),
                                             icmp_type_code = rule.get('icmp_type_code'))

    def add_rule_to_security_group(self, group_name, **kwds):
        security_group = SecurityGroup(self.driver)
        security_group.add_rule_to_group(group_name=group_name,
                                         protocol=kwds.get('protocol'),
                                         cidr_ip =kwds.get('cidr_ip'),
                                         icmp_type_code = kwds.get('icmp_type_code'))

    def reserve_instance(self, image_id, kernel_id, ramdisk_id, \
                         instance_type, key_name, user_data):
        conn  = self.driver.euca_shell
        logger.info('Reserving an instance: image: %s, kernel: ' + \
                    '%s, ramdisk: %s, type: %s, key: %s' % \
                    (image_id, kernel_id, ramdisk_id,
                    instance_type, key_name))
        try:
            reservation = conn.run_instances(image_id=image_id,
                                             kernel_id=kernel_id,
                                             ramdisk_id=ramdisk_id,
                                             instance_type=instance_type,
                                             key_name=key_name,
                                             user_data = user_data,
                                             security_groups=group_names)
                                             #placement=zone,
                                             #min_count=min_count,
                                             #max_count=max_count,

        except EC2ResponseError, ec2RespError:
            logger.log_exc(ec2RespError)

    def run_instances(self, slicename, rspec, keyname, pubkeys):
        """
        Create the security groups and instances.
        """
        # the default image to use for instnaces that dont
        # explicitly request an image.
        # Just choose the first available image for now.
        image_manager = Image(self.driver)
        available_images = image_manager.get_available_disk_images()
        default_image = image_manager.get_disk_images()[0]
        default_ami_id = CloudController.image_ec2_id(default_image['ami']['id'])
        default_aki_id = CloudController.image_ec2_id(default_image['aki']['id'])
        default_ari_id = CloudController.image_ec2_id(default_image['ari']['id'])

        # get requested slivers
        rspec = RSpec(rspec)
        requested_instances = defaultdict(list)
        # iterate over clouds/zones/nodes
        for node in rspec.version.get_nodes_with_slivers():
            instance_types = node.get('slivers', [])
            if isinstance(instance_types, list):
                # iterate over sliver/instance types
                for instance_type in instance_types:
                    fw_rules = instance_type.get('fw_rules', [])
                    # Each sliver get's its own security group.
                    # Keep security group names unique by appending some random
                    # characters on end.
                    group_name = "_".join([slicename,
                                           instance_type['name'],
                                           base64.b64encode(os.urandom(6))])
                    self.create_security_group(group_name, fw_rules)
                    ami_id = default_ami_id
                    aki_id = default_aki_id
                    ari_id = default_ari_id
                    req_image = instance_type.get('disk_images')
                    if req_image and isinstance(req_image, list):
                        req_image_name = req_image[0]['name']
                        disk_image = image_manager.get_disk_image(name=req_image_name)
                        if disk_image:
                            ami_id = CloudController.image_ec2_id(disk_image['ami']['id'])
                            aki_id = CloudController.image_ec2_id(disk_image['aki']['id'])
                            ari_id = CloudController.image_ec2_id(disk_image['ari']['id'])
                    # start the instance
                    self.reserve_instance(image_id=ami_id,
                                          kernel_id=aki_id,
                                          ramdisk_id=ari_id,
                                          instance_type=instance_type['name'],
                                          key_name=keyname,
                                          user_data=pubkeys,
                                          group_name=group_name)

    def delete_instances(self, project_name):
        instances = self.driver.shell.db.instance_get_all_by_project(project_name)
        security_group_manager = SecurityGroup(self)
        for instance in instances:
            # deleate this instance's security groups
            for security_group in instance.security_groups:
                # dont delete the default security group
                if security_group.name != 'default':
                    security_group_manager.delete_security_group(security_group.name)
            # destroy instance
            self.driver.shell.db.instance_destroy(instance.instance_id)
        return 1

    def stop_instances(self, project_name):
        instances = self.driver.shell.db.instance_get_all_by_project(project_name)
        for instance in instances:
            self.driver.shell.db.instance_stop(instance.instance_id)
        return 1

    def update_instances(self, project_name):
        pass
<commit_msg>fix attribute error<commit_after>import base64

from nova.exception import ImageNotFound
from nova.api.ec2.cloud import CloudController
from sfa.util.faults import SfaAPIError
from sfa.rspecs.rspec import RSpec
from sfa.rspecs.elements.hardware_type import HardwareType
from sfa.rspecs.elements.node import Node
from sfa.rspecs.elements.sliver import Sliver
from sfa.rspecs.elements.login import Login
from sfa.rspecs.elements.disk_image import DiskImage
from sfa.rspecs.elements.services import Services
from sfa.util.xrn import Xrn
from sfa.util.osxrn import OSXrn
from sfa.rspecs.version_manager import VersionManager
from sfa.openstack.image import Image
from sfa.openstack.security_group import SecurityGroup

def instance_to_sliver(instance, slice_xrn=None):
    # should include?
    # * instance.image_ref
    # * instance.kernel_id
    # * instance.ramdisk_id
    import nova.db.sqlalchemy.models
    name=None
    type=None
    sliver_id = None
    if isinstance(instance, dict):
        # this is an isntance type dict
        name = instance['name']
        type = instance['name']
    elif isinstance(instance, nova.db.sqlalchemy.models.Instance):
        # this is an object that describes a running instance
        name = instance.display_name
        type = instance.instance_type.name
    else:
        raise SfaAPIError("instnace must be an instance_type dict or" + \
                           " a nova.db.sqlalchemy.models.Instance object")
    if slice_xrn:
        xrn = Xrn(slice_xrn, 'slice')
        sliver_id = xrn.get_sliver_id(instance.project_id, instance.hostname, instance.id)

    sliver = Sliver({'slice_id': sliver_id,
                     'name': name,
                     'type': 'plos-' + type,
                     'tags': []})
    return sliver

class OSAggregate:

    def __init__(self, driver):
        self.driver = driver

    def get_rspec(self, slice_xrn=None, version=None, options={}):
        version_manager = VersionManager()
        version = version_manager.get_version(version)
        if not slice_xrn:
            rspec_version = version_manager._get_version(version.type, version.version, 'ad')
            nodes = self.get_aggregate_nodes()
        else:
            rspec_version = version_manager._get_version(version.type, version.version, 'manifest')
            nodes = self.get_slice_nodes(slice_xrn)
        rspec = RSpec(version=rspec_version, user_options=options)
        rspec.version.add_nodes(nodes)
        return rspec.toxml()

    def get_slice_nodes(self, slice_xrn):
        image_manager = Image(self.driver)
        name = OSXrn(xrn = slice_xrn).name
        instances = self.driver.shell.db.instance_get_all_by_project(name)
        rspec_nodes = []
        for instance in instances:
            rspec_node = Node()
            xrn = OSXrn(instance.hostname, 'node')
            rspec_node['component_id'] = xrn.urn
            rspec_node['component_name'] = xrn.name
            rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
            sliver = instance_to_sliver(instance)
            disk_image = image_manager.get_disk_image(instance.image_ref)
            sliver['disk_images'] = [Image.disk_image_to_rspec_object(disk_image)]
            rspec_node['slivers'] = [sliver]
            rspec_nodes.append(rspec_node)
        return rspec_nodes

    def get_aggregate_nodes(self):

        zones = self.driver.shell.db.zone_get_all()
        if not zones:
            zones = ['cloud']
        else:
            zones = [zone.name for zone in zones]

        # available sliver/instance/vm types
        instances = self.driver.shell.db.instance_type_get_all().values()
        # available images
        image_manager = Image(self.driver)
        disk_images = image_manager.get_available_disk_images()
        disk_image_objects = [Image.disk_image_to_rspec_object(image) \
                               for image in disk_images]
        rspec_nodes = []
        for zone in zones:
            rspec_node = Node()
            xrn = OSXrn(zone, 'node')
            rspec_node['component_id'] = xrn.urn
            rspec_node['component_name'] = xrn.name
            rspec_node['component_manager_id'] = Xrn(self.driver.hrn, 'authority+cm').get_urn()
            rspec_node['exclusive'] = 'false'
            rspec_node['hardware_types'] = [HardwareType({'name': 'plos-pc'}),
                                                HardwareType({'name': 'pc'})]
            slivers = []
            for instance in instances:
                sliver = instance_to_sliver(instance)
                sliver['disk_images'] = disk_image_objects
                slivers.append(sliver)

            rspec_node['slivers'] = slivers
            rspec_nodes.append(rspec_node)

        return rspec_nodes

    def create_project(self, slicename, users, options={}):
        """
        Create the slice if it doesn't alredy exist. Create user
        accounts that don't already exist
        """
        import nova.exception.ProjectNotFound
        try:
            slice = self.driver.shell.auth_manager.get_project(slicename)
        except nova.exception.ProjectNotFound:
            # assume that the first user is the project manager
            proj_manager = Xrn(users[0]['urn']).get_leaf()
            self.driver.shell.auth_manager.create_project(slicename, proj_manager)

        for user in users:
            username = Xrn(user['urn']).get_leaf()
            try:
                self.driver.shell.auth_manager.get_user(username)
            except nova.exception.UserNotFound:
                self.driver.shell.auth_manager.create_user(username)
            self.verify_user_keys(username, user['keys'], options)

    def verify_user_keys(self, username, keys, options={}):
        """
        Add requested keys.
        """
        append = options.get('append', True)
        existing_keys = self.driver.shell.db.key_pair_get_all_by_user(username)
        existing_pub_keys = [key.public_key for key in existing_keys]
        removed_pub_keys = set(existing_pub_keys).difference(keys)
        added_pub_keys = set(keys).difference(existing_pub_keys)
        pubkeys = []
        # add new keys
        for public_key in added_pub_keys:
            key = {}
            key['user_id'] = username
            key['name'] =  username
            key['public'] = public_key
            self.driver.shell.db.key_pair_create(key)

        # remove old keys
        if not append:
            for key in existing_keys:
                if key.public_key in removed_pub_keys:
                    self.driver.shell.db.key_pair_destroy(username, key.name)

    def create_security_group(self, group_name, fw_rules=[]):
        security_group = SecurityGroup(self.driver)
        security_group.create_security_group(group_name)
        for rule in fw_rules:
            security_group.add_rule_to_group(group_name,
                                             protocol = rule.get('protocol'),
                                             cidr_ip = rule.get('cidr_ip'),
                                             port_range = rule.get('port_range'),
                                             icmp_type_code = rule.get('icmp_type_code'))

    def add_rule_to_security_group(self, group_name, **kwds):
        security_group = SecurityGroup(self.driver)
        security_group.add_rule_to_group(group_name=group_name,
                                         protocol=kwds.get('protocol'),
                                         cidr_ip =kwds.get('cidr_ip'),
                                         icmp_type_code = kwds.get('icmp_type_code'))

    def reserve_instance(self, image_id, kernel_id, ramdisk_id, \
                         instance_type, key_name, user_data):
        conn  = self.driver.euca_shell
        logger.info('Reserving an instance: image: %s, kernel: ' + \
                    '%s, ramdisk: %s, type: %s, key: %s' % \
                    (image_id, kernel_id, ramdisk_id,
                    instance_type, key_name))
        try:
            reservation = conn.run_instances(image_id=image_id,
                                             kernel_id=kernel_id,
                                             ramdisk_id=ramdisk_id,
                                             instance_type=instance_type,
                                             key_name=key_name,
                                             user_data = user_data,
                                             security_groups=group_names)
                                             #placement=zone,
                                             #min_count=min_count,
                                             #max_count=max_count,

        except EC2ResponseError, ec2RespError:
            logger.log_exc(ec2RespError)

    def run_instances(self, slicename, rspec, keyname, pubkeys):
        """
        Create the security groups and instances.
        """
        # the default image to use for instnaces that dont
        # explicitly request an image.
        # Just choose the first available image for now.
        image_manager = Image(self.driver)
        available_images = image_manager.get_available_disk_images()
        default_image = image_manager.get_disk_images()[0]
        default_ami_id = CloudController.image_ec2_id(default_image['ami']['id'])
        default_aki_id = CloudController.image_ec2_id(default_image['aki']['id'])
        default_ari_id = CloudController.image_ec2_id(default_image['ari']['id'])

        # get requested slivers
        rspec = RSpec(rspec)
        requested_instances = defaultdict(list)
        # iterate over clouds/zones/nodes
        for node in rspec.version.get_nodes_with_slivers():
            instance_types = node.get('slivers', [])
            if isinstance(instance_types, list):
                # iterate over sliver/instance types
                for instance_type in instance_types:
                    fw_rules = instance_type.get('fw_rules', [])
                    # Each sliver get's its own security group.
                    # Keep security group names unique by appending some random
                    # characters on end.
                    group_name = "_".join([slicename,
                                           instance_type['name'],
                                           base64.b64encode(os.urandom(6))])
                    self.create_security_group(group_name, fw_rules)
                    ami_id = default_ami_id
                    aki_id = default_aki_id
                    ari_id = default_ari_id
                    req_image = instance_type.get('disk_images')
                    if req_image and isinstance(req_image, list):
                        req_image_name = req_image[0]['name']
                        disk_image = image_manager.get_disk_image(name=req_image_name)
                        if disk_image:
                            ami_id = CloudController.image_ec2_id(disk_image['ami']['id'])
                            aki_id = CloudController.image_ec2_id(disk_image['aki']['id'])
                            ari_id = CloudController.image_ec2_id(disk_image['ari']['id'])
                    # start the instance
                    self.reserve_instance(image_id=ami_id,
                                          kernel_id=aki_id,
                                          ramdisk_id=ari_id,
                                          instance_type=instance_type['name'],
                                          key_name=keyname,
                                          user_data=pubkeys,
                                          group_name=group_name)

    def delete_instances(self, project_name):
        instances = self.driver.shell.db.instance_get_all_by_project(project_name)
        security_group_manager = SecurityGroup(self)
        for instance in instances:
            # deleate this instance's security groups
            for security_group in instance.security_groups:
                # dont delete the default security group
                if security_group.name != 'default':
                    security_group_manager.delete_security_group(security_group.name)
            # destroy instance
            self.driver.shell.db.instance_destroy(instance.id)
        return 1

    def stop_instances(self, project_name):
        instances = self.driver.shell.db.instance_get_all_by_project(project_name)
        for instance in instances:
            self.driver.shell.db.instance_stop(instance.id)
        return 1

    def update_instances(self, project_name):
        pass
<|endoftext|>


<commit_before>import inspect
import time
import itertools
import asyncio
import signal
import sys
import numbers

import numpy as np
import scipy as sp
import pandas as pd
import networkx as nx
import h5py

from pycontrol.instruments.instrument import Instrument
from pycontrol.stream import DataStream, DataAxis, DataStreamDescriptor, InputConnector, OutputConnector
from pycontrol.filters.plot import Plotter
from pycontrol.logging import logger

class Parameter(object):
    """ Encapsulates the information for an experiment parameter"""

    def __init__(self, name=None, unit=None, default=None):
        self.name     = name
        self._value   = default
        self.unit     = unit
        self.default  = default
        self.method   = None

        # Hooks to be called before or after updating a sweep parameter
        self.pre_push_hooks = []
        self.post_push_hooks = []

    def add_pre_push_hook(self, hook):
        self.pre_push_hooks.append(hook)

    def add_post_push_hook(self, hook):
        self.post_push_hooks.append(hook)

    @property
    def value(self):
        return self._value

    @value.setter
    def value(self, value):
        self._value = value

    def __str__(self):
        result = ""
        result += "%s" % str(self.value)
        if self.unit:
            result += " %s" % self.unit
        return result

    def __repr__(self):
        result = "<Parameter(name='%s'" % self.name
        result += ",value=%s" % repr(self.value)
        if self.unit:
            result += ",unit='%s'" % self.unit
        return result + ")>"

    def assign_method(self, method):
        logger.debug("Setting method of Parameter %s to %s" % (self.name, str(method)) )
        self.method = method

    def push(self):
        if self.method is not None:
            # logger.debug("Calling pre_push_hooks of Parameter %s with value %s" % (self.name, self._value) )
            for pph in self.pre_push_hooks:
                pph()
            # logger.debug("Calling method of Parameter %s with value %s" % (self.name, self._value) )
            self.method(self._value)
            # logger.debug("Calling post_push_hooks of Parameter %s with value %s" % (self.name, self._value) )
            for pph in self.post_push_hooks:
                pph()

class FloatParameter(Parameter):

    @property
    def value(self):
        return self._value

    @value.setter
    def value(self, value):
        try:
            self._value = float(value)
        except ValueError:
            raise ValueError("FloatParameter given non-float value of "
                             "type '%s'" % type(value))

    def __repr__(self):
        result = super(FloatParameter, self).__repr__()
        return result.replace("<Parameter", "<FloatParameter", 1)

class IntParameter(Parameter):

    @property
    def value(self):
        return self._value

    @value.setter
    def value(self, value):
        try:
            self._value = int(value)
        except ValueError:
            raise ValueError("IntParameter given non-int value of "
                             "type '%s'" % type(value))

    def __repr__(self):
        result = super(IntParameter, self).__repr__()
        return result.replace("<Parameter", "<IntParameter", 1)

class SweptParameter(object):
    """Data structure for a swept Parameters, contains the Parameter
    object rather than subclassing it since we just need to keep track
    of some values"""
    def __init__(self, parameter, values):
        self.parameter = parameter
        self.associated_axes = []
        self.update_values(values)
        self.push = self.parameter.push

    def update_values(self, values):
        self.values = values
        self.length = len(values)
        for axis in self.associated_axes:
            axis.points = self.values

    @property
    def value(self):
        return self.parameter.value
    @value.setter
    def value(self, value):
        self.parameter.value = value

    def __repr__(self):
        return "<SweptParameter: {}>".format(self.parameter.name)

class SweptParameterGroup(object):
    """For unstructured (meshed) coordinate tuples. The actual values
    are stored locally as _values, and we acces each tuples by indexing
    into that array."""
    def __init__(self, parameters, values):
        self.parameters = parameters
        self.associated_axes = []
        self.update_values(values)

    def push(self):
        # Values here will just be the index
        for p in self.parameters:
            p.push()

    def update_values(self, values):
        self._values = values
        self.length = len(values)
        self.values = list(range(self.length)) # Dummy index list for sweeper
        for axis in self.associated_axes:
            axis.points = self._values

    @property
    def value(self):
        return [p.value for p in self.parameters]
    @value.setter
    def value(self, index):
        for i, p in enumerate(self.parameters):
            p.value = self._values[index][i]

    def __repr__(self):
        return "<SweptParameter: {}>".format([p.name for p in self.parameters])

class ExperimentGraph(object):
    def __init__(self, edges, loop):
        self.dag = None
        self.edges = []
        self.loop = loop
        self.create_graph(edges)

    def dfs_edges(self):
        # Edge depth-first traversal of the graph

        # Find the input nodes
        input_nodes = [n for n in self.dag.nodes() if self.dag.in_degree(n) == 0]
        logger.debug("Input nodes for DFS are '%s'", input_nodes)

        dfs_edge_iters  = [nx.edge_dfs(self.dag, input_node) for input_node in input_nodes]
        processed_edges = [] # Keep track of what we've initialized

        for ei in dfs_edge_iters:
            for edge in ei:
                if edge not in processed_edges:
                    processed_edges.append(edge)
                    yield edge

    def create_graph(self, edges):
        dag = nx.DiGraph()
        self.edges = []
        for edge in edges:
            obj = DataStream(name="{}_TO_{}".format(edge[0].name, edge[1].name),
                             loop=self.loop)
            edge[0].add_output_stream(obj)
            edge[1].add_input_stream(obj)
            self.edges.append(obj)
            dag.add_edge(edge[0].parent, edge[1].parent, object=obj)

        self.dag = dag

class MetaExperiment(type):
    """Meta class to bake the instrument objects into a class description
    """

    def __init__(self, name, bases, dct):
        type.__init__(self, name, bases, dct)
        logger.debug("Adding controls to %s", name)
        self._parameters        = {}
        self._instruments       = {}
        self._constants         = {}

        # Beware, passing objects won't work at parse time
        self._output_connectors = []

        for k,v in dct.items():
            if isinstance(v, Instrument):
                logger.debug("Found '%s' instrument", k)
                self._instruments[k] = v
            elif isinstance(v, Parameter):
                logger.debug("Found '%s' parameter", k)
                if v.name is None:
                    v.name = k
                self._parameters[k] = v
            elif isinstance(v, OutputConnector):
                logger.debug("Found '%s' output connector.", k)
                self._output_connectors.append(k)
            elif isinstance(v, numbers.Number) or isinstance(v, str):
                self._constants[k] = v
                # Keep track of numerical parameters

class Experiment(metaclass=MetaExperiment):
    """The measurement loop to be run for each set of sweep parameters."""
    def __init__(self):
        super(Experiment, self).__init__()

        # Iterable that yields sweep values
        self._sweep_generator = None

        # Container for patameters that will be swept
        self._swept_parameters = []

        # This holds the experiment graph
        self.graph = None

        # This holds a reference to a bokeh-server instance
        # for plotting, if there is one.
        self.bokeh_server = None

        # Also keep references to all of the plot filters
        self.plotters = []

        # Things we can't metaclass
        self.output_connectors = {}
        for oc in self._output_connectors:
            a = OutputConnector(name=oc, parent=self)
            a.parent = self
            self.output_connectors[oc] = a
            setattr(self, oc, a)

        # Create the asyncio measurement loop
        self.loop = asyncio.get_event_loop()
        self.loop.set_debug(True)

        # Run the stream init
        self.init_streams()
        self.update_descriptors()

    def set_graph(self, edges):
        unique_nodes = []
        for eb, ee in edges:
            if eb.parent not in unique_nodes:
                unique_nodes.append(eb.parent)
            if ee.parent not in unique_nodes:
                unique_nodes.append(ee.parent)
        self.nodes = unique_nodes
        self.graph = ExperimentGraph(edges, self.loop)
        self.update_descriptors()

    def init_streams(self):
        """Establish the base descriptors for any internal data streams and connectors."""
        pass

    def init_instruments(self):
        """Gets run before a sweep starts"""
        pass

    def shutdown_instruments(self):
        """Gets run after a sweep ends, or when the program is terminated."""
        pass

    async def run(self):
        """This is the inner measurement loop, which is the smallest unit that
        is repeated across various sweep variables. For more complicated run control
        than can be provided by the automatic sweeping, the full experimental
        operation should be defined here"""
        pass

    def run_loop(self):
        """This runs the asyncio main loop."""
        # Set any static parameters
        for p in self._parameters.values():
            p.push()
        tasks = [n.run() for n in self.nodes]
        self.loop.run_until_complete(asyncio.wait(tasks))

    def reset(self):
        for edge in self.graph.edges:
            edge.reset()
        self.update_descriptors()
        self.generate_sweep()

    def update_descriptors(self):
        logger.debug("Starting descriptor update in experiment.")
        for oc in self.output_connectors.values():
            for k in dir(self):
                v = getattr(self,k)
                if isinstance(v, Parameter):
                    oc.descriptor.add_param(k, v.value)
                    if v.unit is not None:
                        oc.descriptor.add_param('unit_'+k, v.unit)
                if isinstance(v, numbers.Number) or isinstance(v, str):
                    oc.descriptor.add_param(k, v)
            oc.update_descriptors()
        # TODO: have this push any changes to JSON file
        # if we're using Quince to define connections.

    async def sweep(self):
        # Set any static parameters
        for p in self._parameters.values():
            p.push()

        # Keep track of the previous values
        logger.debug("Waiting for filters.")
        await asyncio.sleep(1.0)
        last_param_values = None
        logger.debug("Starting experiment sweep.")
        for param_values in self._sweep_generator:

            # Update the parameter values. Unles set and push if there has been a change
            # in the value from the previous iteration.
            for i, sp in enumerate(self._swept_parameters):
                if last_param_values is None or param_values[i] != last_param_values[i]:
                    sp.value = param_values[i]
                    sp.push()

            # update previous values
            last_param_values = param_values

            # Run the procedure
            await self.run()

    def run_sweeps(self):
        # Go and find any plotters and keep track of them.
        # Launch the bokeh-server if necessary.
        self.plotters = [n for n in self.nodes if isinstance(n, Plotter)]

        if len(self.plotters) > 0:

            from .plotting import BokehServerThread
            from bokeh.client import push_session
            from bokeh.plotting import hplot
            from bokeh.io import curdoc, curstate
            from bokeh.util.session_id import generate_session_id
            from bokeh.document import Document
            from bokeh.models.widgets import Panel, Tabs

            bokeh_thread = BokehServerThread(notebook=notebook)
            bokeh_thread.start()

            #On some systems there is a possibility we try to `push_session` before the
            #the server on the BokehServerThread has started. Wait a second, here.
            time.sleep(1)

            tabs = True # Tabs seem a bit sluggish in jupyter notebooks...
            if tabs:
                h = Tabs(tabs=[Panel(child=p, title=p.name) for p in self.plotters])
            else:
                h = hplot(*[p.figure for p in self.plotters])

            curdoc().clear()
            sid = generate_session_id()
            doc = Document()
            doc.add_root(h)
            session = push_session(doc, session_id=sid)

            if notebook:
                from bokeh.embed import autoload_server, components
                from bokeh.io import output_notebook
                from IPython.display import display, HTML

                output_notebook()
                script = autoload_server(model=None, session_id=sid)
                html = \
                        """
                        <html>
                        <head></head>
                        <body>
                        {}
                        </body>
                        </html>
                        """.format(script)
                display(HTML(html))
            else:
                session.show(doc)

        def shutdown():
            if len(self.plotters) > 0:
                time.sleep(0.5)
                bokeh_thread.join()
            self.shutdown_instruments()

        def catch_ctrl_c(signum, frame):
            logger.info("Caught SIGINT. Shutting down.")
            shutdown()
            sys.exit(0)

        signal.signal(signal.SIGINT, catch_ctrl_c)

        # We want to wait for the sweep method above,
        # not the experiment's run method, so replace this
        # in the list of tasks.
        other_nodes = self.nodes[:]
        other_nodes.remove(self)
        tasks = [n.run() for n in other_nodes]
        tasks.append(self.sweep())
        self.loop.run_until_complete(asyncio.wait(tasks))

    def add_sweep(self, param, sweep_list):
        """Add a good-old-fasioned one-variable sweep."""
        p = SweptParameter(param, sweep_list)
        self._swept_parameters.append(p)
        self.generate_sweep()
        for oc in self.output_connectors.values():
            ax = DataAxis(param.name, sweep_list)
            logger.debug("Adding sweep axis %s to connector %s.", ax, oc.name)
            oc.descriptor.add_axis(ax)
            p.associated_axes.append(ax)
        self.update_descriptors()
        param.value = sweep_list[0]
        return p

    def add_unstructured_sweep(self, parameters, coords):
        p = SweptParameterGroup(parameters, coords)
        self._swept_parameters.append(p)
        self.generate_sweep()
        for oc in self.output_connectors.values():
            ax = DataAxis("Unstructured", coords, unstructured=True, coord_names=[p.name for p in parameters])
            logger.debug("Adding unstructred sweep axis %s to connector %s.", ax, oc.name)
            oc.descriptor.add_axis(ax)
            p.associated_axes.append(ax)
        self.update_descriptors()
        for pp, c in zip(parameters, coords[0]):
            pp.value = c
        return p

    def generate_sweep(self):
        self._sweep_generator = itertools.product(*[sp.values for sp in self._swept_parameters])
<commit_msg>Faster way to go through the constants list<commit_after>import inspect
import time
import itertools
import asyncio
import signal
import sys
import numbers

import numpy as np
import scipy as sp
import pandas as pd
import networkx as nx
import h5py

from pycontrol.instruments.instrument import Instrument
from pycontrol.stream import DataStream, DataAxis, DataStreamDescriptor, InputConnector, OutputConnector
from pycontrol.filters.plot import Plotter
from pycontrol.logging import logger

class Parameter(object):
    """ Encapsulates the information for an experiment parameter"""

    def __init__(self, name=None, unit=None, default=None):
        self.name     = name
        self._value   = default
        self.unit     = unit
        self.default  = default
        self.method   = None

        # Hooks to be called before or after updating a sweep parameter
        self.pre_push_hooks = []
        self.post_push_hooks = []

    def add_pre_push_hook(self, hook):
        self.pre_push_hooks.append(hook)

    def add_post_push_hook(self, hook):
        self.post_push_hooks.append(hook)

    @property
    def value(self):
        return self._value

    @value.setter
    def value(self, value):
        self._value = value

    def __str__(self):
        result = ""
        result += "%s" % str(self.value)
        if self.unit:
            result += " %s" % self.unit
        return result

    def __repr__(self):
        result = "<Parameter(name='%s'" % self.name
        result += ",value=%s" % repr(self.value)
        if self.unit:
            result += ",unit='%s'" % self.unit
        return result + ")>"

    def assign_method(self, method):
        logger.debug("Setting method of Parameter %s to %s" % (self.name, str(method)) )
        self.method = method

    def push(self):
        if self.method is not None:
            # logger.debug("Calling pre_push_hooks of Parameter %s with value %s" % (self.name, self._value) )
            for pph in self.pre_push_hooks:
                pph()
            # logger.debug("Calling method of Parameter %s with value %s" % (self.name, self._value) )
            self.method(self._value)
            # logger.debug("Calling post_push_hooks of Parameter %s with value %s" % (self.name, self._value) )
            for pph in self.post_push_hooks:
                pph()

class FloatParameter(Parameter):

    @property
    def value(self):
        return self._value

    @value.setter
    def value(self, value):
        try:
            self._value = float(value)
        except ValueError:
            raise ValueError("FloatParameter given non-float value of "
                             "type '%s'" % type(value))

    def __repr__(self):
        result = super(FloatParameter, self).__repr__()
        return result.replace("<Parameter", "<FloatParameter", 1)

class IntParameter(Parameter):

    @property
    def value(self):
        return self._value

    @value.setter
    def value(self, value):
        try:
            self._value = int(value)
        except ValueError:
            raise ValueError("IntParameter given non-int value of "
                             "type '%s'" % type(value))

    def __repr__(self):
        result = super(IntParameter, self).__repr__()
        return result.replace("<Parameter", "<IntParameter", 1)

class SweptParameter(object):
    """Data structure for a swept Parameters, contains the Parameter
    object rather than subclassing it since we just need to keep track
    of some values"""
    def __init__(self, parameter, values):
        self.parameter = parameter
        self.associated_axes = []
        self.update_values(values)
        self.push = self.parameter.push

    def update_values(self, values):
        self.values = values
        self.length = len(values)
        for axis in self.associated_axes:
            axis.points = self.values

    @property
    def value(self):
        return self.parameter.value
    @value.setter
    def value(self, value):
        self.parameter.value = value

    def __repr__(self):
        return "<SweptParameter: {}>".format(self.parameter.name)

class SweptParameterGroup(object):
    """For unstructured (meshed) coordinate tuples. The actual values
    are stored locally as _values, and we acces each tuples by indexing
    into that array."""
    def __init__(self, parameters, values):
        self.parameters = parameters
        self.associated_axes = []
        self.update_values(values)

    def push(self):
        # Values here will just be the index
        for p in self.parameters:
            p.push()

    def update_values(self, values):
        self._values = values
        self.length = len(values)
        self.values = list(range(self.length)) # Dummy index list for sweeper
        for axis in self.associated_axes:
            axis.points = self._values

    @property
    def value(self):
        return [p.value for p in self.parameters]
    @value.setter
    def value(self, index):
        for i, p in enumerate(self.parameters):
            p.value = self._values[index][i]

    def __repr__(self):
        return "<SweptParameter: {}>".format([p.name for p in self.parameters])

class ExperimentGraph(object):
    def __init__(self, edges, loop):
        self.dag = None
        self.edges = []
        self.loop = loop
        self.create_graph(edges)

    def dfs_edges(self):
        # Edge depth-first traversal of the graph

        # Find the input nodes
        input_nodes = [n for n in self.dag.nodes() if self.dag.in_degree(n) == 0]
        logger.debug("Input nodes for DFS are '%s'", input_nodes)

        dfs_edge_iters  = [nx.edge_dfs(self.dag, input_node) for input_node in input_nodes]
        processed_edges = [] # Keep track of what we've initialized

        for ei in dfs_edge_iters:
            for edge in ei:
                if edge not in processed_edges:
                    processed_edges.append(edge)
                    yield edge

    def create_graph(self, edges):
        dag = nx.DiGraph()
        self.edges = []
        for edge in edges:
            obj = DataStream(name="{}_TO_{}".format(edge[0].name, edge[1].name),
                             loop=self.loop)
            edge[0].add_output_stream(obj)
            edge[1].add_input_stream(obj)
            self.edges.append(obj)
            dag.add_edge(edge[0].parent, edge[1].parent, object=obj)

        self.dag = dag

class MetaExperiment(type):
    """Meta class to bake the instrument objects into a class description
    """

    def __init__(self, name, bases, dct):
        type.__init__(self, name, bases, dct)
        logger.debug("Adding controls to %s", name)
        self._parameters        = {}
        self._instruments       = {}
        self._constants         = {}

        # Beware, passing objects won't work at parse time
        self._output_connectors = []

        for k,v in dct.items():
            if isinstance(v, Instrument):
                logger.debug("Found '%s' instrument", k)
                self._instruments[k] = v
            elif isinstance(v, Parameter):
                logger.debug("Found '%s' parameter", k)
                if v.name is None:
                    v.name = k
                self._parameters[k] = v
            elif isinstance(v, OutputConnector):
                logger.debug("Found '%s' output connector.", k)
                self._output_connectors.append(k)
            elif isinstance(v, numbers.Number) or isinstance(v, str):
                self._constants[k] = v
                # Keep track of numerical parameters

class Experiment(metaclass=MetaExperiment):
    """The measurement loop to be run for each set of sweep parameters."""
    def __init__(self):
        super(Experiment, self).__init__()

        # Iterable that yields sweep values
        self._sweep_generator = None

        # Container for patameters that will be swept
        self._swept_parameters = []

        # This holds the experiment graph
        self.graph = None

        # This holds a reference to a bokeh-server instance
        # for plotting, if there is one.
        self.bokeh_server = None

        # Also keep references to all of the plot filters
        self.plotters = []

        # Things we can't metaclass
        self.output_connectors = {}
        for oc in self._output_connectors:
            a = OutputConnector(name=oc, parent=self)
            a.parent = self
            self.output_connectors[oc] = a
            setattr(self, oc, a)

        # Create the asyncio measurement loop
        self.loop = asyncio.get_event_loop()
        self.loop.set_debug(True)

        # Run the stream init
        self.init_streams()
        self.update_descriptors()

    def set_graph(self, edges):
        unique_nodes = []
        for eb, ee in edges:
            if eb.parent not in unique_nodes:
                unique_nodes.append(eb.parent)
            if ee.parent not in unique_nodes:
                unique_nodes.append(ee.parent)
        self.nodes = unique_nodes
        self.graph = ExperimentGraph(edges, self.loop)
        self.update_descriptors()

    def init_streams(self):
        """Establish the base descriptors for any internal data streams and connectors."""
        pass

    def init_instruments(self):
        """Gets run before a sweep starts"""
        pass

    def shutdown_instruments(self):
        """Gets run after a sweep ends, or when the program is terminated."""
        pass

    async def run(self):
        """This is the inner measurement loop, which is the smallest unit that
        is repeated across various sweep variables. For more complicated run control
        than can be provided by the automatic sweeping, the full experimental
        operation should be defined here"""
        pass

    def run_loop(self):
        """This runs the asyncio main loop."""
        # Set any static parameters
        for p in self._parameters.values():
            p.push()
        tasks = [n.run() for n in self.nodes]
        self.loop.run_until_complete(asyncio.wait(tasks))

    def reset(self):
        for edge in self.graph.edges:
            edge.reset()
        self.update_descriptors()
        self.generate_sweep()

    def update_descriptors(self):
        logger.debug("Starting descriptor update in experiment.")
        for oc in self.output_connectors.values():
            for k,v in self._parameters.items():
                oc.descriptor.add_param(k, v.value)
                if v.unit is not None:
                    oc.descriptor.add_param('unit_'+k, v.unit)
            for k in self._constants.keys():
                if hasattr(self,k):
                    v = getattr(self,k)
                    oc.descriptor.add_param(k, v)
            oc.update_descriptors()
        # TODO: have this push any changes to JSON file
        # if we're using Quince to define connections.

    async def sweep(self):
        # Set any static parameters
        for p in self._parameters.values():
            p.push()

        # Keep track of the previous values
        logger.debug("Waiting for filters.")
        await asyncio.sleep(1.0)
        last_param_values = None
        logger.debug("Starting experiment sweep.")
        for param_values in self._sweep_generator:

            # Update the parameter values. Unles set and push if there has been a change
            # in the value from the previous iteration.
            for i, sp in enumerate(self._swept_parameters):
                if last_param_values is None or param_values[i] != last_param_values[i]:
                    sp.value = param_values[i]
                    sp.push()

            # update previous values
            last_param_values = param_values

            # Run the procedure
            await self.run()

    def run_sweeps(self):
        # Go and find any plotters and keep track of them.
        # Launch the bokeh-server if necessary.
        self.plotters = [n for n in self.nodes if isinstance(n, Plotter)]

        if len(self.plotters) > 0:

            from .plotting import BokehServerThread
            from bokeh.client import push_session
            from bokeh.plotting import hplot
            from bokeh.io import curdoc, curstate
            from bokeh.util.session_id import generate_session_id
            from bokeh.document import Document
            from bokeh.models.widgets import Panel, Tabs

            bokeh_thread = BokehServerThread(notebook=notebook)
            bokeh_thread.start()

            #On some systems there is a possibility we try to `push_session` before the
            #the server on the BokehServerThread has started. Wait a second, here.
            time.sleep(1)

            tabs = True # Tabs seem a bit sluggish in jupyter notebooks...
            if tabs:
                h = Tabs(tabs=[Panel(child=p, title=p.name) for p in self.plotters])
            else:
                h = hplot(*[p.figure for p in self.plotters])

            curdoc().clear()
            sid = generate_session_id()
            doc = Document()
            doc.add_root(h)
            session = push_session(doc, session_id=sid)

            if notebook:
                from bokeh.embed import autoload_server, components
                from bokeh.io import output_notebook
                from IPython.display import display, HTML

                output_notebook()
                script = autoload_server(model=None, session_id=sid)
                html = \
                        """
                        <html>
                        <head></head>
                        <body>
                        {}
                        </body>
                        </html>
                        """.format(script)
                display(HTML(html))
            else:
                session.show(doc)

        def shutdown():
            if len(self.plotters) > 0:
                time.sleep(0.5)
                bokeh_thread.join()
            self.shutdown_instruments()

        def catch_ctrl_c(signum, frame):
            logger.info("Caught SIGINT. Shutting down.")
            shutdown()
            sys.exit(0)

        signal.signal(signal.SIGINT, catch_ctrl_c)

        # We want to wait for the sweep method above,
        # not the experiment's run method, so replace this
        # in the list of tasks.
        other_nodes = self.nodes[:]
        other_nodes.remove(self)
        tasks = [n.run() for n in other_nodes]
        tasks.append(self.sweep())
        self.loop.run_until_complete(asyncio.wait(tasks))

    def add_sweep(self, param, sweep_list):
        """Add a good-old-fasioned one-variable sweep."""
        p = SweptParameter(param, sweep_list)
        self._swept_parameters.append(p)
        self.generate_sweep()
        for oc in self.output_connectors.values():
            ax = DataAxis(param.name, sweep_list)
            logger.debug("Adding sweep axis %s to connector %s.", ax, oc.name)
            oc.descriptor.add_axis(ax)
            p.associated_axes.append(ax)
        self.update_descriptors()
        param.value = sweep_list[0]
        return p

    def add_unstructured_sweep(self, parameters, coords):
        p = SweptParameterGroup(parameters, coords)
        self._swept_parameters.append(p)
        self.generate_sweep()
        for oc in self.output_connectors.values():
            ax = DataAxis("Unstructured", coords, unstructured=True, coord_names=[p.name for p in parameters])
            logger.debug("Adding unstructred sweep axis %s to connector %s.", ax, oc.name)
            oc.descriptor.add_axis(ax)
            p.associated_axes.append(ax)
        self.update_descriptors()
        for pp, c in zip(parameters, coords[0]):
            pp.value = c
        return p

    def generate_sweep(self):
        self._sweep_generator = itertools.product(*[sp.values for sp in self._swept_parameters])
<|endoftext|>