<commit_before>import collections
import datetime
import logging
import math
import time

import botocore
import boto3
import botocore.exceptions
import datadog
import pykube

import autoscaler.autoscaling_groups as autoscaling_groups
import autoscaler.azure as azure
import autoscaler.capacity as capacity
from autoscaler.kube import KubePod, KubeNode, KubeResource, KubePodStatus
import autoscaler.reservations as reservations
import autoscaler.utils as utils

# we are interested in all pods, incl. system ones
pykube.Pod.objects.namespace = None

# HACK: https://github.com/kelproject/pykube/issues/29#issuecomment-230026930
import backports.ssl_match_hostname
# Monkey-patch match_hostname with backports's match_hostname, allowing for IP addresses
# XXX: the exception that this might raise is
# backports.ssl_match_hostname.CertificateError
pykube.http.requests.packages.urllib3.connection.match_hostname = backports.ssl_match_hostname.match_hostname

logger = logging.getLogger(__name__)


class ClusterNodeState(object):
    INSTANCE_TERMINATED = 'instance-terminated'
    ASG_MIN_SIZE = 'asg-min-size'
    POD_PENDING = 'pod-pending'
    GRACE_PERIOD = 'grace-period'
    TYPE_GRACE_PERIOD = 'type-grace-period'
    IDLE_SCHEDULABLE = 'idle-schedulable'
    IDLE_UNSCHEDULABLE = 'idle-unschedulable'
    BUSY_UNSCHEDULABLE = 'busy-unschedulable'
    BUSY = 'busy'
    UNDER_UTILIZED_DRAINABLE = 'under-utilized-drainable'
    UNDER_UTILIZED_UNDRAINABLE = 'under-utilized-undrainable'
    LAUNCH_HR_GRACE_PERIOD = 'launch-hr-grace-period'
    RESERVED = 'reserved'
    DETACHED = 'detached'


class Cluster(object):

    # the number of instances per type that is allowed to be idle
    # this is for keeping some spare capacity around for faster
    # pod scheduling, instead of keeping the cluster at capacity
    # and having to spin up nodes for every job submission
    TYPE_IDLE_COUNT = 5

    # the utilization threshold under which to consider a node
    # under utilized and drainable
    UTIL_THRESHOLD = 0.3

    # since we pay for the full hour, don't prematurely kill instances
    # the number of minutes into the launch hour at which an instance
    # is fine to kill
    LAUNCH_HOUR_THRESHOLD = 60 * 30  # 40 minutes

    # HACK: before we're ready to favor bigger instances in all cases
    # just prioritize the ones that we're confident about
    _GROUP_DEFAULT_PRIORITY = 10
    _GROUP_PRIORITIES = {
        'g2.8xlarge': 2,
        'm4.xlarge': 0,
        'm4.2xlarge': 0,
        'm4.4xlarge': 0,
        'm4.10xlarge': 0
    }

    def __init__(self, aws_regions, aws_access_key, aws_secret_key,
                 azure_regions, kubeconfig,
                 idle_threshold, type_idle_threshold,
                 instance_init_time, cluster_name, notifier,
                 scale_up=True, maintainance=True,
                 datadog_api_key=None,
                 over_provision=5, dry_run=False):
        if kubeconfig:
            # for using locally
            logger.debug('Using kubeconfig %s', kubeconfig)
            self.api = pykube.HTTPClient(
                pykube.KubeConfig.from_file(kubeconfig))
        else:
            # for using on kube
            logger.debug('Using kube service account')
            self.api = pykube.HTTPClient(
                pykube.KubeConfig.from_service_account())

        self._drained = {}
        self.session = boto3.session.Session(
            aws_access_key_id=aws_access_key,
            aws_secret_access_key=aws_secret_key,
            region_name=aws_regions[0])  # provide a default region
        self.autoscaling_groups = autoscaling_groups.AutoScalingGroups(
            session=self.session, regions=aws_regions,
            cluster_name=cluster_name)
        self.autoscaling_timeouts = autoscaling_groups.AutoScalingTimeouts(
            self.session)

        self.azure_groups = azure.AzureGroups(azure_regions)

        self.reservation_client = reservations.ReservationClient()

        # config
        self.azure_regions = azure_regions
        self.aws_regions = aws_regions
        self.idle_threshold = idle_threshold
        self.instance_init_time = instance_init_time
        self.type_idle_threshold = type_idle_threshold
        self.over_provision = over_provision

        self.scale_up = scale_up
        self.maintainance = maintainance

        self.notifier = notifier

        if datadog_api_key:
            datadog.initialize(api_key=datadog_api_key)
            logger.info('Datadog initialized')
        self.stats = datadog.ThreadStats()
        self.stats.start()

        self.dry_run = dry_run

    def scale_loop(self):
        """
        runs one loop of scaling to current needs.
        returns True if successfully scaled.
        """
        logger.info("++++++++++++++ Running Scaling Loop ++++++++++++++++")
        try:
            start_time = time.time()

            pykube_nodes = pykube.Node.objects(self.api)
            if not pykube_nodes:
                logger.warn('Failed to list nodes. Please check kube configuration. Terminating scale loop.')
                return False

            all_nodes = map(KubeNode, pykube_nodes)
            managed_nodes = [node for node in all_nodes if node.is_managed()]

            self._disable_azure = False

            running_insts_map = self.get_running_instances_map(managed_nodes)

            pods = map(KubePod, pykube.Pod.objects(self.api))

            running_or_pending_assigned_pods = [
                p for p in pods if (p.status == KubePodStatus.RUNNING or p.status == KubePodStatus.CONTAINER_CREATING) or (
                    p.status == KubePodStatus.PENDING and p.node_name
                )
            ]

            for node in all_nodes:
                for pod in running_or_pending_assigned_pods:
                    if pod.node_name == node.name:
                        node.count_pod(pod)

            asgs = self.autoscaling_groups.get_all_groups(all_nodes)
            azure_groups = self.azure_groups.get_all_groups(all_nodes)
            scaling_groups = asgs + azure_groups

            pods_to_schedule = self.get_pods_to_schedule(pods)

            res_data = self.reservation_client.list_reservations()
            reservations_map = dict((r['id'], reservations.Reservation(r, running_insts_map, managed_nodes))
                                    for r in res_data['reservations'])

            if self.scale_up:
                logger.info(
                    "++++++++++++++ Scaling Up Begins ++++++++++++++++")
                self.scale(
                    pods_to_schedule, all_nodes, scaling_groups,
                    running_insts_map, reservations_map)
                logger.info("++++++++++++++ Scaling Up Ends ++++++++++++++++")
            if self.maintainance:
                logger.info(
                    "++++++++++++++ Maintenance Begins ++++++++++++++++")
                self.maintain(
                    managed_nodes, running_insts_map,
                    pods_to_schedule, running_or_pending_assigned_pods,
                    scaling_groups, reservations_map)
                logger.info("++++++++++++++ Maintenance Ends ++++++++++++++++")<commit_msg>Add more counters to autoscaler to monitor performance<commit_after>import collections
import datetime
import logging
import math
import time

import botocore
import boto3
import botocore.exceptions
import datadog
import pykube

import autoscaler.autoscaling_groups as autoscaling_groups
import autoscaler.azure as azure
import autoscaler.capacity as capacity
from autoscaler.kube import KubePod, KubeNode, KubeResource, KubePodStatus
import autoscaler.reservations as reservations
import autoscaler.utils as utils

# we are interested in all pods, incl. system ones
pykube.Pod.objects.namespace = None

# HACK: https://github.com/kelproject/pykube/issues/29#issuecomment-230026930
import backports.ssl_match_hostname
# Monkey-patch match_hostname with backports's match_hostname, allowing for IP addresses
# XXX: the exception that this might raise is
# backports.ssl_match_hostname.CertificateError
pykube.http.requests.packages.urllib3.connection.match_hostname = backports.ssl_match_hostname.match_hostname

logger = logging.getLogger(__name__)


class ClusterNodeState(object):
    INSTANCE_TERMINATED = 'instance-terminated'
    ASG_MIN_SIZE = 'asg-min-size'
    POD_PENDING = 'pod-pending'
    GRACE_PERIOD = 'grace-period'
    TYPE_GRACE_PERIOD = 'type-grace-period'
    IDLE_SCHEDULABLE = 'idle-schedulable'
    IDLE_UNSCHEDULABLE = 'idle-unschedulable'
    BUSY_UNSCHEDULABLE = 'busy-unschedulable'
    BUSY = 'busy'
    UNDER_UTILIZED_DRAINABLE = 'under-utilized-drainable'
    UNDER_UTILIZED_UNDRAINABLE = 'under-utilized-undrainable'
    LAUNCH_HR_GRACE_PERIOD = 'launch-hr-grace-period'
    RESERVED = 'reserved'
    DETACHED = 'detached'


class Cluster(object):

    # the number of instances per type that is allowed to be idle
    # this is for keeping some spare capacity around for faster
    # pod scheduling, instead of keeping the cluster at capacity
    # and having to spin up nodes for every job submission
    TYPE_IDLE_COUNT = 5

    # the utilization threshold under which to consider a node
    # under utilized and drainable
    UTIL_THRESHOLD = 0.3

    # since we pay for the full hour, don't prematurely kill instances
    # the number of minutes into the launch hour at which an instance
    # is fine to kill
    LAUNCH_HOUR_THRESHOLD = 60 * 30  # 40 minutes

    # HACK: before we're ready to favor bigger instances in all cases
    # just prioritize the ones that we're confident about
    _GROUP_DEFAULT_PRIORITY = 10
    _GROUP_PRIORITIES = {
        'g2.8xlarge': 2,
        'm4.xlarge': 0,
        'm4.2xlarge': 0,
        'm4.4xlarge': 0,
        'm4.10xlarge': 0
    }

    def __init__(self, aws_regions, aws_access_key, aws_secret_key,
                 azure_regions, kubeconfig,
                 idle_threshold, type_idle_threshold,
                 instance_init_time, cluster_name, notifier,
                 scale_up=True, maintainance=True,
                 datadog_api_key=None,
                 over_provision=5, dry_run=False):
        if kubeconfig:
            # for using locally
            logger.debug('Using kubeconfig %s', kubeconfig)
            self.api = pykube.HTTPClient(
                pykube.KubeConfig.from_file(kubeconfig))
        else:
            # for using on kube
            logger.debug('Using kube service account')
            self.api = pykube.HTTPClient(
                pykube.KubeConfig.from_service_account())

        self._drained = {}
        self.session = boto3.session.Session(
            aws_access_key_id=aws_access_key,
            aws_secret_access_key=aws_secret_key,
            region_name=aws_regions[0])  # provide a default region
        self.autoscaling_groups = autoscaling_groups.AutoScalingGroups(
            session=self.session, regions=aws_regions,
            cluster_name=cluster_name)
        self.autoscaling_timeouts = autoscaling_groups.AutoScalingTimeouts(
            self.session)

        self.azure_groups = azure.AzureGroups(azure_regions)

        self.reservation_client = reservations.ReservationClient()

        # config
        self.azure_regions = azure_regions
        self.aws_regions = aws_regions
        self.idle_threshold = idle_threshold
        self.instance_init_time = instance_init_time
        self.type_idle_threshold = type_idle_threshold
        self.over_provision = over_provision

        self.scale_up = scale_up
        self.maintainance = maintainance

        self.notifier = notifier

        if datadog_api_key:
            datadog.initialize(api_key=datadog_api_key)
            logger.info('Datadog initialized')
        self.stats = datadog.ThreadStats()
        self.stats.start()

        self.dry_run = dry_run

    def scale_loop(self):
        """
        runs one loop of scaling to current needs.
        returns True if successfully scaled.
        """
        logger.info("++++++++++++++ Running Scaling Loop ++++++++++++++++")
        try:
            start_time = time.time()

            kube_lookup_start_time = time.time()
            pykube_nodes = pykube.Node.objects(self.api)
            if not pykube_nodes:
                logger.warn('Failed to list nodes. Please check kube configuration. Terminating scale loop.')
                return False

            all_nodes = map(KubeNode, pykube_nodes)
            managed_nodes = [node for node in all_nodes if node.is_managed()]

            self._disable_azure = False

            running_insts_map = self.get_running_instances_map(managed_nodes)

            pods = map(KubePod, pykube.Pod.objects(self.api))

            running_or_pending_assigned_pods = [
                p for p in pods if (p.status == KubePodStatus.RUNNING or p.status == KubePodStatus.CONTAINER_CREATING) or (
                    p.status == KubePodStatus.PENDING and p.node_name
                )
            ]

            for node in all_nodes:
                for pod in running_or_pending_assigned_pods:
                    if pod.node_name == node.name:
                        node.count_pod(pod)
            self.stats.gauge('autoscaler.scaling_loop.kube_lookup_time', time.time() - kube_lookup_start_time)

            scaling_group_lookup_start_time = time.time()
            asgs = self.autoscaling_groups.get_all_groups(all_nodes)
            azure_groups = self.azure_groups.get_all_groups(all_nodes)
            scaling_groups = asgs + azure_groups
            self.stats.gauge('autoscaler.scaling_loop.scaling_group_lookup_time', time.time() - scaling_group_lookup_start_time)

            pods_to_schedule = self.get_pods_to_schedule(pods)

            res_data = self.reservation_client.list_reservations()
            reservations_map = dict((r['id'], reservations.Reservation(r, running_insts_map, managed_nodes))
                                    for r in res_data['reservations'])

            if self.scale_up:
                logger.info(
                    "++++++++++++++ Scaling Up Begins ++++++++++++++++")
                self.scale(
                    pods_to_schedule, all_nodes, scaling_groups,
                    running_insts_map, reservations_map)
                logger.info("++++++++++++++ Scaling Up Ends ++++++++++++++++")
            if self.maintainance:
                logger.info(
                    "++++++++++++++ Maintenance Begins ++++++++++++++++")
                self.maintain(
                    managed_nodes, running_insts_map,
                    pods_to_schedule, running_or_pending_assigned_pods,
                    scaling_groups, reservations_map)
                logger.info("++++++++++++++ Maintenance Ends ++++++++++++++++")
{'fraction_non_alphanumeric': 0.08780224233418885, 'fraction_numerical': 0.005673375658516818, 'mean_word_length': 2.4872821479039096, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, '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': '2318696', 'n_tokens_mistral': 4550, 'n_tokens_neox': 4300, 'n_words': 953}