<commit_before>#!/usr/bin/env python
# Copyright 2015-2016 Yelp Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
import logging
import math
import os
import time
from collections import defaultdict
from collections import namedtuple
from datetime import datetime
from datetime import timedelta
from math import ceil
from math import floor
from typing import Any
from typing import Dict
from typing import Iterable
from typing import List
from typing import NewType
from typing import Optional
from typing import Tuple
from typing import Union

import boto3
from botocore.exceptions import ClientError
from mypy_extensions import TypedDict
from requests.exceptions import HTTPError

from paasta_tools.autoscaling import cluster_boost
from paasta_tools.autoscaling import ec2_fitness
from paasta_tools.mesos_maintenance import drain
from paasta_tools.mesos_maintenance import undrain
from paasta_tools.mesos_tools import get_mesos_master
from paasta_tools.mesos_tools import get_mesos_task_count_by_slave
from paasta_tools.mesos_tools import slave_pid_to_ip
from paasta_tools.mesos_tools import SlaveTaskCount
from paasta_tools.metrics.metastatus_lib import get_resource_utilization_by_grouping
from paasta_tools.metrics.metastatus_lib import ResourceInfo
from paasta_tools.metrics.metrics_lib import get_metrics_interface
from paasta_tools.paasta_maintenance import is_safe_to_kill
from paasta_tools.utils import load_system_paasta_config
from paasta_tools.utils import SystemPaastaConfig
from paasta_tools.utils import Timeout
from paasta_tools.utils import TimeoutError


AutoscalingInfo = namedtuple(
    "AutoscalingInfo",
    [
        "resource_id",
        "pool",
        "state",
        "current",
        "target",
        "min_capacity",
        "max_capacity",
        "instances",
    ],
)

ClusterAutoscalingResource = TypedDict(
    'ClusterAutoscalingResource',
    {
        'type': str,
        'id': str,
        'region': str,
        'pool': str,
        'min_capacity': int,
        'max_capacity': int,
    },
)

ResourcePoolSetting = TypedDict(
    'ResourcePoolSetting',
    {
        'target_utilization': float,
        'drain_timeout': int,
    },
)

MesosState = NewType('MesosState', Dict)

CLUSTER_METRICS_PROVIDER_KEY = 'cluster_metrics_provider'
DEFAULT_TARGET_UTILIZATION = 0.8  # decimal fraction
DEFAULT_DRAIN_TIMEOUT = 600  # seconds

AWS_SPOT_MODIFY_TIMEOUT = 30
MISSING_SLAVE_PANIC_THRESHOLD = .3
MAX_CLUSTER_DELTA = .2

log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())


class Timer(object):
    def __init__(self, timeout: int) -> None:
        self.timeout = timedelta(seconds=timeout)
        self.start()

    def start(self) -> None:
        self.last_start = datetime.now()

    def ready(self) -> bool:
        return datetime.now() > self.last_start + self.timeout

    def left(self) -> timedelta:
        return self.last_start + self.timeout - datetime.now()


class ClusterAutoscaler(object):

    def __init__(
        self,
        resource: ClusterAutoscalingResource,
        pool_settings: ResourcePoolSetting,
        config_folder: str,
        dry_run: bool,
        utilization_error: float,
        log_level: str=None,
        draining_enabled: bool=True,
        enable_metrics: bool=False,
    ) -> None:
        self.resource = resource
        self.pool_settings = pool_settings
        self.config_folder = config_folder
        self.dry_run = dry_run
        self.utilization_error = utilization_error
        self.draining_enabled = draining_enabled
        if log_level is not None:
            self.log.setLevel(log_level)
        self.instances: List[Dict] = []
        self.sfr: Optional[Dict[str, Any]] = None
        self.enable_metrics = enable_metrics

        self.setup_metrics()

    @property
    def log(self) -> logging.Logger:
        resource_id = self.resource.get("id", "unknown")
        name = '.'.join([__name__, self.__class__.__name__, resource_id])
        return logging.getLogger(name)

    def setup_metrics(self) -> None:
        if not self.enable_metrics:
            return None
        config = load_system_paasta_config()
        dims = {
            'paasta_cluster': config.get_cluster(),
            'region': self.resource.get('region', 'unknown'),
            'pool': self.resource.get('pool', 'unknown'),
            'resource_id': self.resource.get('id', 'unknown'),
            'resource_type': self.__class__.__name__,
        }
        self.metrics = get_metrics_interface('paasta.cluster_autoscaler')
        self.target_gauge = self.metrics.create_gauge('target_capacity', **dims)
        self.current_gauge = self.metrics.create_gauge('current_capacity', **dims)
        self.ideal_gauge = self.metrics.create_gauge('ideal_capacity', **dims)
        self.max_gauge = self.metrics.create_gauge('max_capacity', **dims)
        self.min_gauge = self.metrics.create_gauge('min_capacity', **dims)
        self.mesos_error_gauge = self.metrics.create_gauge('mesos_error', **dims)

    def emit_metrics(
        self,
        current: float,
        target: float,
        ideal: float,
    ) -> None:
        if not self.enable_metrics:
            return None
        self.current_gauge.set(current)
        self.target_gauge.set(target)
        self.ideal_gauge.set(ideal)
        self.min_gauge.set(self.resource['min_capacity'])
        self.max_gauge.set(self.resource['max_capacity'])
        self.mesos_error_gauge.set(self.utilization_error)

    def set_capacity(self, capacity: float) -> Optional[Any]:
        pass

    def get_instance_type_weights(self) -> Optional[Dict[str, float]]:
        pass

    def is_resource_cancelled(self) -> bool:
        raise NotImplementedError()

    def metrics_provider(self, mesos_state: MesosState) -> Tuple[float, int]:
        raise NotImplementedError()

    def describe_instances(
        self,
        instance_ids: List[str],
        region: Optional[str]=None,
        instance_filters: Optional[List[Dict]]=None,
    ) -> Optional[List[Dict]]:
        """This wraps ec2.describe_instances and catches instance not
        found errors. It returns a list of instance description
        dictionaries.  Optionally, a filter can be passed through to
        the ec2 call

        :param instance_ids: a list of instance ids, [] means all
        :param instance_filters: a list of ec2 filters
        :param region to connect to ec2
        :returns: a list of instance description dictionaries"""
        if not instance_filters:
            instance_filters = []
        ec2_client = boto3.client('ec2', region_name=region)
        try:
            instance_descriptions = ec2_client.describe_instances(InstanceIds=instance_ids, Filters=instance_filters)
        except ClientError as e:
            if e.response['Error']['Code'] == 'InvalidInstanceID.NotFound':
                self.log.warn('Cannot find one or more instance from IDs {}'.format(instance_ids))
                return None
            else:
                raise
        instance_reservations = [reservation['Instances'] for reservation in instance_descriptions['Reservations']]
        instances = [instance for reservation in instance_reservations for instance in reservation]
        return instances

    def describe_instance_status(
        self,
        instance_ids: List[str],
        region: Optional[str]=None,
        instance_filters: Optional[List[Dict]]=None,
    ) -> Optional[Dict]:
        """This wraps ec2.describe_instance_status and catches instance not
        found errors. It returns a list of instance description
        dictionaries.  Optionally, a filter can be passed through to
        the ec2 call

        :param instance_ids: a list of instance ids, [] means all
        :param instance_filters: a list of ec2 filters
        :param region to connect to ec2
        :returns: a list of instance description dictionaries"""
        if not instance_filters:
            instance_filters = []
        ec2_client = boto3.client('ec2', region_name=region)
        try:
            instance_descriptions = ec2_client.describe_instance_status(
                InstanceIds=instance_ids,
                Filters=instance_filters,
            )
        except ClientError as e:
            if e.response['Error']['Code'] == 'InvalidInstanceID.NotFound':
                self.log.warn('Cannot find one or more instance from IDs {}'.format(instance_ids))
                return None
            else:
                raise
        return instance_descriptions

    def get_instance_ips(
        self,
        instances: List[Dict],
        region: Optional[str]=None,
    ) -> List[str]:
        instance_descriptions = self.describe_instances(
            [instance['InstanceId'] for instance in instances],
            region=region,
        )
        instance_ips = []
        for instance in instance_descriptions:
            try:
                instance_ips.append(instance['PrivateIpAddress'])
            except KeyError:
                self.log.warning("Instance {} does not have an IP. This normally means it has been"
                                 " terminated".format(instance['InstanceId']))
        return instance_ips

    def cleanup_cancelled_config(
        self,
        resource_id: str,
        config_folder: str,
        dry_run: bool=False,
    ) -> None:
        file_name = "{}.json".format(resource_id)
        configs_to_delete = [os.path.join(walk[0], file_name)
                             for walk in os.walk(config_folder) if file_name in walk[2]]
        if not configs_to_delete:
            self.log.info("Resource config for {} not found".format(resource_id))
            return
        if not dry_run:
            os.remove(configs_to_delete[0])
        self.log.info("Deleted Resource config {}".format(configs_to_delete[0]))

    def get_aws_slaves(self, mesos_state: MesosState) -> Dict[str, Dict]:
        instance_ips = self.get_instance_ips(self.instances, region=self.resource['region'])
        slaves = {
            slave['id']: slave for slave in mesos_state.get('slaves', [])
            if slave_pid_to_ip(slave['pid']) in instance_ips and
            slave['attributes'].get('pool', 'default') == self.resource['pool']
        }
        return slaves

    def get_pool_slaves(self, mesos_state: MesosState) -> Dict[str, Dict]:
        slaves = {
            slave['id']: slave for slave in mesos_state.get('slaves', [])
            if slave['attributes'].get('pool', 'default') == self.resource['pool']
        }
        return slaves

    def check_expected_slaves(
        self,
        slaves: Dict[str, Dict],
        expected_instances: Optional[int],
    ) -> None:
        current_instances = len(slaves)
        if current_instances == 0:
            error_message = ("No instances are active, not scaling until the instances are attached to mesos")
            raise ClusterAutoscalingError(error_message)
        if expected_instances:
            self.log.info("Found %.2f%% slaves registered in mesos for this resource (%d/%d)" % (
                float(float(current_instances) / float(expected_instances)) * 100,
                current_instances,
                expected_instances,
            ))
            if float(current_instances) / expected_instances < (1.00 - MISSING_SLAVE_PANIC_THRESHOLD):
                error_message = (
                    "We currently have %d instances active in mesos out of a desired %d.\n"
                    "Refusing to scale because we either need to wait for the requests to be "
                    "filled, or the new instances are not healthy for some reason.\n"
                ) % (
                    current_instances, expected_instances,
                )
                if self.sfr and self.sfr['SpotFleetRequestState'] == 'cancelled_running':
                    log.warn(
                        error_message + "But this is an sfr in cancelled_running state, so continuing anyways",
                    )
                    return None
                error_message += "(refusing to go past %.2f%% missing instances)" % MISSING_SLAVE_PANIC_THRESHOLD
                raise ClusterAutoscalingError(error_message)

    def can_kill(
        self,
        hostname: str,
        should_drain: bool,
        dry_run: bool,
        timer: Timer,
    ) -> bool:
        if dry_run:
            return True
        if timer.ready():
            self.log.warning("Timer expired before slave ready to kill, proceding to terminate anyways")
            timer.start()
            raise TimeoutError
        if not should_drain:
            self.log.info("Not draining, waiting %s longer before killing" % timer.left())
            return False
        if is_safe_to_kill(hostname):
            self.log.info("Slave %s is ready to kill, with %s left on timer" % (hostname, timer.left()))
            timer.start()
            return True
        return False

    async def wait_and_terminate(
        self,
        slave: 'PaastaAwsSlave',
        drain_timeout: int,
        dry_run: bool,
        timer: Timer,
        region: Optional[str]=None,
        should_drain: bool=True,
    ) -> None:
        """Waits for slave to be drained and then terminate

        :param slave: dict of slave to kill
        :param drain_timeout: how long to wait before terminating
            even if not drained
        :param region to connect to ec2
        :param dry_run: Don't drain or make changes to spot fleet if True
        :param should_drain: whether we should drain hosts before waiting to stop them
        :param timer: a Timer object to keep terminates happening once every n seconds accross co-routines
        """
        ec2_client = boto3.client('ec2', region_name=region)
        self.log.info("Starting TERMINATING: {} (Hostname = {}, IP = {})".format(
            slave.instance_id,
            slave.hostname,
            slave.ip,
        ))
        try:
            # This loop should always finish because the maintenance window should trigger is_ready_to_kill
            # being true. Just in case though we set a timeout (from timer) and terminate anyway
            while True:
                instance_id = slave.instance_id
                if not instance_id:
                    self.log.warning(
                        "Didn't find instance ID for slave: {}. Skipping terminating".format(slave.pid),
                    )
                    break
                # Check if no tasks are running or we have reached the maintenance window
                if self.can_kill(slave.hostname, should_drain, dry_run, timer):
                    self.log.info("TERMINATING: {} (Hostname = {}, IP = {})".format(
                        instance_id,
                        slave.hostname,
                        slave.ip,
                    ))
                    try:
                        ec2_client.terminate_instances(InstanceIds=[instance_id], DryRun=dry_run)
                    except ClientError as e:
                        if e.response['Error'].get('Code') == 'DryRunOperation':
                            pass
                        else:
                            raise
                    break
                else:
                    self.log.info("Instance {}: NOT ready to kill".format(instance_id))
                self.log.debug("Waiting 5 seconds and then checking again")
                await asyncio.sleep(5)
        except TimeoutError:
            self.log.error("Timed out after {} waiting to drain {}, now terminating anyway".format(
                timer.timeout,
                slave.pid,
            ))
            try:
                ec2_client.terminate_instances(InstanceIds=[instance_id], DryRun=dry_run)
            except ClientError as e:
                if e.response['Error'].get('Code') == 'DryRunOperation':
                    pass
                else:
                    raise

    async def scale_resource(
        self,
        current_capacity: float,
        target_capacity: int,
    ) -> None:
        """Scales an AWS resource based on current and target capacity
        If scaling up we just set target capacity and let AWS take care of the rest
        If scaling down we pick the slaves we'd prefer to kill, put them in maintenance
        mode and drain them (via paasta_maintenance and setup_marathon_jobs). We then kill
        them once they are running 0 tasks or once a timeout is reached

        :param current_capacity: integer current resource capacity
        :param target_capacity: target resource capacity
        """
        target_capacity = int(target_capacity)
        delta = target_capacity - current_capacity
        if delta == 0:
            self.log.info("Already at target capacity: {}".format(target_capacity))
            return
        elif delta > 0:
            self.log.info("Increasing resource capacity to: {}".format(target_capacity))
            self.set_capacity(target_capacity)
            return
        elif delta < 0:
            mesos_state = get_mesos_master().state_summary()
            slaves_list = get_mesos_task_count_by_slave(mesos_state, pool=self.resource['pool'])
            filtered_slaves = self.filter_aws_slaves(slaves_list)
            killable_capacity = round(sum([slave.instance_weight for slave in filtered_slaves]), 2)
            amount_to_decrease = delta * -1
            if amount_to_decrease > killable_capacity:
                self.log.error(
                    "Didn't find enough candidates to kill. This shouldn't happen so let's not kill anything! "<commit_msg>Rounding capacity to terminate too<commit_after>#!/usr/bin/env python
# Copyright 2015-2016 Yelp Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import asyncio
import logging
import math
import os
import time
from collections import defaultdict
from collections import namedtuple
from datetime import datetime
from datetime import timedelta
from math import ceil
from math import floor
from typing import Any
from typing import Dict
from typing import Iterable
from typing import List
from typing import NewType
from typing import Optional
from typing import Tuple
from typing import Union

import boto3
from botocore.exceptions import ClientError
from mypy_extensions import TypedDict
from requests.exceptions import HTTPError

from paasta_tools.autoscaling import cluster_boost
from paasta_tools.autoscaling import ec2_fitness
from paasta_tools.mesos_maintenance import drain
from paasta_tools.mesos_maintenance import undrain
from paasta_tools.mesos_tools import get_mesos_master
from paasta_tools.mesos_tools import get_mesos_task_count_by_slave
from paasta_tools.mesos_tools import slave_pid_to_ip
from paasta_tools.mesos_tools import SlaveTaskCount
from paasta_tools.metrics.metastatus_lib import get_resource_utilization_by_grouping
from paasta_tools.metrics.metastatus_lib import ResourceInfo
from paasta_tools.metrics.metrics_lib import get_metrics_interface
from paasta_tools.paasta_maintenance import is_safe_to_kill
from paasta_tools.utils import load_system_paasta_config
from paasta_tools.utils import SystemPaastaConfig
from paasta_tools.utils import Timeout
from paasta_tools.utils import TimeoutError


AutoscalingInfo = namedtuple(
    "AutoscalingInfo",
    [
        "resource_id",
        "pool",
        "state",
        "current",
        "target",
        "min_capacity",
        "max_capacity",
        "instances",
    ],
)

ClusterAutoscalingResource = TypedDict(
    'ClusterAutoscalingResource',
    {
        'type': str,
        'id': str,
        'region': str,
        'pool': str,
        'min_capacity': int,
        'max_capacity': int,
    },
)

ResourcePoolSetting = TypedDict(
    'ResourcePoolSetting',
    {
        'target_utilization': float,
        'drain_timeout': int,
    },
)

MesosState = NewType('MesosState', Dict)

CLUSTER_METRICS_PROVIDER_KEY = 'cluster_metrics_provider'
DEFAULT_TARGET_UTILIZATION = 0.8  # decimal fraction
DEFAULT_DRAIN_TIMEOUT = 600  # seconds

AWS_SPOT_MODIFY_TIMEOUT = 30
MISSING_SLAVE_PANIC_THRESHOLD = .3
MAX_CLUSTER_DELTA = .2

log = logging.getLogger(__name__)
log.addHandler(logging.NullHandler())


class Timer(object):
    def __init__(self, timeout: int) -> None:
        self.timeout = timedelta(seconds=timeout)
        self.start()

    def start(self) -> None:
        self.last_start = datetime.now()

    def ready(self) -> bool:
        return datetime.now() > self.last_start + self.timeout

    def left(self) -> timedelta:
        return self.last_start + self.timeout - datetime.now()


class ClusterAutoscaler(object):

    def __init__(
        self,
        resource: ClusterAutoscalingResource,
        pool_settings: ResourcePoolSetting,
        config_folder: str,
        dry_run: bool,
        utilization_error: float,
        log_level: str=None,
        draining_enabled: bool=True,
        enable_metrics: bool=False,
    ) -> None:
        self.resource = resource
        self.pool_settings = pool_settings
        self.config_folder = config_folder
        self.dry_run = dry_run
        self.utilization_error = utilization_error
        self.draining_enabled = draining_enabled
        if log_level is not None:
            self.log.setLevel(log_level)
        self.instances: List[Dict] = []
        self.sfr: Optional[Dict[str, Any]] = None
        self.enable_metrics = enable_metrics

        self.setup_metrics()

    @property
    def log(self) -> logging.Logger:
        resource_id = self.resource.get("id", "unknown")
        name = '.'.join([__name__, self.__class__.__name__, resource_id])
        return logging.getLogger(name)

    def setup_metrics(self) -> None:
        if not self.enable_metrics:
            return None
        config = load_system_paasta_config()
        dims = {
            'paasta_cluster': config.get_cluster(),
            'region': self.resource.get('region', 'unknown'),
            'pool': self.resource.get('pool', 'unknown'),
            'resource_id': self.resource.get('id', 'unknown'),
            'resource_type': self.__class__.__name__,
        }
        self.metrics = get_metrics_interface('paasta.cluster_autoscaler')
        self.target_gauge = self.metrics.create_gauge('target_capacity', **dims)
        self.current_gauge = self.metrics.create_gauge('current_capacity', **dims)
        self.ideal_gauge = self.metrics.create_gauge('ideal_capacity', **dims)
        self.max_gauge = self.metrics.create_gauge('max_capacity', **dims)
        self.min_gauge = self.metrics.create_gauge('min_capacity', **dims)
        self.mesos_error_gauge = self.metrics.create_gauge('mesos_error', **dims)

    def emit_metrics(
        self,
        current: float,
        target: float,
        ideal: float,
    ) -> None:
        if not self.enable_metrics:
            return None
        self.current_gauge.set(current)
        self.target_gauge.set(target)
        self.ideal_gauge.set(ideal)
        self.min_gauge.set(self.resource['min_capacity'])
        self.max_gauge.set(self.resource['max_capacity'])
        self.mesos_error_gauge.set(self.utilization_error)

    def set_capacity(self, capacity: float) -> Optional[Any]:
        pass

    def get_instance_type_weights(self) -> Optional[Dict[str, float]]:
        pass

    def is_resource_cancelled(self) -> bool:
        raise NotImplementedError()

    def metrics_provider(self, mesos_state: MesosState) -> Tuple[float, int]:
        raise NotImplementedError()

    def describe_instances(
        self,
        instance_ids: List[str],
        region: Optional[str]=None,
        instance_filters: Optional[List[Dict]]=None,
    ) -> Optional[List[Dict]]:
        """This wraps ec2.describe_instances and catches instance not
        found errors. It returns a list of instance description
        dictionaries.  Optionally, a filter can be passed through to
        the ec2 call

        :param instance_ids: a list of instance ids, [] means all
        :param instance_filters: a list of ec2 filters
        :param region to connect to ec2
        :returns: a list of instance description dictionaries"""
        if not instance_filters:
            instance_filters = []
        ec2_client = boto3.client('ec2', region_name=region)
        try:
            instance_descriptions = ec2_client.describe_instances(InstanceIds=instance_ids, Filters=instance_filters)
        except ClientError as e:
            if e.response['Error']['Code'] == 'InvalidInstanceID.NotFound':
                self.log.warn('Cannot find one or more instance from IDs {}'.format(instance_ids))
                return None
            else:
                raise
        instance_reservations = [reservation['Instances'] for reservation in instance_descriptions['Reservations']]
        instances = [instance for reservation in instance_reservations for instance in reservation]
        return instances

    def describe_instance_status(
        self,
        instance_ids: List[str],
        region: Optional[str]=None,
        instance_filters: Optional[List[Dict]]=None,
    ) -> Optional[Dict]:
        """This wraps ec2.describe_instance_status and catches instance not
        found errors. It returns a list of instance description
        dictionaries.  Optionally, a filter can be passed through to
        the ec2 call

        :param instance_ids: a list of instance ids, [] means all
        :param instance_filters: a list of ec2 filters
        :param region to connect to ec2
        :returns: a list of instance description dictionaries"""
        if not instance_filters:
            instance_filters = []
        ec2_client = boto3.client('ec2', region_name=region)
        try:
            instance_descriptions = ec2_client.describe_instance_status(
                InstanceIds=instance_ids,
                Filters=instance_filters,
            )
        except ClientError as e:
            if e.response['Error']['Code'] == 'InvalidInstanceID.NotFound':
                self.log.warn('Cannot find one or more instance from IDs {}'.format(instance_ids))
                return None
            else:
                raise
        return instance_descriptions

    def get_instance_ips(
        self,
        instances: List[Dict],
        region: Optional[str]=None,
    ) -> List[str]:
        instance_descriptions = self.describe_instances(
            [instance['InstanceId'] for instance in instances],
            region=region,
        )
        instance_ips = []
        for instance in instance_descriptions:
            try:
                instance_ips.append(instance['PrivateIpAddress'])
            except KeyError:
                self.log.warning("Instance {} does not have an IP. This normally means it has been"
                                 " terminated".format(instance['InstanceId']))
        return instance_ips

    def cleanup_cancelled_config(
        self,
        resource_id: str,
        config_folder: str,
        dry_run: bool=False,
    ) -> None:
        file_name = "{}.json".format(resource_id)
        configs_to_delete = [os.path.join(walk[0], file_name)
                             for walk in os.walk(config_folder) if file_name in walk[2]]
        if not configs_to_delete:
            self.log.info("Resource config for {} not found".format(resource_id))
            return
        if not dry_run:
            os.remove(configs_to_delete[0])
        self.log.info("Deleted Resource config {}".format(configs_to_delete[0]))

    def get_aws_slaves(self, mesos_state: MesosState) -> Dict[str, Dict]:
        instance_ips = self.get_instance_ips(self.instances, region=self.resource['region'])
        slaves = {
            slave['id']: slave for slave in mesos_state.get('slaves', [])
            if slave_pid_to_ip(slave['pid']) in instance_ips and
            slave['attributes'].get('pool', 'default') == self.resource['pool']
        }
        return slaves

    def get_pool_slaves(self, mesos_state: MesosState) -> Dict[str, Dict]:
        slaves = {
            slave['id']: slave for slave in mesos_state.get('slaves', [])
            if slave['attributes'].get('pool', 'default') == self.resource['pool']
        }
        return slaves

    def check_expected_slaves(
        self,
        slaves: Dict[str, Dict],
        expected_instances: Optional[int],
    ) -> None:
        current_instances = len(slaves)
        if current_instances == 0:
            error_message = ("No instances are active, not scaling until the instances are attached to mesos")
            raise ClusterAutoscalingError(error_message)
        if expected_instances:
            self.log.info("Found %.2f%% slaves registered in mesos for this resource (%d/%d)" % (
                float(float(current_instances) / float(expected_instances)) * 100,
                current_instances,
                expected_instances,
            ))
            if float(current_instances) / expected_instances < (1.00 - MISSING_SLAVE_PANIC_THRESHOLD):
                error_message = (
                    "We currently have %d instances active in mesos out of a desired %d.\n"
                    "Refusing to scale because we either need to wait for the requests to be "
                    "filled, or the new instances are not healthy for some reason.\n"
                ) % (
                    current_instances, expected_instances,
                )
                if self.sfr and self.sfr['SpotFleetRequestState'] == 'cancelled_running':
                    log.warn(
                        error_message + "But this is an sfr in cancelled_running state, so continuing anyways",
                    )
                    return None
                error_message += "(refusing to go past %.2f%% missing instances)" % MISSING_SLAVE_PANIC_THRESHOLD
                raise ClusterAutoscalingError(error_message)

    def can_kill(
        self,
        hostname: str,
        should_drain: bool,
        dry_run: bool,
        timer: Timer,
    ) -> bool:
        if dry_run:
            return True
        if timer.ready():
            self.log.warning("Timer expired before slave ready to kill, proceding to terminate anyways")
            timer.start()
            raise TimeoutError
        if not should_drain:
            self.log.info("Not draining, waiting %s longer before killing" % timer.left())
            return False
        if is_safe_to_kill(hostname):
            self.log.info("Slave %s is ready to kill, with %s left on timer" % (hostname, timer.left()))
            timer.start()
            return True
        return False

    async def wait_and_terminate(
        self,
        slave: 'PaastaAwsSlave',
        drain_timeout: int,
        dry_run: bool,
        timer: Timer,
        region: Optional[str]=None,
        should_drain: bool=True,
    ) -> None:
        """Waits for slave to be drained and then terminate

        :param slave: dict of slave to kill
        :param drain_timeout: how long to wait before terminating
            even if not drained
        :param region to connect to ec2
        :param dry_run: Don't drain or make changes to spot fleet if True
        :param should_drain: whether we should drain hosts before waiting to stop them
        :param timer: a Timer object to keep terminates happening once every n seconds accross co-routines
        """
        ec2_client = boto3.client('ec2', region_name=region)
        self.log.info("Starting TERMINATING: {} (Hostname = {}, IP = {})".format(
            slave.instance_id,
            slave.hostname,
            slave.ip,
        ))
        try:
            # This loop should always finish because the maintenance window should trigger is_ready_to_kill
            # being true. Just in case though we set a timeout (from timer) and terminate anyway
            while True:
                instance_id = slave.instance_id
                if not instance_id:
                    self.log.warning(
                        "Didn't find instance ID for slave: {}. Skipping terminating".format(slave.pid),
                    )
                    break
                # Check if no tasks are running or we have reached the maintenance window
                if self.can_kill(slave.hostname, should_drain, dry_run, timer):
                    self.log.info("TERMINATING: {} (Hostname = {}, IP = {})".format(
                        instance_id,
                        slave.hostname,
                        slave.ip,
                    ))
                    try:
                        ec2_client.terminate_instances(InstanceIds=[instance_id], DryRun=dry_run)
                    except ClientError as e:
                        if e.response['Error'].get('Code') == 'DryRunOperation':
                            pass
                        else:
                            raise
                    break
                else:
                    self.log.info("Instance {}: NOT ready to kill".format(instance_id))
                self.log.debug("Waiting 5 seconds and then checking again")
                await asyncio.sleep(5)
        except TimeoutError:
            self.log.error("Timed out after {} waiting to drain {}, now terminating anyway".format(
                timer.timeout,
                slave.pid,
            ))
            try:
                ec2_client.terminate_instances(InstanceIds=[instance_id], DryRun=dry_run)
            except ClientError as e:
                if e.response['Error'].get('Code') == 'DryRunOperation':
                    pass
                else:
                    raise

    async def scale_resource(
        self,
        current_capacity: float,
        target_capacity: int,
    ) -> None:
        """Scales an AWS resource based on current and target capacity
        If scaling up we just set target capacity and let AWS take care of the rest
        If scaling down we pick the slaves we'd prefer to kill, put them in maintenance
        mode and drain them (via paasta_maintenance and setup_marathon_jobs). We then kill
        them once they are running 0 tasks or once a timeout is reached

        :param current_capacity: integer current resource capacity
        :param target_capacity: target resource capacity
        """
        target_capacity = int(target_capacity)
        delta = target_capacity - current_capacity
        if delta == 0:
            self.log.info("Already at target capacity: {}".format(target_capacity))
            return
        elif delta > 0:
            self.log.info("Increasing resource capacity to: {}".format(target_capacity))
            self.set_capacity(target_capacity)
            return
        elif delta < 0:
            mesos_state = get_mesos_master().state_summary()
            slaves_list = get_mesos_task_count_by_slave(mesos_state, pool=self.resource['pool'])
            filtered_slaves = self.filter_aws_slaves(slaves_list)
            killable_capacity = round(sum([slave.instance_weight for slave in filtered_slaves]), 2)
            amount_to_decrease = round(delta * -1, 2)
            if amount_to_decrease > killable_capacity:
                self.log.error(
                    "Didn't find enough candidates to kill. This shouldn't happen so let's not kill anything! "
{'fraction_non_alphanumeric': 0.0837074783087729, 'fraction_numerical': 0.003663407244181242, 'mean_word_length': 2.3065573770491805, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 53, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1211086', 'n_tokens_mistral': 10022, 'n_tokens_neox': 9637, 'n_words': 2737}