<commit_before>import logging
import time
import uuid
import os
import tempfile
import yaml

from botocore.exceptions import WaiterError, ClientError
import boto3.session
from avocado.utils import runtime as avocado_runtime
from textwrap import dedent
from threading import Thread

import cluster
import ec2_client
from sdcm.utils import retrying

from . import wait

logger = logging.getLogger(__name__)

INSTANCE_PROVISION_ON_DEMAND = 'on_demand'
INSTANCE_PROVISION_SPOT_FLEET = 'spot_fleet'
INSTANCE_PROVISION_SPOT_LOW_PRICE = 'spot_low_price'
INSTANCE_PROVISION_SPOT_DURATION = 'spot_duration'
SPOT_CNT_LIMIT = 20
SPOT_FLEET_LIMIT = 50


def _prepend_user_prefix(user_prefix, base_name):
    if not user_prefix:
        user_prefix = cluster.DEFAULT_USER_PREFIX
    return '%s-%s' % (user_prefix, base_name)


def clean_aws_credential(region_name, credential_key_name, credential_key_file):
    try:
        session = boto3.session.Session(region_name=region_name)
        service = session.resource('ec2')
        key_pair_info = service.KeyPair(credential_key_name)
        key_pair_info.delete()
        cluster.remove_if_exists(credential_key_file)
    except Exception as details:
        logger.exception(str(details))


class PublicIpNotReady(Exception):
    pass


class AWSCluster(cluster.BaseCluster):

    """
    Cluster of Node objects, started on Amazon EC2.
    """

    def __init__(self, ec2_ami_id, ec2_subnet_id, ec2_security_group_ids,
                 services, credentials, cluster_uuid=None,
                 ec2_instance_type='c4.xlarge', ec2_ami_username='root',
                 ec2_user_data='', ec2_block_device_mappings=None,
                 cluster_prefix='cluster',
                 node_prefix='node', n_nodes=10, params=None):
        region_names = params.get('region_name').split()
        if len(credentials) > 1 or len(region_names) > 1:
            assert len(credentials) == len(region_names)
        for idx, region_name in enumerate(region_names):
            credential = credentials[idx]
            if credential.type == 'generated':
                credential_key_name = credential.key_pair_name
                credential_key_file = credential.key_file
                if params.get('failure_post_behavior') == 'destroy':
                    avocado_runtime.CURRENT_TEST.runner_queue.put({'func_at_exit': clean_aws_credential,
                                                                   'args': (region_name,
                                                                            credential_key_name,
                                                                            credential_key_file),
                                                                   'once': True})
            cluster.CREDENTIALS.append(credential)

        self._ec2_ami_id = ec2_ami_id
        self._ec2_subnet_id = ec2_subnet_id
        self._ec2_security_group_ids = ec2_security_group_ids
        self._ec2_services = services
        self._credentials = credentials
        self._reuse_credentials = None
        self._ec2_instance_type = ec2_instance_type
        self._ec2_ami_username = ec2_ami_username
        if ec2_block_device_mappings is None:
            ec2_block_device_mappings = []
        self._ec2_block_device_mappings = ec2_block_device_mappings
        self._ec2_user_data = ec2_user_data
        self._ec2_ami_id = ec2_ami_id
        self.region_names = region_names
        self.instance_provision = params.get('instance_provision', default=INSTANCE_PROVISION_ON_DEMAND)
        super(AWSCluster, self).__init__(cluster_uuid=cluster_uuid,
                                         cluster_prefix=cluster_prefix,
                                         node_prefix=node_prefix,
                                         n_nodes=n_nodes,
                                         params=params,
                                         region_names=self.region_names)

    def __str__(self):
        return 'Cluster %s (AMI: %s Type: %s)' % (self.name,
                                                  self._ec2_ami_id,
                                                  self._ec2_instance_type)

    def write_node_public_ip_file(self):
        public_ip_file_path = os.path.join(self.logdir, 'public_ips')
        with open(public_ip_file_path, 'w') as public_ip_file:
            public_ip_file.write("%s" % "\n".join(self.get_node_public_ips()))
            public_ip_file.write("\n")

    def write_node_private_ip_file(self):
        private_ip_file_path = os.path.join(self.logdir, 'private_ips')
        with open(private_ip_file_path, 'w') as private_ip_file:
            private_ip_file.write("%s" % "\n".join(self.get_node_private_ips()))
            private_ip_file.write("\n")

    def _create_on_demand_instances(self, count, interfaces, ec2_user_data, dc_idx=0):
        ami_id = self._ec2_ami_id[dc_idx]
        self.log.debug("Creating {count} on-demand instances using AMI id '{ami_id}'... ".format(**locals()))
        instances = self._ec2_services[dc_idx].create_instances(ImageId=ami_id,
                                                                UserData=ec2_user_data,
                                                                MinCount=count,
                                                                MaxCount=count,
                                                                KeyName=self._credentials[dc_idx].key_pair_name,
                                                                BlockDeviceMappings=self._ec2_block_device_mappings,
                                                                NetworkInterfaces=interfaces,
                                                                InstanceType=self._ec2_instance_type)
        self.log.debug("Created instances: %s." % instances)
        return instances

    def _create_spot_instances(self, count,  interfaces, ec2_user_data='', dc_idx=0):
        ec2 = ec2_client.EC2Client(region_name=self.region_names[dc_idx])
        subnet_info = ec2.get_subnet_info(self._ec2_subnet_id[dc_idx])
        spot_params = dict(instance_type=self._ec2_instance_type,
                           image_id=self._ec2_ami_id[dc_idx],
                           region_name=subnet_info['AvailabilityZone'],
                           network_if=interfaces,
                           key_pair=self._credentials[dc_idx].key_pair_name,
                           user_data=ec2_user_data,
                           count=count,
                           block_device_mappings=self._ec2_block_device_mappings)
        if self.instance_provision == INSTANCE_PROVISION_SPOT_DURATION:
            # duration value must be a multiple of 60
            spot_params.update({'duration': cluster.TEST_DURATION / 60 * 60 + 60})

        limit = SPOT_FLEET_LIMIT if self.instance_provision == INSTANCE_PROVISION_SPOT_FLEET else SPOT_CNT_LIMIT
        request_cnt = 1
        tail_cnt = 0
        if count > limit:
            # pass common reservationid
            spot_params['user_data'] += (' --customreservation %s' % str(uuid.uuid4())[:18])
            self.log.debug("User_data to spot instances: '%s'", spot_params['user_data'])
            request_cnt = count / limit
            spot_params['count'] = limit
            tail_cnt = count % limit
            if tail_cnt:
                request_cnt += 1
        instances = []
        for i in range(request_cnt):
            if tail_cnt and i == request_cnt - 1:
                spot_params['count'] = tail_cnt
            try:
                if self.instance_provision == INSTANCE_PROVISION_SPOT_FLEET and count > 1:
                    instances_i = ec2.create_spot_fleet(**spot_params)
                else:
                    instances_i = ec2.create_spot_instances(**spot_params)
                instances.extend(instances_i)
            except ClientError as cl_ex:
                if ec2_client.MAX_SPOT_EXCEEDED_ERROR in cl_ex.message:
                    self.log.debug('Cannot create spot instance(-s): %s.'
                                   'Creating on demand instance(-s) instead.', cl_ex)
                    instances_i = self._create_on_demand_instances(count, interfaces, ec2_user_data, dc_idx)
                    instances.extend(instances_i)
                else:
                    raise

        return instances

    def _create_instances(self, count, ec2_user_data='', dc_idx=0):
        if not ec2_user_data:
            ec2_user_data = self._ec2_user_data
        self.log.debug("Passing user_data '%s' to create_instances", ec2_user_data)
        interfaces = [{'DeviceIndex': 0,
                       'SubnetId': self._ec2_subnet_id[dc_idx],
                       'AssociatePublicIpAddress': True,
                       'Groups': self._ec2_security_group_ids[dc_idx]}]

        if self.instance_provision == INSTANCE_PROVISION_ON_DEMAND:
            instances = self._create_on_demand_instances(count, interfaces, ec2_user_data, dc_idx)
        else:
            instances = self._create_spot_instances(count, interfaces, ec2_user_data, dc_idx)

        return instances

    def _get_instances(self, dc_idx):
        ec2 = ec2_client.EC2Client(region_name=self.region_names[dc_idx])
        return [ec2.get_instance_by_private_ip(ip) for ip in self._node_private_ips]

    def update_bootstrap(self, ec2_user_data, enable_auto_bootstrap):
        """
        Update --bootstrap argument inside ec2_user_data string.
        """
        if enable_auto_bootstrap:
            if '--bootstrap ' in ec2_user_data:
                ec2_user_data.replace('--bootstrap false', '--bootstrap true')
            else:
                ec2_user_data += ' --bootstrap true '
        else:
            if '--bootstrap ' in ec2_user_data:
                ec2_user_data.replace('--bootstrap true', '--bootstrap false')
            else:
                ec2_user_data += ' --bootstrap false '
        return ec2_user_data

    def add_nodes(self, count, ec2_user_data='', dc_idx=0, enable_auto_bootstrap=False):
        if cluster.Setup.REUSE_CLUSTER:
            instances = self._get_instances(dc_idx)
        else:
            instances = self._create_instances(count, ec2_user_data, dc_idx)

        added_nodes = [self._create_node(instance, self._ec2_ami_username,
                                         self.node_prefix, node_index,
                                         self.logdir, dc_idx=dc_idx)
                       for node_index, instance in
                       enumerate(instances, start=self._node_index + 1)]
        for node in added_nodes:
            node.enable_auto_bootstrap = enable_auto_bootstrap
        self._node_index += len(added_nodes)
        self.nodes += added_nodes
        self.write_node_public_ip_file()
        self.write_node_private_ip_file()
        return added_nodes

    def _create_node(self, instance, ami_username, node_prefix, node_index,
                     base_logdir, dc_idx):
        return AWSNode(ec2_instance=instance, ec2_service=self._ec2_services[dc_idx],
                       credentials=self._credentials[dc_idx], ami_username=ami_username,
                       node_prefix=node_prefix, node_index=node_index,
                       base_logdir=base_logdir, dc_idx=dc_idx)


class AWSNode(cluster.BaseNode):

    """
    Wraps EC2.Instance, so that we can also control the instance through SSH.
    """

    def __init__(self, ec2_instance, ec2_service, credentials,
                 node_prefix='node', node_index=1, ami_username='root',
                 base_logdir=None, dc_idx=0):
        name = '%s-%s' % (node_prefix, node_index)
        self._instance = ec2_instance
        self._ec2_service = ec2_service
        logger.debug("Waiting until instance {0._instance} starts running...".format(self))
        self._instance_wait_safe(self._instance.wait_until_running)
        self._wait_public_ip()
        ssh_login_info = {'hostname': None,
                          'user': ami_username,
                          'key_file': credentials.key_file}
        super(AWSNode, self).__init__(name=name,
                                      ssh_login_info=ssh_login_info,
                                      base_logdir=base_logdir,
                                      node_prefix=node_prefix,
                                      dc_idx=dc_idx)

        if not cluster.Setup.REUSE_CLUSTER:
            tags_list = [{'Key': 'Name', 'Value': name},
                         {'Key': 'workspace', 'Value': cluster.WORKSPACE},
                         {'Key': 'uname', 'Value': ' | '.join(os.uname())}]
            if cluster.TEST_DURATION >= 24 * 60 or cluster.Setup.KEEP_ALIVE:
                self.log.info('Test duration set to %s. '
                              'Keep cluster on failure %s. '
                              'Tagging node with {"keep": "alive"}',
                              cluster.TEST_DURATION, cluster.Setup.KEEP_ALIVE)
                tags_list.append({'Key': 'keep', 'Value': 'alive'})

            self._ec2_service.create_tags(Resources=[self._instance.id],
                                          Tags=tags_list)

    @property
    def public_ip_address(self):
        return self._instance.public_ip_address

    @property<commit_msg>Fixing missing spot termination task thread attribute in the AWS node constructor<commit_after>import logging
import time
import uuid
import os
import tempfile
import yaml

from botocore.exceptions import WaiterError, ClientError
import boto3.session
from avocado.utils import runtime as avocado_runtime
from textwrap import dedent
from threading import Thread

import cluster
import ec2_client
from sdcm.utils import retrying

from . import wait

logger = logging.getLogger(__name__)

INSTANCE_PROVISION_ON_DEMAND = 'on_demand'
INSTANCE_PROVISION_SPOT_FLEET = 'spot_fleet'
INSTANCE_PROVISION_SPOT_LOW_PRICE = 'spot_low_price'
INSTANCE_PROVISION_SPOT_DURATION = 'spot_duration'
SPOT_CNT_LIMIT = 20
SPOT_FLEET_LIMIT = 50


def _prepend_user_prefix(user_prefix, base_name):
    if not user_prefix:
        user_prefix = cluster.DEFAULT_USER_PREFIX
    return '%s-%s' % (user_prefix, base_name)


def clean_aws_credential(region_name, credential_key_name, credential_key_file):
    try:
        session = boto3.session.Session(region_name=region_name)
        service = session.resource('ec2')
        key_pair_info = service.KeyPair(credential_key_name)
        key_pair_info.delete()
        cluster.remove_if_exists(credential_key_file)
    except Exception as details:
        logger.exception(str(details))


class PublicIpNotReady(Exception):
    pass


class AWSCluster(cluster.BaseCluster):

    """
    Cluster of Node objects, started on Amazon EC2.
    """

    def __init__(self, ec2_ami_id, ec2_subnet_id, ec2_security_group_ids,
                 services, credentials, cluster_uuid=None,
                 ec2_instance_type='c4.xlarge', ec2_ami_username='root',
                 ec2_user_data='', ec2_block_device_mappings=None,
                 cluster_prefix='cluster',
                 node_prefix='node', n_nodes=10, params=None):
        region_names = params.get('region_name').split()
        if len(credentials) > 1 or len(region_names) > 1:
            assert len(credentials) == len(region_names)
        for idx, region_name in enumerate(region_names):
            credential = credentials[idx]
            if credential.type == 'generated':
                credential_key_name = credential.key_pair_name
                credential_key_file = credential.key_file
                if params.get('failure_post_behavior') == 'destroy':
                    avocado_runtime.CURRENT_TEST.runner_queue.put({'func_at_exit': clean_aws_credential,
                                                                   'args': (region_name,
                                                                            credential_key_name,
                                                                            credential_key_file),
                                                                   'once': True})
            cluster.CREDENTIALS.append(credential)

        self._ec2_ami_id = ec2_ami_id
        self._ec2_subnet_id = ec2_subnet_id
        self._ec2_security_group_ids = ec2_security_group_ids
        self._ec2_services = services
        self._credentials = credentials
        self._reuse_credentials = None
        self._ec2_instance_type = ec2_instance_type
        self._ec2_ami_username = ec2_ami_username
        if ec2_block_device_mappings is None:
            ec2_block_device_mappings = []
        self._ec2_block_device_mappings = ec2_block_device_mappings
        self._ec2_user_data = ec2_user_data
        self._ec2_ami_id = ec2_ami_id
        self.region_names = region_names
        self.instance_provision = params.get('instance_provision', default=INSTANCE_PROVISION_ON_DEMAND)
        super(AWSCluster, self).__init__(cluster_uuid=cluster_uuid,
                                         cluster_prefix=cluster_prefix,
                                         node_prefix=node_prefix,
                                         n_nodes=n_nodes,
                                         params=params,
                                         region_names=self.region_names)

    def __str__(self):
        return 'Cluster %s (AMI: %s Type: %s)' % (self.name,
                                                  self._ec2_ami_id,
                                                  self._ec2_instance_type)

    def write_node_public_ip_file(self):
        public_ip_file_path = os.path.join(self.logdir, 'public_ips')
        with open(public_ip_file_path, 'w') as public_ip_file:
            public_ip_file.write("%s" % "\n".join(self.get_node_public_ips()))
            public_ip_file.write("\n")

    def write_node_private_ip_file(self):
        private_ip_file_path = os.path.join(self.logdir, 'private_ips')
        with open(private_ip_file_path, 'w') as private_ip_file:
            private_ip_file.write("%s" % "\n".join(self.get_node_private_ips()))
            private_ip_file.write("\n")

    def _create_on_demand_instances(self, count, interfaces, ec2_user_data, dc_idx=0):
        ami_id = self._ec2_ami_id[dc_idx]
        self.log.debug("Creating {count} on-demand instances using AMI id '{ami_id}'... ".format(**locals()))
        instances = self._ec2_services[dc_idx].create_instances(ImageId=ami_id,
                                                                UserData=ec2_user_data,
                                                                MinCount=count,
                                                                MaxCount=count,
                                                                KeyName=self._credentials[dc_idx].key_pair_name,
                                                                BlockDeviceMappings=self._ec2_block_device_mappings,
                                                                NetworkInterfaces=interfaces,
                                                                InstanceType=self._ec2_instance_type)
        self.log.debug("Created instances: %s." % instances)
        return instances

    def _create_spot_instances(self, count,  interfaces, ec2_user_data='', dc_idx=0):
        ec2 = ec2_client.EC2Client(region_name=self.region_names[dc_idx])
        subnet_info = ec2.get_subnet_info(self._ec2_subnet_id[dc_idx])
        spot_params = dict(instance_type=self._ec2_instance_type,
                           image_id=self._ec2_ami_id[dc_idx],
                           region_name=subnet_info['AvailabilityZone'],
                           network_if=interfaces,
                           key_pair=self._credentials[dc_idx].key_pair_name,
                           user_data=ec2_user_data,
                           count=count,
                           block_device_mappings=self._ec2_block_device_mappings)
        if self.instance_provision == INSTANCE_PROVISION_SPOT_DURATION:
            # duration value must be a multiple of 60
            spot_params.update({'duration': cluster.TEST_DURATION / 60 * 60 + 60})

        limit = SPOT_FLEET_LIMIT if self.instance_provision == INSTANCE_PROVISION_SPOT_FLEET else SPOT_CNT_LIMIT
        request_cnt = 1
        tail_cnt = 0
        if count > limit:
            # pass common reservationid
            spot_params['user_data'] += (' --customreservation %s' % str(uuid.uuid4())[:18])
            self.log.debug("User_data to spot instances: '%s'", spot_params['user_data'])
            request_cnt = count / limit
            spot_params['count'] = limit
            tail_cnt = count % limit
            if tail_cnt:
                request_cnt += 1
        instances = []
        for i in range(request_cnt):
            if tail_cnt and i == request_cnt - 1:
                spot_params['count'] = tail_cnt
            try:
                if self.instance_provision == INSTANCE_PROVISION_SPOT_FLEET and count > 1:
                    instances_i = ec2.create_spot_fleet(**spot_params)
                else:
                    instances_i = ec2.create_spot_instances(**spot_params)
                instances.extend(instances_i)
            except ClientError as cl_ex:
                if ec2_client.MAX_SPOT_EXCEEDED_ERROR in cl_ex.message:
                    self.log.debug('Cannot create spot instance(-s): %s.'
                                   'Creating on demand instance(-s) instead.', cl_ex)
                    instances_i = self._create_on_demand_instances(count, interfaces, ec2_user_data, dc_idx)
                    instances.extend(instances_i)
                else:
                    raise

        return instances

    def _create_instances(self, count, ec2_user_data='', dc_idx=0):
        if not ec2_user_data:
            ec2_user_data = self._ec2_user_data
        self.log.debug("Passing user_data '%s' to create_instances", ec2_user_data)
        interfaces = [{'DeviceIndex': 0,
                       'SubnetId': self._ec2_subnet_id[dc_idx],
                       'AssociatePublicIpAddress': True,
                       'Groups': self._ec2_security_group_ids[dc_idx]}]

        if self.instance_provision == INSTANCE_PROVISION_ON_DEMAND:
            instances = self._create_on_demand_instances(count, interfaces, ec2_user_data, dc_idx)
        else:
            instances = self._create_spot_instances(count, interfaces, ec2_user_data, dc_idx)

        return instances

    def _get_instances(self, dc_idx):
        ec2 = ec2_client.EC2Client(region_name=self.region_names[dc_idx])
        return [ec2.get_instance_by_private_ip(ip) for ip in self._node_private_ips]

    def update_bootstrap(self, ec2_user_data, enable_auto_bootstrap):
        """
        Update --bootstrap argument inside ec2_user_data string.
        """
        if enable_auto_bootstrap:
            if '--bootstrap ' in ec2_user_data:
                ec2_user_data.replace('--bootstrap false', '--bootstrap true')
            else:
                ec2_user_data += ' --bootstrap true '
        else:
            if '--bootstrap ' in ec2_user_data:
                ec2_user_data.replace('--bootstrap true', '--bootstrap false')
            else:
                ec2_user_data += ' --bootstrap false '
        return ec2_user_data

    def add_nodes(self, count, ec2_user_data='', dc_idx=0, enable_auto_bootstrap=False):
        if cluster.Setup.REUSE_CLUSTER:
            instances = self._get_instances(dc_idx)
        else:
            instances = self._create_instances(count, ec2_user_data, dc_idx)

        added_nodes = [self._create_node(instance, self._ec2_ami_username,
                                         self.node_prefix, node_index,
                                         self.logdir, dc_idx=dc_idx)
                       for node_index, instance in
                       enumerate(instances, start=self._node_index + 1)]
        for node in added_nodes:
            node.enable_auto_bootstrap = enable_auto_bootstrap
        self._node_index += len(added_nodes)
        self.nodes += added_nodes
        self.write_node_public_ip_file()
        self.write_node_private_ip_file()
        return added_nodes

    def _create_node(self, instance, ami_username, node_prefix, node_index,
                     base_logdir, dc_idx):
        return AWSNode(ec2_instance=instance, ec2_service=self._ec2_services[dc_idx],
                       credentials=self._credentials[dc_idx], ami_username=ami_username,
                       node_prefix=node_prefix, node_index=node_index,
                       base_logdir=base_logdir, dc_idx=dc_idx)


class AWSNode(cluster.BaseNode):

    """
    Wraps EC2.Instance, so that we can also control the instance through SSH.
    """

    def __init__(self, ec2_instance, ec2_service, credentials,
                 node_prefix='node', node_index=1, ami_username='root',
                 base_logdir=None, dc_idx=0):
        name = '%s-%s' % (node_prefix, node_index)
        self._instance = ec2_instance
        self._ec2_service = ec2_service
        logger.debug("Waiting until instance {0._instance} starts running...".format(self))
        self._instance_wait_safe(self._instance.wait_until_running)
        self._wait_public_ip()
        ssh_login_info = {'hostname': None,
                          'user': ami_username,
                          'key_file': credentials.key_file}
        self._spot_aws_termination_task = None
        super(AWSNode, self).__init__(name=name,
                                      ssh_login_info=ssh_login_info,
                                      base_logdir=base_logdir,
                                      node_prefix=node_prefix,
                                      dc_idx=dc_idx)

        if not cluster.Setup.REUSE_CLUSTER:
            tags_list = [{'Key': 'Name', 'Value': name},
                         {'Key': 'workspace', 'Value': cluster.WORKSPACE},
                         {'Key': 'uname', 'Value': ' | '.join(os.uname())}]
            if cluster.TEST_DURATION >= 24 * 60 or cluster.Setup.KEEP_ALIVE:
                self.log.info('Test duration set to %s. '
                              'Keep cluster on failure %s. '
                              'Tagging node with {"keep": "alive"}',
                              cluster.TEST_DURATION, cluster.Setup.KEEP_ALIVE)
                tags_list.append({'Key': 'keep', 'Value': 'alive'})

            self._ec2_service.create_tags(Resources=[self._instance.id],
                                          Tags=tags_list)

    @property
    def public_ip_address(self):
        return self._instance.public_ip_address

    @property
{'fraction_non_alphanumeric': 0.08248041090241068, 'fraction_numerical': 0.009447756157912495, 'mean_word_length': 1.8836756756756756, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 98, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6469089', 'n_tokens_mistral': 7584, 'n_tokens_neox': 7381, 'n_words': 1306}